[yugabyte#23744] docdb: Call TS count callback after loading sys catalog.

Summary:
After loading the persisted tablet server registry from the sys catalog, this diff adds logic to arrange for the ts count callback to be called if there are enough tablet servers in the registry.
Jira: DB-12649

Test Plan:
```
./yb_build.sh release --cxx-test-filter-re master_heartbeat-itest --cxx-test master_heartbeat-itest --gtest_filter GlobalTransactionTableCreationTest.CreateGlobalTransactionTableAfterFailover
```

Reviewers: asrivastava

Reviewed By: asrivastava

Subscribers: ybase, slingam

Differential Revision: https://phorge.dev.yugabyte.com/D40001

Author: druzac

src/yb/integration-tests/master_heartbeat-itest.cc

@@ -52,6 +52,8 @@
 #include "yb/util/flags.h"
 #include "yb/util/tostring.h"
 
+#include "yb/yql/pgwrapper/libpq_utils.h"
+
 using namespace std::literals;
 
 DECLARE_bool(enable_load_balancing);
@@ -672,4 +674,60 @@ TEST_F(MasterHeartbeatITestWithExternal, ReRegisterRemovedPeers) {
       original_uuids, 60s, "Wait for master to register original uuids"));
 }
 
+// This test class sets up a cluster in an inconsistent state.  The tservers have a placement uuid
+// set but the masters do not have the placement uuid in the cluster config. This prevents the
+// master leader from creating the global transaction status table after registering the tservers
+// which is required for a test case.
+class GlobalTransactionTableCreationTest : public YBTest {
+ public:
+  void SetUp() override;
+
+  void TearDown() override;
+
+  std::unique_ptr<ExternalMiniCluster> cluster_;
+  std::string placement_uuid_;
+};
+
+TEST_F(GlobalTransactionTableCreationTest, CreateGlobalTransactionTableAfterFailover) {
+  // In this test we validate the global transaction table is created by the callback scheduled by
+  // the ts manager at catalog load time when it detects there are enough tservers in the registry.
+  // Normally the master leader that first registers enough tservers will execute the
+  // callback. However because this test's setup sets a placement uuid for the tservers without
+  // setting it in the master's cluster config, the master leader will fail to create the global
+  // transaction table. After failover we fix the cluster config which unblocks the attempt to
+  // create the global transaction able.
+  ASSERT_OK(cluster_->StepDownMasterLeaderAndWaitForNewLeader());
+  // Sanity check that we cannot create a table yet.
+  std::string stmt = "CREATE TABLE test_table (k INT PRIMARY KEY, v INT)";
+  auto pgconn = ASSERT_RESULT(cluster_->ConnectToDB("yugabyte"));
+  ASSERT_NOK(pgconn.ExecuteFormat(stmt));
+  master::MasterClusterClient cluster_client(
+      cluster_->GetLeaderMasterProxy<master::MasterClusterProxy>());
+  auto config = ASSERT_RESULT(cluster_client.GetMasterClusterConfig());
+  config.mutable_replication_info()->mutable_live_replicas()->set_placement_uuid(placement_uuid_);
+  ASSERT_OK(cluster_client.ChangeMasterClusterConfig(std::move(config)));
+
+  // Now try to create a table. Table creation through pg will fail unless the transaction table
+  // already exists.
+  ASSERT_OK(WaitFor([&pgconn, &stmt]() -> Result<bool> {
+        return pgconn.ExecuteFormat(stmt).ok();
+      },
+      MonoDelta::FromSeconds(60), "Could not create table"));
+}
+
+void GlobalTransactionTableCreationTest::SetUp() {
+  placement_uuid_ = Uuid::Generate().ToString();
+  auto opts = ExternalMiniClusterOptions();
+  opts.num_masters = 3;
+  opts.num_tablet_servers = 3;
+  opts.enable_ysql = true;
+  opts.extra_tserver_flags = {Format("--placement_uuid=$0", placement_uuid_)};
+  cluster_ = std::make_unique<ExternalMiniCluster>(opts);
+  ASSERT_OK(cluster_->Start());
+}
+
+void GlobalTransactionTableCreationTest::TearDown() {
+  cluster_->Shutdown();
+}
+
 }  // namespace yb::integration_tests

src/yb/master/CMakeLists.txt

@@ -76,6 +76,7 @@ set(MASTER_SRCS
   catalog_entity_info.cc
   catalog_entity_tasks.cc
   catalog_loaders.cc
+  catalog_loading_state.cc
   catalog_manager.cc
   catalog_manager_bg_tasks.cc
   catalog_manager_ext.cc

src/yb/master/catalog_loaders.h

@@ -36,42 +36,15 @@
 
 #include <boost/preprocessor/cat.hpp>
 
-#include "yb/master/master_fwd.h"
+#include "yb/master/catalog_loading_state.h"
 #include "yb/master/catalog_manager.h"
+#include "yb/master/master_fwd.h"
 #include "yb/master/permissions_manager.h"
 #include "yb/master/sys_catalog.h"
 
 namespace yb {
 namespace master {
 
-struct SysCatalogLoadingState {
-  std::unordered_map<TableId, std::vector<TableId>> parent_to_child_tables;
-  std::vector<std::pair<std::function<void()>, std::string>> post_load_tasks;
-
-  // The tables which require their memory state to write to disk.
-  TableIdSet write_to_disk_tables;
-
-  // Index tables which require backfill status validation (by checking GC delete markers state).
-  // The tables are grouped by the indexed table id for performance reasons.
-  std::unordered_map<TableId, TableIdSet> validate_backfill_status_index_tables;
-
-  const LeaderEpoch epoch;
-
-  explicit SysCatalogLoadingState(LeaderEpoch leader_epoch = {}) : epoch(std::move(leader_epoch))
-  {}
-
-  void AddPostLoadTask(std::function<void()>&& func, std::string&& msg) {
-    post_load_tasks.push_back({std::move(func), std::move(msg)});
-  }
-
-  void Reset() {
-    parent_to_child_tables.clear();
-    post_load_tasks.clear();
-    write_to_disk_tables.clear();
-    validate_backfill_status_index_tables.clear();
-  }
-};
-
 #define DECLARE_LOADER_CLASS(name, key_type, entry_pb_name, mutex) \
   class BOOST_PP_CAT(name, Loader) : \
       public Visitor<BOOST_PP_CAT(BOOST_PP_CAT(Persistent, name), Info)> { \

src/yb/master/catalog_loading_state.cc

@@ -0,0 +1,31 @@
+// Copyright (c) YugabyteDB, 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 "yb/master/catalog_loading_state.h"
+
+namespace yb::master {
+
+SysCatalogLoadingState::SysCatalogLoadingState(LeaderEpoch leader_epoch)
+    : epoch(std::move(leader_epoch)) {}
+
+void SysCatalogLoadingState::AddPostLoadTask(std::function<void()>&& func, std::string&& msg) {
+  post_load_tasks.push_back({std::move(func), std::move(msg)});
+}
+
+void SysCatalogLoadingState::Reset() {
+  parent_to_child_tables.clear();
+  post_load_tasks.clear();
+  write_to_disk_tables.clear();
+  validate_backfill_status_index_tables.clear();
+}
+}  // namespace yb::master

src/yb/master/catalog_loading_state.h

@@ -0,0 +1,41 @@
+// Copyright (c) YugabyteDB, 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.
+//
+
+#pragma once
+
+#include "yb/common/entity_ids_types.h"
+
+#include "yb/master/leader_epoch.h"
+
+namespace yb::master {
+
+struct SysCatalogLoadingState {
+  std::unordered_map<TableId, std::vector<TableId>> parent_to_child_tables;
+  std::vector<std::pair<std::function<void()>, std::string>> post_load_tasks;
+
+  // The tables which require their memory state to write to disk.
+  TableIdSet write_to_disk_tables;
+
+  // Index tables which require backfill status validation (by checking GC delete markers state).
+  // The tables are grouped by the indexed table id for performance reasons.
+  std::unordered_map<TableId, TableIdSet> validate_backfill_status_index_tables;
+
+  const LeaderEpoch epoch;
+
+  explicit SysCatalogLoadingState(LeaderEpoch leader_epoch = {});
+
+  void AddPostLoadTask(std::function<void()>&& func, std::string&& msg);
+
+  void Reset();
+};
+}  // namespace yb::master

src/yb/master/catalog_manager.cc

@@ -1320,31 +1320,6 @@ Status CatalogManager::VisitSysCatalog(SysCatalogLoadingState* state) {
     // it's important to end their tasks now.
     AbortAndWaitForAllTasksUnlocked();
 
-    // Clear internal maps and run data loaders.
-    RETURN_NOT_OK(RunLoaders(state));
-
-    // Prepare various default system configurations.
-    RETURN_NOT_OK(PrepareDefaultSysConfig(term));
-
-    RETURN_NOT_OK(MaybeRestoreInitialSysCatalogSnapshotAndReloadSysCatalog(state));
-
-    // Create the system namespaces (created only if they don't already exist).
-    RETURN_NOT_OK(PrepareDefaultNamespaces(term));
-
-    // Create the system tables (created only if they don't already exist).
-    RETURN_NOT_OK(PrepareSystemTables(state->epoch));
-
-    // Create the default cassandra (created only if they don't already exist).
-    RETURN_NOT_OK(permissions_manager_->PrepareDefaultRoles(term));
-
-    // If this is the first time we start up, we have no config information as default. We write an
-    // empty version 0.
-    RETURN_NOT_OK(PrepareDefaultClusterConfig(term));
-
-    RETURN_NOT_OK(xcluster_manager_->PrepareDefaultXClusterConfig(term, /* recreate = */ false));
-
-    permissions_manager_->BuildRecursiveRoles();
-
     if (FLAGS_enable_ysql) {
       // Number of TS to wait for before creating the txn table.
       auto wait_ts_count = std::max(FLAGS_txn_table_wait_min_ts_count, FLAGS_replication_factor);
@@ -1382,6 +1357,31 @@ Status CatalogManager::VisitSysCatalog(SysCatalogLoadingState* state) {
           });
     }
 
+    // Clear internal maps and run data loaders.
+    RETURN_NOT_OK(RunLoaders(state));
+
+    // Prepare various default system configurations.
+    RETURN_NOT_OK(PrepareDefaultSysConfig(term));
+
+    RETURN_NOT_OK(MaybeRestoreInitialSysCatalogSnapshotAndReloadSysCatalog(state));
+
+    // Create the system namespaces (created only if they don't already exist).
+    RETURN_NOT_OK(PrepareDefaultNamespaces(term));
+
+    // Create the system tables (created only if they don't already exist).
+    RETURN_NOT_OK(PrepareSystemTables(state->epoch));
+
+    // Create the default cassandra (created only if they don't already exist).
+    RETURN_NOT_OK(permissions_manager_->PrepareDefaultRoles(term));
+
+    // If this is the first time we start up, we have no config information as default. We write an
+    // empty version 0.
+    RETURN_NOT_OK(PrepareDefaultClusterConfig(term));
+
+    RETURN_NOT_OK(xcluster_manager_->PrepareDefaultXClusterConfig(term, /* recreate = */ false));
+
+    permissions_manager_->BuildRecursiveRoles();
+
     if (!VERIFY_RESULT(StartRunningInitDbIfNeeded(state->epoch))) {
       // If we are not running initdb, this is an existing cluster, and we need to check whether we
       // need to do a one-time migration to make YSQL system catalog tables transactional.
@@ -1491,8 +1491,8 @@ Status CatalogManager::RunLoaders(SysCatalogLoadingState* state) {
 
   RETURN_NOT_OK(xcluster_manager_->RunLoaders(hidden_tablets_));
   RETURN_NOT_OK(master_->clone_state_manager().ClearAndRunLoaders(state->epoch));
-  RETURN_NOT_OK(
-      master_->ts_manager()->RunLoader(master_->MakeCloudInfoPB(), &master_->proxy_cache()));
+  RETURN_NOT_OK(master_->ts_manager()->RunLoader(
+      master_->MakeCloudInfoPB(), &master_->proxy_cache(), *state));
 
   return Status::OK();
 }

src/yb/master/catalog_manager_bg_tasks.cc

@@ -366,10 +366,7 @@ void CatalogManagerBgTasks::Run() {
       }
 
       if (FLAGS_enable_ysql) {
-        // Start the tablespace background task.
         catalog_manager_->StartTablespaceBgTaskIfStopped();
-
-        // Start the pg catalog versions background task.
         catalog_manager_->StartPgCatalogVersionsBgTaskIfStopped();
       }
 

src/yb/master/ts_manager.cc

@@ -433,7 +433,8 @@ Status TSManager::MarkUnresponsiveTServers(const LeaderEpoch& epoch) {
   return Status::OK();
 }
 
-Status TSManager::RunLoader(const CloudInfoPB& cloud_info, rpc::ProxyCache* proxy_cache) {
+Status TSManager::RunLoader(
+    const CloudInfoPB& cloud_info, rpc::ProxyCache* proxy_cache, SysCatalogLoadingState& state) {
   if (!GetAtomicFlag(&FLAGS_persist_tserver_registry)) {
     return Status::OK();
   }
@@ -442,10 +443,10 @@ Status TSManager::RunLoader(const CloudInfoPB& cloud_info, rpc::ProxyCache* prox
   MutexLock l_reg(registration_lock_);
   std::lock_guard l_map(map_lock_);
   servers_by_id_ = loader->TakeMap();
-  // todo(zdrudi): cluster can be wedged (permanently?) if we crash after registering enough
-  // tservers but before calling the callback.  so add a check to potentially schedule the callback
-  // here.
-  // https://github.com/yugabyte/yugabyte-db/issues/23744
+  if (servers_by_id_.size() >= ts_count_callback_min_count_ && ts_count_callback_) {
+    state.AddPostLoadTask(
+        std::move(ts_count_callback_), "Callback run at minimum number of tservers");
+  }
   return Status::OK();
 }
 

src/yb/master/ts_manager.h

@@ -42,6 +42,7 @@
 #include "yb/gutil/macros.h"
 #include "yb/gutil/thread_annotations.h"
 
+#include "yb/master/catalog_loading_state.h"
 #include "yb/master/master_cluster.fwd.h"
 #include "yb/master/master_fwd.h"
 #include "yb/master/ts_descriptor.h"
@@ -160,7 +161,8 @@ class TSManager {
   // Transition all such TServers into the UNRESPONSIVE state.
   Status MarkUnresponsiveTServers(const LeaderEpoch& epoch);
 
-  Status RunLoader(const CloudInfoPB& cloud_info, rpc::ProxyCache* proxy_cache);
+  Status RunLoader(
+      const CloudInfoPB& cloud_info, rpc::ProxyCache* proxy_cache, SysCatalogLoadingState& state);
 
   Status RemoveTabletServer(
       const std::string& permanent_uuid, const BlacklistSet& blacklist,