YugabyteDB (2.13.0.0-b42, bfc6a6643e7399ac8a0e81d06a3ee6d6571b33ab)

// Copyright (c) YugaByte, Inc.

//

// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except

// in compliance with the License.  You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software distributed under the License

// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express

// or implied.  See the License for the specific language governing permissions and limitations

// under the License.

//

#include <algorithm>

#include <map>

#include <string>

#include <utility>

#include <chrono>

#include <boost/assign.hpp>

#include <gflags/gflags.h>

#include <gtest/gtest.h>

#include "yb/common/common.pb.h"

#include "yb/common/entity_ids.h"

#include "yb/common/ql_value.h"

#include "yb/common/schema.h"

#include "yb/common/wire_protocol.h"

#include "yb/cdc/cdc_service.h"

#include "yb/cdc/cdc_service.pb.h"

#include "yb/cdc/cdc_service.proxy.h"

#include "yb/client/client.h"

#include "yb/client/client-test-util.h"

#include "yb/client/meta_cache.h"

#include "yb/client/schema.h"

#include "yb/client/session.h"

#include "yb/client/table.h"

#include "yb/client/table_alterer.h"

#include "yb/client/table_creator.h"

#include "yb/client/table_handle.h"

#include "yb/client/transaction.h"

#include "yb/client/yb_op.h"

#include "yb/gutil/stl_util.h"

#include "yb/gutil/strings/join.h"

#include "yb/gutil/strings/substitute.h"

#include "yb/integration-tests/cdc_test_util.h"

#include "yb/integration-tests/mini_cluster.h"

#include "yb/integration-tests/twodc_test_base.h"

#include "yb/integration-tests/yb_mini_cluster_test_base.h"

#include "yb/master/cdc_consumer_registry_service.h"

#include "yb/master/mini_master.h"

#include "yb/master/master.h"

#include "yb/master/master_cluster.proxy.h"

#include "yb/master/master_ddl.proxy.h"

#include "yb/master/master_replication.proxy.h"

#include "yb/master/master-test-util.h"

#include "yb/master/sys_catalog_initialization.h"

#include "yb/rpc/rpc_controller.h"

#include "yb/tablet/tablet.h"

#include "yb/tablet/tablet_peer.h"

#include "yb/tserver/mini_tablet_server.h"

#include "yb/tserver/tablet_server.h"

#include "yb/tserver/ts_tablet_manager.h"

#include "yb/tserver/cdc_consumer.h"

#include "yb/util/atomic.h"

#include "yb/util/faststring.h"

#include "yb/util/format.h"

#include "yb/util/monotime.h"

#include "yb/util/random.h"

#include "yb/util/random_util.h"

#include "yb/util/result.h"

#include "yb/util/stopwatch.h"

#include "yb/util/test_util.h"

#include "yb/util/test_macros.h"

#include "yb/yql/pgwrapper/libpq_utils.h"

#include "yb/yql/pgwrapper/pg_wrapper.h"

DECLARE_int32(replication_factor);

DECLARE_int32(cdc_max_apply_batch_num_records);

DECLARE_int32(client_read_write_timeout_ms);

DECLARE_int32(pgsql_proxy_webserver_port);

DECLARE_bool(enable_ysql);

DECLARE_bool(hide_pg_catalog_table_creation_logs);

DECLARE_bool(master_auto_run_initdb);

DECLARE_int32(pggate_rpc_timeout_secs);

DECLARE_bool(enable_delete_truncate_xcluster_replicated_table);

namespace yb {

using client::YBClient;

using client::YBClientBuilder;

using client::YBColumnSchema;

using client::YBError;

using client::YBSchema;

using client::YBSchemaBuilder;

using client::YBSession;

using client::YBTable;

using client::YBTableAlterer;

using client::YBTableCreator;

using client::YBTableType;

using client::YBTableName;

using master::GetNamespaceInfoResponsePB;

using master::MiniMaster;

using tserver::MiniTabletServer;

using tserver::enterprise::CDCConsumer;

using pgwrapper::ToString;

using pgwrapper::GetInt32;

using pgwrapper::PGConn;

using pgwrapper::PGResultPtr;

using pgwrapper::PgSupervisor;

namespace enterprise {

constexpr static const char* const kKeyColumnName = "key";

class TwoDCYsqlTest : public TwoDCTestBase, public testing::WithParamInterface<TwoDCTestParams> {

 public:

  Status Initialize(uint32_t replication_factor, uint32_t num_masters = 1) {

    master::SetDefaultInitialSysCatalogSnapshotFlags();

    TwoDCTestBase::SetUp();

    FLAGS_enable_ysql = true;

    FLAGS_master_auto_run_initdb = true;

    FLAGS_hide_pg_catalog_table_creation_logs = true;

    FLAGS_pggate_rpc_timeout_secs = 120;

    FLAGS_cdc_max_apply_batch_num_records = GetParam().batch_size;

    FLAGS_cdc_enable_replicate_intents = GetParam().enable_replicate_intents;

    MiniClusterOptions opts;

    opts.num_tablet_servers = replication_factor;

    opts.num_masters = num_masters;

    FLAGS_replication_factor = replication_factor;

    opts.cluster_id = "producer";

    producer_cluster_.mini_cluster_ = std::make_unique<MiniCluster>(opts);

    RETURN_NOT_OK(producer_cluster()->StartSync());

    RETURN_NOT_OK(producer_cluster()->WaitForTabletServerCount(replication_factor));

    RETURN_NOT_OK(WaitForInitDb(producer_cluster()));

    opts.cluster_id = "consumer";

    consumer_cluster_.mini_cluster_ = std::make_unique<MiniCluster>(opts);

    RETURN_NOT_OK(consumer_cluster()->StartSync());

    RETURN_NOT_OK(consumer_cluster()->WaitForTabletServerCount(replication_factor));

    RETURN_NOT_OK(WaitForInitDb(consumer_cluster()));

    producer_cluster_.client_ = VERIFY_RESULT(producer_cluster()->CreateClient());

    consumer_cluster_.client_ = VERIFY_RESULT(consumer_cluster()->CreateClient());

    RETURN_NOT_OK(InitPostgres(&producer_cluster_));

    RETURN_NOT_OK(InitPostgres(&consumer_cluster_));

    return Status::OK();

  }

  Result<std::vector<std::shared_ptr<client::YBTable>>>

      SetUpWithParams(std::vector<uint32_t> num_consumer_tablets,

                      std::vector<uint32_t> num_producer_tablets,

                      uint32_t replication_factor,

                      uint32_t num_masters = 1,

                      bool colocated = false,

                      boost::optional<std::string> tablegroup_name = boost::none) {

    RETURN_NOT_OK(Initialize(replication_factor, num_masters));

    if (num_consumer_tablets.size() != num_producer_tablets.size()) {

      return STATUS(IllegalState,

                    Format("Num consumer tables: $0 num producer tables: $1 must be equal.",

                           num_consumer_tablets.size(), num_producer_tablets.size()));

    }

    RETURN_NOT_OK(CreateDatabase(&producer_cluster_, kNamespaceName, colocated));

    RETURN_NOT_OK(CreateDatabase(&consumer_cluster_, kNamespaceName, colocated));

    if (tablegroup_name.has_value()) {

      RETURN_NOT_OK(CreateTablegroup(&producer_cluster_, kNamespaceName, tablegroup_name.get()));

      RETURN_NOT_OK(CreateTablegroup(&consumer_cluster_, kNamespaceName, tablegroup_name.get()));

    }

    std::vector<YBTableName> tables;

    std::vector<std::shared_ptr<client::YBTable>> yb_tables;

    for (uint32_t i = 0; i < num_consumer_tablets.size(); i++) {

      RETURN_NOT_OK(CreateTable(i, num_producer_tablets[i], &producer_cluster_,

                                &tables, tablegroup_name, colocated));

      std::shared_ptr<client::YBTable> producer_table;

      RETURN_NOT_OK(producer_client()->OpenTable(tables[i * 2], &producer_table));

      yb_tables.push_back(producer_table);

      RETURN_NOT_OK(CreateTable(i, num_consumer_tablets[i], &consumer_cluster_,

                                &tables, tablegroup_name, colocated));

      std::shared_ptr<client::YBTable> consumer_table;

      RETURN_NOT_OK(consumer_client()->OpenTable(tables[(i * 2) + 1], &consumer_table));

      yb_tables.push_back(consumer_table);

    }

    return yb_tables;

  }

  Status InitPostgres(Cluster* cluster) {

    auto pg_ts = RandomElement(cluster->mini_cluster_->mini_tablet_servers());

    auto port = cluster->mini_cluster_->AllocateFreePort();

    yb::pgwrapper::PgProcessConf pg_process_conf =

        VERIFY_RESULT(yb::pgwrapper::PgProcessConf::CreateValidateAndRunInitDb(

            yb::ToString(Endpoint(pg_ts->bound_rpc_addr().address(), port)),

            pg_ts->options()->fs_opts.data_paths.front() + "/pg_data",

            pg_ts->server()->GetSharedMemoryFd()));

    pg_process_conf.master_addresses = pg_ts->options()->master_addresses_flag;

    pg_process_conf.force_disable_log_file = true;

    FLAGS_pgsql_proxy_webserver_port = cluster->mini_cluster_->AllocateFreePort();

    LOG(INFO) << "Starting PostgreSQL server listening on "

              << pg_process_conf.listen_addresses << ":" << pg_process_conf.pg_port << ", data: "

              << pg_process_conf.data_dir

              << ", pgsql webserver port: " << FLAGS_pgsql_proxy_webserver_port;

    cluster->pg_supervisor_ = std::make_unique<yb::pgwrapper::PgSupervisor>(pg_process_conf);

    RETURN_NOT_OK(cluster->pg_supervisor_->Start());

    cluster->pg_host_port_ = HostPort(pg_process_conf.listen_addresses, pg_process_conf.pg_port);

    return Status::OK();

  }

  Status CreateDatabase(Cluster* cluster,

                        const std::string& namespace_name = kNamespaceName,

                        bool colocated = false) {

    auto conn = EXPECT_RESULT(cluster->Connect());

    EXPECT_OK(conn.ExecuteFormat("CREATE DATABASE $0$1",

                                 namespace_name, colocated ? " colocated = true" : ""));

    return Status::OK();

  }

  Result<string> GetUniverseId(Cluster* cluster) {

    master::GetMasterClusterConfigRequestPB req;

    master::GetMasterClusterConfigResponsePB resp;

    master::MasterClusterProxy master_proxy(

        &cluster->client_->proxy_cache(),

        VERIFY_RESULT(cluster->mini_cluster_->GetLeaderMasterBoundRpcAddr()));

    rpc::RpcController rpc;

    rpc.set_timeout(MonoDelta::FromSeconds(kRpcTimeout));

    RETURN_NOT_OK(master_proxy.GetMasterClusterConfig(req, &resp, &rpc));

    if (resp.has_error()) {

      return STATUS(IllegalState, "Error getting cluster config");

    }

    return resp.cluster_config().cluster_uuid();

  }

  Result<YBTableName> CreateTable(Cluster* cluster,

                                  const std::string& namespace_name,

                                  const std::string& table_name,

                                  const boost::optional<std::string>& tablegroup_name,

                                  uint32_t num_tablets,

                                  bool colocated = false,

                                  const int table_oid = 0) {

    auto conn = EXPECT_RESULT(cluster->ConnectToDB(namespace_name));

    std::string table_oid_string = "";

    if (table_oid > 0) {

      // Need to turn on session flag to allow for CREATE WITH table_oid.

      EXPECT_OK(conn.Execute("set yb_enable_create_with_table_oid=true"));

      table_oid_string = Format("table_oid = $0", table_oid);

    }

    std::string query = Format("CREATE TABLE $0($1 int PRIMARY KEY) ", table_name, kKeyColumnName);

    // One cannot use tablegroup together with split into tablets.

    if (tablegroup_name.has_value()) {

      std::string with_clause =

          table_oid_string.empty() ? "" : Format("WITH ($0) ", table_oid_string);

      std::string tablegroup_clause = Format("TABLEGROUP $0", tablegroup_name.value());

      query += Format("$0$1", with_clause, tablegroup_clause);

    } else {

      std::string colocated_clause = Format("colocated = $0", colocated);

      std::string with_clause =

          table_oid_string.empty() ? colocated_clause

                                   : Format("$0, $1", table_oid_string, colocated_clause);

      query += Format("WITH ($0) SPLIT INTO $1 TABLETS", with_clause, num_tablets);

    }

    EXPECT_OK(conn.Execute(query));

    return GetTable(cluster, namespace_name, table_name);

  }

  Status CreateTable(uint32_t idx, uint32_t num_tablets, Cluster* cluster,

                     std::vector<YBTableName>* tables,

                     const boost::optional<std::string>& tablegroup_name,

                     bool colocated = false) {

    /*

     * If we either have tablegroup name or colocated flag

     * Generate table_oid based on index so that we have the same table_oid for producer/consumer.

     */

    const int table_oid = (tablegroup_name.has_value() || colocated) ? (idx + 1) * 111111 : 0;

    auto table = VERIFY_RESULT(CreateTable(cluster, kNamespaceName, Format("test_table_$0", idx),

                                           tablegroup_name, num_tablets, colocated, table_oid));

    tables->push_back(table);

    return Status::OK();

  }

  Result<YBTableName> GetTable(Cluster* cluster,

                               const std::string& namespace_name,

                               const std::string& table_name,

                               bool verify_table_name = true,

                               bool exclude_system_tables = true) {

    master::ListTablesRequestPB req;

    master::ListTablesResponsePB resp;

    req.set_name_filter(table_name);

    req.mutable_namespace_()->set_name(namespace_name);

    req.mutable_namespace_()->set_database_type(YQL_DATABASE_PGSQL);

    if (!exclude_system_tables) {

      req.set_exclude_system_tables(true);

      req.add_relation_type_filter(master::USER_TABLE_RELATION);

    }

    master::MasterDdlProxy master_proxy(

        &cluster->client_->proxy_cache(),

        VERIFY_RESULT(cluster->mini_cluster_->GetLeaderMiniMaster())->bound_rpc_addr());

    rpc::RpcController rpc;

    rpc.set_timeout(MonoDelta::FromSeconds(kRpcTimeout));

    RETURN_NOT_OK(master_proxy.ListTables(req, &resp, &rpc));

    if (resp.has_error()) {

      return STATUS(IllegalState, "Failed listing tables");

    }

    // Now need to find the table and return it.

    for (const auto& table : resp.tables()) {

      // If !verify_table_name, just return the first table.

      if (!verify_table_name ||

          (table.name() == table_name && table.namespace_().name() == namespace_name)) {

        YBTableName yb_table;

        yb_table.set_table_id(table.id());

        yb_table.set_namespace_id(table.namespace_().id());

        yb_table.set_pgschema_name(table.pgschema_name());

        return yb_table;

      }

    }

    return STATUS(IllegalState,

                  strings::Substitute("Unable to find table $0 in namespace $1",

                                      table_name, namespace_name));

  }

  /*

   * TODO (#11597): Given one is not able to get tablegroup ID by name, currently this works by

   * getting the first available tablegroup appearing in the namespace.

   */

  Result<TablegroupId> GetTablegroup(Cluster* cluster, const std::string& namespace_name) {

    // Lookup the namespace id from the namespace name.

    std::string namespace_id;

    {

      master::ListNamespacesRequestPB req;

      master::ListNamespacesResponsePB resp;

      master::MasterDdlProxy master_proxy(

          &cluster->client_->proxy_cache(),

          VERIFY_RESULT(cluster->mini_cluster_->GetLeaderMiniMaster())->bound_rpc_addr());

      rpc::RpcController rpc;

      rpc.set_timeout(MonoDelta::FromSeconds(kRpcTimeout));

      RETURN_NOT_OK(master_proxy.ListNamespaces(req, &resp, &rpc));

      if (resp.has_error()) {

        return STATUS(IllegalState, "Failed to get namespace info");

      }

      // Find and return the namespace id.

      bool namespaceFound = false;

      for (const auto& entry : resp.namespaces()) {

        if (entry.name() == namespace_name) {

          namespaceFound = true;

          namespace_id = entry.id();

          break;

        }

      }

      if (!namespaceFound) {

        return STATUS(IllegalState, "Failed to find namespace");

      }

    }

    master::ListTablegroupsRequestPB req;

    master::ListTablegroupsResponsePB resp;

    master::MasterDdlProxy master_proxy(

        &cluster->client_->proxy_cache(),

        VERIFY_RESULT(cluster->mini_cluster_->GetLeaderMiniMaster())->bound_rpc_addr());

    req.set_namespace_id(namespace_id);

    rpc::RpcController rpc;

    rpc.set_timeout(MonoDelta::FromSeconds(kRpcTimeout));

    RETURN_NOT_OK(master_proxy.ListTablegroups(req, &resp, &rpc));

    if (resp.has_error()) {

      return STATUS(IllegalState, "Failed listing tablegroups");

    }

    // Find and return the tablegroup.

    if (resp.tablegroups().empty()) {

      return STATUS(IllegalState,

                    Format("Unable to find tablegroup in namespace $0", namespace_name));

    }

    return resp.tablegroups()[0].id() + master::kTablegroupParentTableIdSuffix;

  }

  Status CreateTablegroup(Cluster* cluster,

                          const std::string& namespace_name,

                          const std::string& tablegroup_name) {

    auto conn = EXPECT_RESULT(cluster->ConnectToDB(namespace_name));

    EXPECT_OK(conn.ExecuteFormat("CREATE TABLEGROUP $0", tablegroup_name));

    return Status::OK();

  }

  void WriteWorkload(uint32_t start, uint32_t end, Cluster* cluster, const YBTableName& table,

                     bool delete_op = false) {

    auto conn = EXPECT_RESULT(cluster->ConnectToDB(table.namespace_name()));

    LOG(INFO) << "Writing " << end-start << (delete_op ? " deletes" : " inserts");

    for (uint32_t i = start; i < end; i++) {

      if (delete_op) {

        EXPECT_OK(conn.ExecuteFormat("DELETE FROM $0 WHERE $1 = $2",

                                     table.table_name(), kKeyColumnName, i));

      } else {

        // TODO(#6582) transactions currently don't work, so don't use ON CONFLICT DO NOTHING now.

        EXPECT_OK(conn.ExecuteFormat("INSERT INTO $0($1) VALUES ($2)", // ON CONFLICT DO NOTHING",

                                     table.table_name(), kKeyColumnName, i));

      }

    }

  }

  void WriteTransactionalWorkload(uint32_t start, uint32_t end, Cluster* cluster,

                                  const YBTableName& table) {

    auto conn = EXPECT_RESULT(cluster->ConnectToDB(table.namespace_name()));

    EXPECT_OK(conn.Execute("BEGIN"));

    for (uint32_t i = start; i < end; i++) {

      EXPECT_OK(conn.ExecuteFormat("INSERT INTO $0($1) VALUES ($2) ON CONFLICT DO NOTHING",

                                   table.table_name(), kKeyColumnName, i));

    }

    EXPECT_OK(conn.Execute("COMMIT"));

  }

  void DeleteWorkload(uint32_t start, uint32_t end, Cluster* cluster, const YBTableName& table) {

    WriteWorkload(start, end, cluster, table, true /* delete_op */);

  }

  PGResultPtr ScanToStrings(const YBTableName& table_name, Cluster* cluster) {

    auto conn = EXPECT_RESULT(cluster->ConnectToDB(table_name.namespace_name()));

    auto result =

        EXPECT_RESULT(conn.FetchFormat("SELECT * FROM $0 ORDER BY key", table_name.table_name()));

    return result;

  }

  Status VerifyWrittenRecords(const YBTableName& producer_table,

                              const YBTableName& consumer_table) {

    return LoggedWaitFor([=]() -> Result<bool> {

      auto producer_results = ScanToStrings(producer_table, &producer_cluster_);

      auto consumer_results = ScanToStrings(consumer_table, &consumer_cluster_);

      if (PQntuples(producer_results.get()) != PQntuples(consumer_results.get())) {

        return false;

      }

      for (int i = 0; i < PQntuples(producer_results.get()); ++i) {

        auto prod_val = EXPECT_RESULT(ToString(producer_results.get(), i, 0));

        auto cons_val = EXPECT_RESULT(ToString(consumer_results.get(), i, 0));

        if (prod_val != cons_val) {

          return false;

        }

      }

      return true;

    }, MonoDelta::FromSeconds(kRpcTimeout), "Verify written records");

  }

  Status VerifyNumRecords(const YBTableName& table, Cluster* cluster, int expected_size) {

    return LoggedWaitFor([=]() -> Result<bool> {

      auto results = ScanToStrings(table, cluster);

      return PQntuples(results.get()) == expected_size;

    }, MonoDelta::FromSeconds(kRpcTimeout), "Verify number of records");

  }

  Status DeleteTable(Cluster* cluster,

                     TableId* table_id /* = nullptr */) {

    RETURN_NOT_OK(cluster->client_->DeleteTable(*table_id));

    return Status::OK();

  }

  Status TruncateTable(Cluster* cluster,

                       std::vector<string> table_ids) {

    RETURN_NOT_OK(cluster->client_->TruncateTables(table_ids));

    return Status::OK();

  }

};

INSTANTIATE_TEST_CASE_P(TwoDCTestParams, TwoDCYsqlTest,

                        ::testing::Values(TwoDCTestParams(1, true), TwoDCTestParams(1, false),

                                          TwoDCTestParams(0, true), TwoDCTestParams(0, false)));

TEST_P(TwoDCYsqlTest, SetupUniverseReplication) {

  YB_SKIP_TEST_IN_TSAN();

  auto tables = ASSERT_RESULT(SetUpWithParams({8, 4}, {6, 6}, 3, 1, false /* colocated */));

  const string kUniverseId = ASSERT_RESULT(GetUniverseId(&producer_cluster_));

  std::vector<std::shared_ptr<client::YBTable>> producer_tables;

  // tables contains both producer and consumer universe tables (alternately).

  // Pick out just the producer tables from the list.

  producer_tables.reserve(tables.size() / 2);

  for (size_t i = 0; i < tables.size(); i += 2) {

    producer_tables.push_back(tables[i]);

  }

  ASSERT_OK(SetupUniverseReplication(

      producer_cluster(), consumer_cluster(), consumer_client(), kUniverseId, producer_tables));

  // Verify that universe was setup on consumer.

  master::GetUniverseReplicationResponsePB resp;

  ASSERT_OK(VerifyUniverseReplication(consumer_cluster(), consumer_client(), kUniverseId, &resp));

  ASSERT_EQ(resp.entry().producer_id(), kUniverseId);

  ASSERT_EQ(resp.entry().tables_size(), producer_tables.size());

  for (size_t i = 0; i < producer_tables.size(); i++) {

    ASSERT_EQ(resp.entry().tables(narrow_cast<int>(i)), producer_tables[i]->id());

  }

  // Verify that CDC streams were created on producer for all tables.

  for (size_t i = 0; i < producer_tables.size(); i++) {

    master::ListCDCStreamsResponsePB stream_resp;

    ASSERT_OK(GetCDCStreamForTable(producer_tables[i]->id(), &stream_resp));

    ASSERT_EQ(stream_resp.streams_size(), 1);

    ASSERT_EQ(stream_resp.streams(0).table_id().Get(0), producer_tables[i]->id());

  }

  ASSERT_OK(DeleteUniverseReplication(kUniverseId));

}

TEST_P(TwoDCYsqlTest, SimpleReplication) {

  YB_SKIP_TEST_IN_TSAN();

  constexpr int kNTabletsPerTable = 1;

  std::vector<uint32_t> tables_vector = {kNTabletsPerTable, kNTabletsPerTable};

  auto tables = ASSERT_RESULT(SetUpWithParams(tables_vector, tables_vector, 1));

  const string kUniverseId = ASSERT_RESULT(GetUniverseId(&producer_cluster_));

  // tables contains both producer and consumer universe tables (alternately).

  // Pick out just the producer tables from the list.

  std::vector<std::shared_ptr<client::YBTable>> producer_tables;

  std::vector<std::shared_ptr<client::YBTable>> consumer_tables;

  producer_tables.reserve(tables.size() / 2);

  consumer_tables.reserve(tables.size() / 2);

  for (size_t i = 0; i < tables.size(); ++i) {

    if (i % 2 == 0) {

      producer_tables.push_back(tables[i]);

    } else {

      consumer_tables.push_back(tables[i]);

    }

  }

  // 1. Write some data.

  for (const auto& producer_table : producer_tables) {

    LOG(INFO) << "Writing records for table " << producer_table->name().ToString();

    WriteWorkload(0, 100, &producer_cluster_, producer_table->name());

  }

  // Verify data is written on the producer.

  for (const auto& producer_table : producer_tables) {

    auto producer_results = ScanToStrings(producer_table->name(), &producer_cluster_);

    ASSERT_EQ(100, PQntuples(producer_results.get()));

    int result;

    for (int i = 0; i < 100; ++i) {

      result = ASSERT_RESULT(GetInt32(producer_results.get(), i, 0));

      ASSERT_EQ(i, result);

    }

  }

  // 2. Setup replication.

  ASSERT_OK(SetupUniverseReplication(producer_cluster(), consumer_cluster(), consumer_client(),

                                     kUniverseId, producer_tables));

  // 3. Verify everything is setup correctly.

  master::GetUniverseReplicationResponsePB get_universe_replication_resp;

  ASSERT_OK(VerifyUniverseReplication(consumer_cluster(), consumer_client(), kUniverseId,

      &get_universe_replication_resp));

  ASSERT_OK(CorrectlyPollingAllTablets(

      consumer_cluster(), narrow_cast<uint32_t>(tables_vector.size() * kNTabletsPerTable)));

  auto data_replicated_correctly = [&](int num_results) -> Result<bool> {

    for (const auto& consumer_table : consumer_tables) {

      LOG(INFO) << "Checking records for table " << consumer_table->name().ToString();

      auto consumer_results = ScanToStrings(consumer_table->name(), &consumer_cluster_);

      if (num_results != PQntuples(consumer_results.get())) {

        return false;

      }

      int result;

      for (int i = 0; i < num_results; ++i) {

        result = VERIFY_RESULT(GetInt32(consumer_results.get(), i, 0));

        if (i != result) {

          return false;

        }

      }

    }

    return true;

  };

  ASSERT_OK(WaitFor([&]() -> Result<bool> { return data_replicated_correctly(100); },

                    MonoDelta::FromSeconds(20), "IsDataReplicatedCorrectly"));

  // 4. Write more data.

  for (const auto& producer_table : producer_tables) {

    WriteWorkload(100, 105, &producer_cluster_, producer_table->name());

  }

  // 5. Make sure this data is also replicated now.

  ASSERT_OK(WaitFor([&]() { return data_replicated_correctly(105); },

                    MonoDelta::FromSeconds(20), "IsDataReplicatedCorrectly"));

}

TEST_P(TwoDCYsqlTest, SetupUniverseReplicationWithProducerBootstrapId) {

  YB_SKIP_TEST_IN_TSAN();

  constexpr int kNTabletsPerTable = 1;

  std::vector<uint32_t> tables_vector = {kNTabletsPerTable, kNTabletsPerTable};

  auto tables = ASSERT_RESULT(SetUpWithParams(tables_vector, tables_vector, 3));

  const string kUniverseId = ASSERT_RESULT(GetUniverseId(&producer_cluster_));

  auto* producer_leader_mini_master = ASSERT_RESULT(producer_cluster()->GetLeaderMiniMaster());

  auto producer_master_proxy = std::make_shared<master::MasterReplicationProxy>(

      &producer_client()->proxy_cache(),

      producer_leader_mini_master->bound_rpc_addr());

  std::unique_ptr<client::YBClient> client;

  std::unique_ptr<cdc::CDCServiceProxy> producer_cdc_proxy;

  client = ASSERT_RESULT(consumer_cluster()->CreateClient());

  producer_cdc_proxy = std::make_unique<cdc::CDCServiceProxy>(

      &client->proxy_cache(),

      HostPort::FromBoundEndpoint(producer_cluster()->mini_tablet_server(0)->bound_rpc_addr()));

  // tables contains both producer and consumer universe tables (alternately).

  // Pick out just the producer tables from the list.

  std::vector<std::shared_ptr<client::YBTable>> producer_tables;

  std::vector<std::shared_ptr<client::YBTable>> consumer_tables;

  producer_tables.reserve(tables.size() / 2);

  consumer_tables.reserve(tables.size() / 2);

  for (size_t i = 0; i < tables.size(); ++i) {

    if (i % 2 == 0) {

      producer_tables.push_back(tables[i]);

    } else {

      consumer_tables.push_back(tables[i]);

    }

  }

  // 1. Write some data so that we can verify that only new records get replicated.

  for (const auto& producer_table : producer_tables) {

    LOG(INFO) << "Writing records for table " << producer_table->name().ToString();

    WriteWorkload(0, 100, &producer_cluster_, producer_table->name());

  }

  SleepFor(MonoDelta::FromSeconds(10));

  cdc::BootstrapProducerRequestPB req;

  cdc::BootstrapProducerResponsePB resp;

  for (const auto& producer_table : producer_tables) {

    req.add_table_ids(producer_table->id());

  }

  rpc::RpcController rpc;

  ASSERT_OK(producer_cdc_proxy->BootstrapProducer(req, &resp, &rpc));

  ASSERT_FALSE(resp.has_error());

  ASSERT_EQ(resp.cdc_bootstrap_ids().size(), producer_tables.size());

  int table_idx = 0;

  for (const auto& bootstrap_id : resp.cdc_bootstrap_ids()) {

    LOG(INFO) << "Got bootstrap id " << bootstrap_id

              << " for table " << producer_tables[table_idx++]->name().table_name();

  }

  std::unordered_map<std::string, int> tablet_bootstraps;

  // Verify that for each of the table's tablets, a new row in cdc_state table with the returned

  // id was inserted.

  client::TableHandle table;

  client::YBTableName cdc_state_table(

      YQL_DATABASE_CQL, master::kSystemNamespaceName, master::kCdcStateTableName);

  ASSERT_OK(table.Open(cdc_state_table, producer_client()));

  // 2 tables with 8 tablets each.

  ASSERT_EQ(tables_vector.size() * kNTabletsPerTable, boost::size(client::TableRange(table)));

  int nrows = 0;

  for (const auto& row : client::TableRange(table)) {

    nrows++;

    string stream_id = row.column(0).string_value();

    tablet_bootstraps[stream_id]++;

    string checkpoint = row.column(2).string_value();

    auto s = OpId::FromString(checkpoint);

    ASSERT_OK(s);

    OpId op_id = *s;

    ASSERT_GT(op_id.index, 0);

    LOG(INFO) << "Bootstrap id " << stream_id

              << " for tablet " << row.column(1).string_value();

  }

  ASSERT_EQ(tablet_bootstraps.size(), producer_tables.size());

  // Check that each bootstrap id has 8 tablets.

  for (const auto& e : tablet_bootstraps) {

    ASSERT_EQ(e.second, kNTabletsPerTable);

  }

  // Map table -> bootstrap_id. We will need when setting up replication.

  std::unordered_map<TableId, std::string> table_bootstrap_ids;

  for (int i = 0; i < resp.cdc_bootstrap_ids_size(); i++) {

    table_bootstrap_ids[req.table_ids(i)] = resp.cdc_bootstrap_ids(i);

  }

  // 2. Setup replication.

  master::SetupUniverseReplicationRequestPB setup_universe_req;

  master::SetupUniverseReplicationResponsePB setup_universe_resp;

  setup_universe_req.set_producer_id(kUniverseId);

  string master_addr = producer_cluster()->GetMasterAddresses();

  auto hp_vec = ASSERT_RESULT(HostPort::ParseStrings(master_addr, 0));

  HostPortsToPBs(hp_vec, setup_universe_req.mutable_producer_master_addresses());

  setup_universe_req.mutable_producer_table_ids()->Reserve(

      narrow_cast<int>(producer_tables.size()));

  for (const auto& producer_table : producer_tables) {

    setup_universe_req.add_producer_table_ids(producer_table->id());

    const auto& iter = table_bootstrap_ids.find(producer_table->id());

    ASSERT_NE(iter, table_bootstrap_ids.end());

    setup_universe_req.add_producer_bootstrap_ids(iter->second);

  }

  auto* consumer_leader_mini_master = ASSERT_RESULT(consumer_cluster()->GetLeaderMiniMaster());

  master::MasterReplicationProxy master_proxy(

      &consumer_client()->proxy_cache(),

      consumer_leader_mini_master->bound_rpc_addr());

  rpc.Reset();

  rpc.set_timeout(MonoDelta::FromSeconds(kRpcTimeout));

  ASSERT_OK(master_proxy.SetupUniverseReplication(setup_universe_req, &setup_universe_resp, &rpc));

  ASSERT_FALSE(setup_universe_resp.has_error());

  // 3. Verify everything is setup correctly.

  master::GetUniverseReplicationResponsePB get_universe_replication_resp;

  ASSERT_OK(VerifyUniverseReplication(consumer_cluster(), consumer_client(), kUniverseId,

      &get_universe_replication_resp));

  ASSERT_OK(CorrectlyPollingAllTablets(

      consumer_cluster(), narrow_cast<uint32_t>(tables_vector.size() * kNTabletsPerTable)));

  // 4. Write more data.

  for (const auto& producer_table : producer_tables) {

    WriteWorkload(1000, 1005, &producer_cluster_, producer_table->name());

  }

  // 5. Verify that only new writes get replicated to consumer since we bootstrapped the producer

  // after we had already written some data, therefore the old data (whatever was there before we

  // bootstrapped the producer) should not be replicated.

  auto data_replicated_correctly = [&]() -> Result<bool> {

    for (const auto& consumer_table : consumer_tables) {

      LOG(INFO) << "Checking records for table " << consumer_table->name().ToString();

      auto consumer_results = ScanToStrings(consumer_table->name(), &consumer_cluster_);

      if (5 != PQntuples(consumer_results.get())) {

        return false;

      }

      int result;

      for (int i = 0; i < 5; ++i) {

        result = VERIFY_RESULT(GetInt32(consumer_results.get(), i, 0));

        if ((1000 + i) != result) {

          return false;

        }

      }

    }

    return true;

  };

  ASSERT_OK(WaitFor([&]() -> Result<bool> { return data_replicated_correctly(); },

                    MonoDelta::FromSeconds(20), "IsDataReplicatedCorrectly"));

}

TEST_P(TwoDCYsqlTest, ColocatedDatabaseReplication) {

  YB_SKIP_TEST_IN_TSAN();

  constexpr int kNTabletsPerColocatedTable = 1;

  constexpr int kNTabletsPerTable = 3;

  std::vector<uint32_t> tables_vector = {kNTabletsPerColocatedTable, kNTabletsPerColocatedTable};

  // Create two colocated tables on each cluster.

  auto colocated_tables =

      ASSERT_RESULT(SetUpWithParams(tables_vector, tables_vector, 3, 1, true /* colocated */));

  const string kUniverseId = ASSERT_RESULT(GetUniverseId(&producer_cluster_));

  // Also create an additional non-colocated table in each database.

  auto non_colocated_table = ASSERT_RESULT(CreateTable(&producer_cluster_,

                                                       kNamespaceName,

                                                       "test_table_2",

                                                       boost::none /* tablegroup */,

                                                       kNTabletsPerTable,

                                                       false /* colocated */));

  std::shared_ptr<client::YBTable> non_colocated_producer_table;

  ASSERT_OK(producer_client()->OpenTable(non_colocated_table, &non_colocated_producer_table));

  non_colocated_table = ASSERT_RESULT(CreateTable(&consumer_cluster_,

                                                  kNamespaceName,

                                                  "test_table_2",

                                                  boost::none /* tablegroup */,

                                                  kNTabletsPerTable,

                                                  false /* colocated */));

  std::shared_ptr<client::YBTable> non_colocated_consumer_table;

  ASSERT_OK(consumer_client()->OpenTable(non_colocated_table, &non_colocated_consumer_table));

  // colocated_tables contains both producer and consumer universe tables (alternately).

  // Pick out just the producer tables from the list.

  std::vector<std::shared_ptr<client::YBTable>> producer_tables;

  std::vector<std::shared_ptr<client::YBTable>> consumer_tables;

  std::vector<std::shared_ptr<client::YBTable>> colocated_producer_tables;

  std::vector<std::shared_ptr<client::YBTable>> colocated_consumer_tables;

  producer_tables.reserve(colocated_tables.size() / 2 + 1);

  consumer_tables.reserve(colocated_tables.size() / 2 + 1);

  colocated_producer_tables.reserve(colocated_tables.size() / 2);

  colocated_consumer_tables.reserve(colocated_tables.size() / 2);

  for (size_t i = 0; i < colocated_tables.size(); ++i) {

    if (i % 2 == 0) {

      producer_tables.push_back(colocated_tables[i]);

      colocated_producer_tables.push_back(colocated_tables[i]);

    } else {

      consumer_tables.push_back(colocated_tables[i]);

      colocated_consumer_tables.push_back(colocated_tables[i]);

    }

  }

  producer_tables.push_back(non_colocated_producer_table);

  consumer_tables.push_back(non_colocated_consumer_table);

  // 1. Write some data to all tables.

  for (const auto& producer_table : producer_tables) {

    LOG(INFO) << "Writing records for table " << producer_table->name().ToString();

    WriteWorkload(0, 100, &producer_cluster_, producer_table->name());

  }

  // 2. Setup replication for only the colocated tables.

  // Get the producer namespace id, so we can construct the colocated parent table id.

  GetNamespaceInfoResponsePB ns_resp;

  ASSERT_OK(producer_client()->GetNamespaceInfo("", kNamespaceName, YQL_DATABASE_PGSQL, &ns_resp));

  rpc::RpcController rpc;

  master::SetupUniverseReplicationRequestPB setup_universe_req;

  master::SetupUniverseReplicationResponsePB setup_universe_resp;

  setup_universe_req.set_producer_id(kUniverseId);

  string master_addr = producer_cluster()->GetMasterAddresses();

  auto hp_vec = ASSERT_RESULT(HostPort::ParseStrings(master_addr, 0));

  HostPortsToPBs(hp_vec, setup_universe_req.mutable_producer_master_addresses());

  // Only need to add the colocated parent table id.

  setup_universe_req.mutable_producer_table_ids()->Reserve(1);

  setup_universe_req.add_producer_table_ids(

      ns_resp.namespace_().id() + master::kColocatedParentTableIdSuffix);

  auto* consumer_leader_mini_master = ASSERT_RESULT(consumer_cluster()->GetLeaderMiniMaster());

  auto master_proxy = std::make_shared<master::MasterReplicationProxy>(

      &consumer_client()->proxy_cache(),

      consumer_leader_mini_master->bound_rpc_addr());

  rpc.Reset();

  rpc.set_timeout(MonoDelta::FromSeconds(kRpcTimeout));

  ASSERT_OK(master_proxy->SetupUniverseReplication(setup_universe_req, &setup_universe_resp, &rpc));

  ASSERT_FALSE(setup_universe_resp.has_error());

  // 3. Verify everything is setup correctly.

  master::GetUniverseReplicationResponsePB get_universe_replication_resp;

  ASSERT_OK(VerifyUniverseReplication(consumer_cluster(), consumer_client(), kUniverseId,

      &get_universe_replication_resp));

  ASSERT_OK(CorrectlyPollingAllTablets(consumer_cluster(), kNTabletsPerColocatedTable));

  // 4. Check that colocated tables are being replicated.

  auto data_replicated_correctly = [&](int num_results, bool onlyColocated) -> Result<bool> {

    auto &tables = onlyColocated ? colocated_consumer_tables : consumer_tables;

    for (const auto& consumer_table : tables) {

      LOG(INFO) << "Checking records for table " << consumer_table->name().ToString();

      auto consumer_results = ScanToStrings(consumer_table->name(), &consumer_cluster_);

      if (num_results != PQntuples(consumer_results.get())) {

        return false;

      }

      int result;

      for (int i = 0; i < num_results; ++i) {

        result = VERIFY_RESULT(GetInt32(consumer_results.get(), i, 0));

        if (i != result) {

          return false;

        }

      }

    }

    return true;

  };

  ASSERT_OK(WaitFor([&]() -> Result<bool> { return data_replicated_correctly(100, true); },

                    MonoDelta::FromSeconds(20), "IsDataReplicatedCorrectly"));

  // Ensure that the non colocated table is not replicated.

  auto non_coloc_results = ScanToStrings(non_colocated_consumer_table->name(), &consumer_cluster_);

  ASSERT_EQ(0, PQntuples(non_coloc_results.get()));

  // 5. Add the regular table to replication.

  // Prepare and send AlterUniverseReplication command.

  master::AlterUniverseReplicationRequestPB alter_req;

  master::AlterUniverseReplicationResponsePB alter_resp;

  alter_req.set_producer_id(kUniverseId);

  alter_req.add_producer_table_ids_to_add(non_colocated_producer_table->id());