[#28828] docdb: Perform initial is-region-local checks on tserver side

Summary:
When using region local transactions, there are two times we check for region locality:
1. We perform a check in postgres code (YBCIsRegionLocal) that takes the relation's tablespace oid,
   calculates the tablespace cost, and checks this is at most the interzone cost. This value is
   propagated up through pggate and eventually passed to the tserver in the Perform RPC, where it is
   checked only if the Perform RPC starts a transaction.
2. We perform checks at the tablet level in YBTransaction to decide whether we need to promote the
   the transaction to global.

(1) does nothing when table-level locking is enabled, because the transaction is now always started
before the Perform RPC (always done during an AcquireObjectLocks RPC), forcing us to use (2)
instead, which has performance implications (#28317).

Additionally, as part of work for tablespace locality, we now pass tablespace oid for read/write
operations inside the Perform, and additionally have tablespace information on the tserver for any
tablespace with transaction tables. So the calculation for (1) can be done entirely on tserver side.

This revision adds calculation of region locality on the tserver side using the above information.
A future change will address #28317 by pushing tablespace oid up in AcquireObjectLocks RPC and
reusing this logic.

The old logic (YBCIsRegionLocal) is left intact for now to avoid severe performance regressions
during upgrade, since the new logic only works after finalizing an upgrade to a version with the
enable_tablespace_based_transaction_placement autoflag (this controls availability of tablespace
information on the tserver). It can be removed once we have a guarantee that all upgrades come
from a verison with this autoflag (2025.1.1.1/2025.1.2/2025.2.0+).

There is also a potential for performance regression if the auto_create_transaction_tables is off
(since the new calculation is based off of tablespaces with transaction tables), but this flag is on
by default and turning it off is ill-advised in general (manual creation of transaction tables via
yb-admin becomes required).
Jira: DB-18538

Test Plan: Jenkins

Reviewers: bkolagani

Reviewed By: bkolagani

Subscribers: yql, ybase

Tags: #jenkins-ready

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

Author: es1024

src/yb/client/transaction_manager.cc

@@ -77,6 +77,20 @@ namespace client {
 
 namespace {
 
+const CloudInfoPB& GetPlacementFromGFlags() {
+  static GoogleOnceType once = GOOGLE_ONCE_INIT;
+  static CloudInfoPB cloud_info;
+  auto set_placement_from_gflags = [](CloudInfoPB* cloud_info) {
+    cloud_info->set_placement_cloud(FLAGS_placement_cloud);
+    cloud_info->set_placement_region(FLAGS_placement_region);
+    cloud_info->set_placement_zone(FLAGS_placement_zone);
+  };
+  GoogleOnceInitArg(
+      &once, static_cast<void (*)(CloudInfoPB*)>(set_placement_from_gflags), &cloud_info);
+
+  return cloud_info;
+}
+
 // Cache of tablet ids of the global transaction table and any transaction tables with
 // the same placement.
 class TransactionTableState {
@@ -110,6 +124,7 @@ class TransactionTableState {
     std::lock_guard lock(mutex_);
     if (!initialized_.load() || status_tablets_version_ < new_version) {
       tablets_ = std::move(tablets);
+      tablespace_region_local_.clear();
       tablespace_contains_tablespace_.clear();
       has_region_local_tablets_.store(!tablets_.region_local_tablets.empty());
       status_tablets_version_ = new_version;
@@ -126,6 +141,32 @@ class TransactionTableState {
     return tablets_.tablespaces.contains(tablespace_oid);
   }
 
+  bool TablespaceIsRegionLocal(PgTablespaceOid tablespace_oid) {
+    if (tablespace_oid == kPgInvalidOid) {
+      return false;
+    }
+
+    {
+      SharedLock lock(mutex_);
+      auto itr = tablespace_region_local_.find(tablespace_oid);
+      if (itr != tablespace_region_local_.end()) {
+        return itr->second;
+      }
+    }
+
+    std::lock_guard lock(mutex_);
+    auto itr = tablets_.tablespaces.find(tablespace_oid);
+    if (itr == tablets_.tablespaces.end()) {
+      return false;
+    }
+
+    auto& placement = itr->second.placement_info;
+    return tablespace_region_local_.try_emplace(
+        tablespace_oid,
+        PlacementInfoContainsCloudInfo(placement, GetPlacementFromGFlags()) &&
+            !PlacementInfoSpansMultipleRegions(placement)).first->second;
+  }
+
   bool TablespaceContainsTablespace(PgTablespaceOid lhs, PgTablespaceOid rhs) {
     if (lhs == rhs) {
       return true;
@@ -241,6 +282,8 @@ class TransactionTableState {
 
   TransactionStatusTablets tablets_ GUARDED_BY(mutex_);
 
+  std::unordered_map<PgTablespaceOid, bool> tablespace_region_local_ GUARDED_BY(mutex_);
+
   struct TablespaceOidPair {
     PgTablespaceOid first;
     PgTablespaceOid second;
@@ -251,20 +294,6 @@ class TransactionTableState {
   std::unordered_map<TablespaceOidPair, bool> tablespace_contains_tablespace_ GUARDED_BY(mutex_);
 };
 
-const CloudInfoPB& GetPlacementFromGFlags() {
-  static GoogleOnceType once = GOOGLE_ONCE_INIT;
-  static CloudInfoPB cloud_info;
-  auto set_placement_from_gflags = [](CloudInfoPB* cloud_info) {
-    cloud_info->set_placement_cloud(FLAGS_placement_cloud);
-    cloud_info->set_placement_region(FLAGS_placement_region);
-    cloud_info->set_placement_zone(FLAGS_placement_zone);
-  };
-  GoogleOnceInitArg(
-      &once, static_cast<void (*)(CloudInfoPB*)>(set_placement_from_gflags), &cloud_info);
-
-  return cloud_info;
-}
-
 // Loads transaction tablets list to cache.
 class LoadStatusTabletsTask {
  public:
@@ -460,6 +489,10 @@ class TransactionManager::Impl {
     return table_state_.HasAnyRegionLocalStatusTablets();
   }
 
+  bool TablespaceIsRegionLocal(PgTablespaceOid tablespace_oid) {
+    return table_state_.TablespaceIsRegionLocal(tablespace_oid);
+  }
+
   bool TablespaceLocalTransactionsPossible(PgTablespaceOid tablespace_oid) {
     return table_state_.HasAnyTransactionLocalStatusTablets(tablespace_oid);
   }
@@ -560,6 +593,10 @@ bool TransactionManager::RegionLocalTransactionsPossible() {
   return impl_->RegionLocalTransactionsPossible();
 }
 
+bool TransactionManager::TablespaceIsRegionLocal(PgTablespaceOid tablespace_oid) {
+  return impl_->TablespaceIsRegionLocal(tablespace_oid);
+}
+
 bool TransactionManager::TablespaceLocalTransactionsPossible(PgTablespaceOid tablespace_oid) {
   return impl_->TablespaceLocalTransactionsPossible(tablespace_oid);
 }

src/yb/client/transaction_manager.h

@@ -66,6 +66,8 @@ class TransactionManager {
 
   bool RegionLocalTransactionsPossible();
 
+  bool TablespaceIsRegionLocal(PgTablespaceOid tablespace_oid);
+
   bool TablespaceLocalTransactionsPossible(PgTablespaceOid tablespace_oid);
 
   bool TablespaceContainsTablespace(PgTablespaceOid lhs, PgTablespaceOid rhs);

src/yb/master/master.cc

@@ -335,8 +335,8 @@ Status Master::RegisterServices() {
       FLAGS_master_svc_queue_length,
       std::make_shared<tserver::PgClientServiceImpl>(
           *master_tablet_server_, client_future(), clock(),
-          std::bind(&Master::TransactionPool, this), mem_tracker(), metric_entity(), messenger(),
-          fs_manager_->uuid(), options())));
+          std::bind(&Master::TransactionManager, this), std::bind(&Master::TransactionPool, this),
+          mem_tracker(), metric_entity(), messenger(), fs_manager_->uuid(), options())));
 
   return Status::OK();
 }

src/yb/tserver/pg_client_service.cc

@@ -540,14 +540,17 @@ class PgClientServiceImpl::Impl : public SessionProvider {
   explicit Impl(
       std::reference_wrapper<const TabletServerIf> tablet_server,
       const std::shared_future<client::YBClient*>& client_future,
-      const scoped_refptr<ClockBase>& clock, TransactionPoolProvider transaction_pool_provider,
+      const scoped_refptr<ClockBase>& clock,
+      TransactionManagerProvider transaction_manager_provider,
+      TransactionPoolProvider transaction_pool_provider,
       rpc::Messenger* messenger, const TserverXClusterContextIf* xcluster_context,
       PgMutationCounter* pg_node_level_mutation_counter, MetricEntity* metric_entity,
       const MemTrackerPtr& parent_mem_tracker, const std::string& permanent_uuid,
       const server::ServerBaseOptions& tablet_server_opts)
       : tablet_server_(tablet_server.get()),
         client_future_(client_future),
         clock_(clock),
+        transaction_manager_provider_(std::move(transaction_manager_provider)),
         transaction_pool_provider_(std::move(transaction_pool_provider)),
         messenger_(*messenger),
         table_cache_(client_future_),
@@ -577,7 +580,8 @@ class PgClientServiceImpl::Impl : public SessionProvider {
             .lock_owner_registry =
                 tablet_server_.ObjectLockSharedStateManager()
                     ? &tablet_server_.ObjectLockSharedStateManager()->registry()
-                    : nullptr},
+                    : nullptr,
+            .transaction_manager_provider = transaction_manager_provider_},
         cdc_state_table_(client_future_),
         txn_snapshot_manager_(
             instance_id_,
@@ -2715,6 +2719,7 @@ class PgClientServiceImpl::Impl : public SessionProvider {
   const TabletServerIf& tablet_server_;
   std::shared_future<client::YBClient*> client_future_;
   const scoped_refptr<ClockBase> clock_;
+  TransactionManagerProvider transaction_manager_provider_;
   TransactionPoolProvider transaction_pool_provider_;
   rpc::Messenger& messenger_;
   PgTableCache table_cache_;
@@ -2791,15 +2796,17 @@ class PgClientServiceImpl::Impl : public SessionProvider {
 PgClientServiceImpl::PgClientServiceImpl(
     std::reference_wrapper<const TabletServerIf> tablet_server,
     const std::shared_future<client::YBClient*>& client_future,
-    const scoped_refptr<ClockBase>& clock, TransactionPoolProvider transaction_pool_provider,
+    const scoped_refptr<ClockBase>& clock, TransactionManagerProvider transaction_manager_provider,
+    TransactionPoolProvider transaction_pool_provider,
     const std::shared_ptr<MemTracker>& parent_mem_tracker,
     const scoped_refptr<MetricEntity>& entity, rpc::Messenger* messenger,
     const std::string& permanent_uuid, const server::ServerBaseOptions& tablet_server_opts,
     const TserverXClusterContextIf* xcluster_context,
     PgMutationCounter* pg_node_level_mutation_counter)
     : PgClientServiceIf(entity),
       impl_(new Impl(
-          tablet_server, client_future, clock, std::move(transaction_pool_provider), messenger,
+          tablet_server, client_future, clock, std::move(transaction_manager_provider),
+          std::move(transaction_pool_provider), messenger,
           xcluster_context, pg_node_level_mutation_counter, entity.get(), parent_mem_tracker,
           permanent_uuid, tablet_server_opts)) {}
 

src/yb/tserver/pg_client_service.h

@@ -121,7 +121,9 @@ class PgClientServiceImpl : public PgClientServiceIf {
   explicit PgClientServiceImpl(
       std::reference_wrapper<const TabletServerIf> tablet_server,
       const std::shared_future<client::YBClient*>& client_future,
-      const scoped_refptr<ClockBase>& clock, TransactionPoolProvider transaction_pool_provider,
+      const scoped_refptr<ClockBase>& clock,
+      TransactionManagerProvider transaction_manager_provider,
+      TransactionPoolProvider transaction_pool_provider,
       const std::shared_ptr<MemTracker>& parent_mem_tracker,
       const scoped_refptr<MetricEntity>& entity, rpc::Messenger* messenger,
       const std::string& permanent_uuid, const server::ServerBaseOptions& tablet_server_opts,

src/yb/tserver/pg_client_session.cc

@@ -38,6 +38,7 @@
 #include "yb/client/table.h"
 #include "yb/client/table_alterer.h"
 #include "yb/client/transaction.h"
+#include "yb/client/transaction_manager.h"
 #include "yb/client/yb_op.h"
 
 #include "yb/common/common.pb.h"
@@ -117,6 +118,11 @@ DEFINE_test_flag(
     "Add several unknown tables while processing perfrom request. "
     "It is expected that opening of such tables will fail");
 
+DEFINE_test_flag(bool, perform_ignore_pg_is_region_local, false,
+    "Ignore the is_all_region_local field of PgPerformOptionsPB. The intended state is for "
+    "everything to work when this field is not set, as the field is only left in for upgrade from "
+    "older versions and to be removed in the future.");
+
 DECLARE_bool(vector_index_dump_stats);
 DECLARE_bool(yb_enable_cdc_consistent_snapshot_streams);
 DECLARE_bool(ysql_enable_db_catalog_version_mode);
@@ -1518,6 +1524,10 @@ template <>
   return DoTrackSharedMemoryPgMethodExecution(wait_state, metadata, "AcquireObjectLock");
 }
 
+PgTablespaceOid GetOpTablespaceOid(const PgPerformOpPB& op) {
+  return op.has_write() ? op.write().tablespace_oid() : op.read().tablespace_oid();
+}
+
 } // namespace
 
 class PgClientSession::Impl {
@@ -3681,23 +3691,52 @@ class PgClientSession::Impl {
   }
 
   TransactionFullLocality GetTargetTransactionLocality(const PgPerformRequestPB& request) const {
-    if (!FLAGS_use_tablespace_based_transaction_placement &&
-        !request.options().force_tablespace_locality()) {
-      return request.options().is_all_region_local()
-          ? TransactionFullLocality::RegionLocal() : TransactionFullLocality::Global();
+    auto tablespace_oids = request.ops() | std::views::transform(GetOpTablespaceOid);
+
+    if (FLAGS_use_tablespace_based_transaction_placement ||
+        request.options().force_tablespace_locality()) {
+      if (auto oid = request.options().force_tablespace_locality_oid()) {
+        return TransactionFullLocality::TablespaceLocal(oid);
+      }
+      return CalculateTablespaceBasedLocality(tablespace_oids);
     }
 
-    if (auto oid = request.options().force_tablespace_locality_oid()) {
-      return TransactionFullLocality::TablespaceLocal(oid);
+    // TODO: is_all_region_local() handles exactly two cases that tablespace oid check does not:
+    // 1. until upgrade is finalized (enable_tablespace_based_transaction_placement autoflag on
+    //    master), tablespace oid check does not work, so it is needed to avoid global latencies
+    //    during the upgrade. This is only for upgrades from before 2025.1.2/2025.2.0.
+    // 2. if auto_create_local_transaction_tables is OFF (not default), and transaction tables
+    //    are manually created, tablespace oid will not be mapped to a transaction table, and
+    //    TablespaceIsRegionLocal() returns false. But this case is not really supported, aside
+    //    from use setting up some unit tests.
+    // Once these are no longer of concern, is_all_region_local and corresponding code in
+    // pggate/pg can be removed.
+    if (!FLAGS_TEST_perform_ignore_pg_is_region_local && request.options().is_all_region_local()) {
+      return TransactionFullLocality::RegionLocal();
     }
+    return CalculateRegionBasedLocality(tablespace_oids);
+  }
+
+  TransactionFullLocality CalculateRegionBasedLocality(
+      std::ranges::range auto&& tablespace_oids) const {
+    auto& transaction_manager = context_.transaction_manager_provider();
+    bool all_region_local = transaction_manager.RegionLocalTransactionsPossible();
+    for (PgTablespaceOid oid : tablespace_oids) {
+      all_region_local = all_region_local && transaction_manager.TablespaceIsRegionLocal(oid);
+    }
+    return all_region_local
+        ? TransactionFullLocality::RegionLocal() : TransactionFullLocality::Global();
+  }
 
+  TransactionFullLocality CalculateTablespaceBasedLocality(
+      std::ranges::range auto&& tablespace_oids) const {
+    auto& transaction_manager = context_.transaction_manager_provider();
     PgTablespaceOid tablespace_oid = kInvalidOid;
-    for (const auto& op : request.ops()) {
-      PgTablespaceOid oid =
-          op.has_write() ? op.write().tablespace_oid() : op.read().tablespace_oid();
-      if (tablespace_oid == kInvalidOid) {
+    for (PgTablespaceOid oid : tablespace_oids) {
+      if (tablespace_oid == kInvalidOid ||
+          transaction_manager.TablespaceContainsTablespace(oid, tablespace_oid)) {
         tablespace_oid = oid;
-      } else if (tablespace_oid != oid) {
+      } else if (!transaction_manager.TablespaceContainsTablespace(tablespace_oid, oid)) {
         return TransactionFullLocality::Global();
       }
     }

src/yb/tserver/pg_client_session.h

@@ -98,6 +98,7 @@ struct PgClientSessionContext {
   const EventStatsPtr& stats_exchange_response_size;
   const std::string& instance_uuid;
   docdb::ObjectLockOwnerRegistry* lock_owner_registry;
+  const TransactionManagerProvider transaction_manager_provider;
 };
 
 class PgClientSession final {

src/yb/tserver/tablet_server.cc

@@ -679,6 +679,7 @@ Status TabletServer::RegisterServices() {
 
   auto pg_client_service_holder = std::make_shared<PgClientServiceHolder>(
         *this, tablet_manager_->client_future(), clock(),
+        std::bind(&TabletServer::TransactionManager, this),
         std::bind(&TabletServer::TransactionPool, this), mem_tracker(), metric_entity(),
         messenger(), permanent_uuid(), options(), xcluster_context_.get(),
         &pg_node_level_mutation_counter_);

src/yb/tserver/tserver_fwd.h

@@ -28,6 +28,7 @@ namespace yb {
 
 namespace client {
 
+class TransactionManager;
 class TransactionPool;
 class YBPgsqlOp;
 
@@ -69,6 +70,7 @@ enum class TabletServerServiceRpcMethodIndexes;
 YB_STRONGLY_TYPED_BOOL(AllowSplitTablet);
 
 using TSLocalLockManagerPtr = std::shared_ptr<TSLocalLockManager>;
+using TransactionManagerProvider = std::function<client::TransactionManager&()>;
 using TransactionPoolProvider = std::function<client::TransactionPool&()>;
 
 template <typename, typename = std::void_t<>>