[#11353] geo: Create transaction tables automatically with tablespace…

… update

Summary:
Changed the tablespace update procedure to start transaction table creation for tablespaces
which have placements set and at least one table in them (to account for tablespaces with bad
placements). This should handle both tablespaces from previous versions with no associated transaction
table, and tablespaces whose first table was through ALTER TABLE SET TABLESPACE. The transaction
table creation hook in CreateTable is left in as an optimization.

Test Plan: `ybd --cxx-test pgwrapper_geo_transactions-test --gtest_filter GeoTransactionsTest.TestAutomaticLocalTransactionTableCreationWithAlter`

Reviewers: dsrinivasan

Reviewed By: dsrinivasan

Subscribers: rthallam, bogdan, ybase

Differential Revision: https://phabricator.dev.yugabyte.com/D15601

src/yb/master/catalog_manager.cc

@@ -2153,6 +2153,33 @@ Result<shared_ptr<TableToTablespaceIdMap>> CatalogManager::GetYsqlTableToTablesp
   return table_to_tablespace_map;
 }
 
+Status CatalogManager::CreateTransactionStatusTablesForTablespaces(
+    const TablespaceIdToReplicationInfoMap& tablespace_info,
+    const TableToTablespaceIdMap& table_to_tablespace_map) {
+  if (!GetAtomicFlag(&FLAGS_enable_ysql_tablespaces_for_placement) ||
+      !GetAtomicFlag(&FLAGS_auto_create_local_transaction_tables)) {
+    return Status::OK();
+  }
+
+  std::unordered_set<TablespaceId> valid_tablespaces;
+  for (const auto& entry : table_to_tablespace_map) {
+    if (entry.second) {
+      valid_tablespaces.insert(*entry.second);
+    }
+  }
+  for (const auto& entry : tablespace_info) {
+    if (!entry.second) {
+      valid_tablespaces.erase(entry.first);
+    }
+  }
+
+  for (const auto& tablespace_id : valid_tablespaces) {
+    RETURN_NOT_OK(CreateLocalTransactionStatusTableIfNeeded(nullptr /* rpc */, tablespace_id));
+  }
+
+  return Status::OK();
+}
+
 void CatalogManager::StartTablespaceBgTaskIfStopped() {
   if (GetAtomicFlag(&FLAGS_ysql_tablespace_info_refresh_secs) <= 0 ||
       !GetAtomicFlag(&FLAGS_enable_ysql_tablespaces_for_placement)) {
@@ -2239,6 +2266,12 @@ Status CatalogManager::DoRefreshTablespaceInfo() {
                                                                   table_to_tablespace_map);
   }
 
+  if (table_to_tablespace_map) {
+    // Trigger transaction table creates for tablespaces with tables and no transaction tables.
+    RETURN_NOT_OK(CreateTransactionStatusTablesForTablespaces(
+        *tablespace_info, *table_to_tablespace_map));
+  }
+
   VLOG(3) << "Refreshed tablespace information in memory";
   return Status::OK();
 }

src/yb/master/catalog_manager.h

@@ -1280,6 +1280,10 @@ class CatalogManager :
                                     const Status& s,
                                     CreateTableResponsePB* resp);
 
+  CHECKED_STATUS CreateTransactionStatusTablesForTablespaces(
+      const TablespaceIdToReplicationInfoMap& tablespace_info,
+      const TableToTablespaceIdMap& table_to_tablespace_map);
+
   void StartTablespaceBgTaskIfStopped();
 
   std::shared_ptr<YsqlTablespaceManager> GetTablespaceManager() const;

src/yb/yql/pgwrapper/geo_transactions-test.cc

@@ -35,6 +35,10 @@
 
 DECLARE_int32(ysql_tablespace_info_refresh_secs);
 DECLARE_int32(TEST_nodes_per_cloud);
+DECLARE_int32(load_balancer_max_concurrent_adds);
+DECLARE_int32(load_balancer_max_concurrent_removals);
+DECLARE_int32(load_balancer_max_concurrent_moves);
+DECLARE_int32(load_balancer_max_concurrent_moves_per_table);
 DECLARE_bool(enable_ysql_tablespaces_for_placement);
 DECLARE_bool(force_global_transactions);
 DECLARE_bool(auto_create_local_transaction_tables);
@@ -58,6 +62,7 @@ YB_STRONGLY_TYPED_BOOL(WaitForHashChange);
 constexpr auto kDatabaseName = "yugabyte";
 constexpr auto kTablePrefix = "test";
 const auto kStatusTabletCacheRefreshTimeout = MonoDelta::FromMilliseconds(20000);
+const auto kWaitLoadBalancerTimeout = MonoDelta::FromMilliseconds(30000);
 
 } // namespace
 
@@ -80,6 +85,12 @@ class GeoTransactionsTest : public pgwrapper::PgMiniTestBase {
     ANNOTATE_UNPROTECTED_WRITE(FLAGS_TEST_nodes_per_cloud) = 5;
     // Reduce time spent waiting for tablespace refresh.
     ANNOTATE_UNPROTECTED_WRITE(FLAGS_ysql_tablespace_info_refresh_secs) = 1;
+    // We wait for the load balancer whenever it gets triggered anyways, so there's
+    // no concerns about the load balancer taking too many resources.
+    ANNOTATE_UNPROTECTED_WRITE(FLAGS_load_balancer_max_concurrent_adds) = 10;
+    ANNOTATE_UNPROTECTED_WRITE(FLAGS_load_balancer_max_concurrent_removals) = 10;
+    ANNOTATE_UNPROTECTED_WRITE(FLAGS_load_balancer_max_concurrent_moves) = 10;
+    ANNOTATE_UNPROTECTED_WRITE(FLAGS_load_balancer_max_concurrent_moves_per_table) = 10;
 
     pgwrapper::PgMiniTestBase::SetUp();
     client_ = ASSERT_RESULT(cluster_->CreateClient());
@@ -186,6 +197,37 @@ class GeoTransactionsTest : public pgwrapper::PgMiniTestBase {
     }
   }
 
+  void SetupTablesWithAlter() {
+    // Create tablespaces and tables.
+    ANNOTATE_UNPROTECTED_WRITE(FLAGS_force_global_transactions) = true;
+    auto conn = ASSERT_RESULT(Connect());
+    bool wait_for_version = ANNOTATE_UNPROTECTED_READ(FLAGS_auto_create_local_transaction_tables);
+    auto current_version = transaction_manager_->GetLoadedStatusTabletsVersion();
+    for (size_t i = 1; i <= NumTabletServers(); ++i) {
+      ASSERT_OK(conn.ExecuteFormat(R"#(
+          CREATE TABLESPACE tablespace$0 WITH (replica_placement='{
+            "num_replicas": 1,
+            "placement_blocks":[{
+              "cloud": "cloud0",
+              "region": "rack$0",
+              "zone": "zone",
+              "min_num_replicas": 1
+            }]
+          }')
+      )#", i));
+      ASSERT_OK(conn.ExecuteFormat(
+          "CREATE TABLE $0$1(value int)", kTablePrefix, i));
+      ASSERT_OK(conn.ExecuteFormat(
+          "ALTER TABLE $0$1 SET TABLESPACE tablespace$1", kTablePrefix, i));
+
+      WaitForLoadBalanceCompletion();
+      if (wait_for_version) {
+        WaitForStatusTabletsVersion(current_version + 1);
+        ++current_version;
+      }
+    }
+  }
+
   void DropTables() {
     // Drop tablespaces and tables.
     ANNOTATE_UNPROTECTED_WRITE(FLAGS_force_global_transactions) = true;
@@ -288,6 +330,17 @@ class GeoTransactionsTest : public pgwrapper::PgMiniTestBase {
         strings::Substitute(error, version)));
   }
 
+  void WaitForLoadBalanceCompletion() {
+    ASSERT_OK(WaitFor([&]() -> Result<bool> {
+      bool is_idle = VERIFY_RESULT(client_->IsLoadBalancerIdle());
+      return !is_idle;
+    }, kWaitLoadBalancerTimeout, "Timeout waiting for load balancer to start"));
+
+    ASSERT_OK(WaitFor([&]() -> Result<bool> {
+      return client_->IsLoadBalancerIdle();
+    }, kWaitLoadBalancerTimeout, "Timeout waiting for load balancer to go idle"));
+  }
+
  private:
   std::unique_ptr<YBClient> client_;
   TransactionManager* transaction_manager_;
@@ -494,5 +547,35 @@ TEST_F(GeoTransactionsTest, YB_DISABLE_TEST_IN_TSAN(TestAutomaticLocalTransactio
       ExpectedLocality::kGlobal);
 }
 
+TEST_F(GeoTransactionsTest,
+       YB_DISABLE_TEST_IN_TSAN(TestAutomaticLocalTransactionTableCreationWithAlter)) {
+  ANNOTATE_UNPROTECTED_WRITE(FLAGS_auto_create_local_transaction_tables) = true;
+  SetupTablesWithAlter();
+
+  // The case of connecting to TS2 with force_global_transactions = false will error out
+  // because it is a global transaction, see #10537.
+  CheckInsert(
+      1, SetGlobalTransactionsGFlag::kFalse, SetGlobalTransactionSessionVar::kFalse,
+      ExpectedLocality::kLocal);
+  CheckInsert(
+      1, SetGlobalTransactionsGFlag::kTrue, SetGlobalTransactionSessionVar::kFalse,
+      ExpectedLocality::kGlobal);
+  CheckInsert(
+      2, SetGlobalTransactionsGFlag::kTrue, SetGlobalTransactionSessionVar::kFalse,
+      ExpectedLocality::kGlobal);
+  CheckInsert(
+      1, SetGlobalTransactionsGFlag::kFalse, SetGlobalTransactionSessionVar::kTrue,
+      ExpectedLocality::kGlobal);
+  CheckInsert(
+      2, SetGlobalTransactionsGFlag::kFalse, SetGlobalTransactionSessionVar::kTrue,
+      ExpectedLocality::kGlobal);
+  CheckInsert(
+      1, SetGlobalTransactionsGFlag::kTrue, SetGlobalTransactionSessionVar::kTrue,
+      ExpectedLocality::kGlobal);
+  CheckInsert(
+      2, SetGlobalTransactionsGFlag::kTrue, SetGlobalTransactionSessionVar::kTrue,
+      ExpectedLocality::kGlobal);
+}
+
 } // namespace client
 } // namespace yb