[#28317] YSQL: Start transactions with correct locality in AcquireObjectLocks RPC handler

Summary:
Given the following transaction:
```sql
BEGIN;
INSERT INTO global_table ...;
COMMIT;
```
Without object locking, the first RPC of this transaction is the Perform. This sets
force_global_transaction in PgPerformOptionsPB based on the force_global_transaction GUC, and
also contains tablespace oids from which we are able to decide the correct locality to start
the transaction as (region-local, tablespace-local(tablespace), global).

But with object locking, the first RPC is to acquire an object lock, and this RPC today does not set
force_global_transaction, nor does it have any tablespace information. Instead, we just always
start as region-local, and then later promoted to global during the Perform RPC as needed.

Without object locking, we mask the latency of the CREATED heartbeat by use of the transaction pool,
but with object locking, we cannot mask the latency of the PROMOTED heartbeat, so this results in an
extra global roundtrip for global transactions.

Additionally, in the context of #11268, it also means that every transaction will be promoted to
global when tablespace locality is used, since region-local cannot promote to any tablespace-local,
and there is also no good default tablespace-local to use instead.

This is also the cause of several geo-partitioning unit tests failure when object locking is
enabled, for exactly the above reasons.

This revision passes force_global_transaction and tablespace oid for AcquireObjectLocks RPCs. This
information is then used to start transactions (if needed) under the appropriate locality.

**Upgrade/Downgrade safety**
N/A, only modifies protobufs for local TS <=> PG communication.
Jira: DB-17999

Test Plan:
Jenkins

Also see D47261 for Jenkins run with object locking on.

Reviewers: myang

Reviewed By: myang

Subscribers: bkolagani, yql, ybase

Tags: #jenkins-ready

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

Author: es1024

src/yb/tserver/pg_client.proto

@@ -1303,6 +1303,7 @@ message PgAcquireObjectLockRequestPB {
     uint64 relation_oid = 2;
     uint64 object_oid = 3;
     uint64 object_sub_oid = 4;
+    uint32 tablespace_oid = 5;
   }
   LockOId lock_oid = 3;
 

src/yb/tserver/pg_client_session.cc

@@ -123,6 +123,9 @@ DEFINE_test_flag(bool, perform_ignore_pg_is_region_local, false,
     "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.");
 
+DEFINE_test_flag(bool, force_initial_region_local, false,
+    "Force transaction to start as region-local initially.");
+
 DECLARE_bool(vector_index_dump_stats);
 DECLARE_bool(yb_enable_cdc_consistent_snapshot_streams);
 DECLARE_bool(ysql_enable_db_catalog_version_mode);
@@ -1183,7 +1186,10 @@ class TransactionProvider {
     if (!next_plain_) {
       return std::nullopt;
     }
-    auto res = NextTxnMetaForPlain(deadline, true /* is_for_release */);
+    // next_plain_ exists, so the locality passed here doesn't matter as we won't be creating
+    // a new transaction.
+    auto res = NextTxnMetaForPlain(
+        TransactionFullLocality::RegionLocal(), deadline, true /* is_for_release */);
     if (!res.ok()) {
       LOG(DFATAL) << "Unexpected error while fetching existing plain re-usable txn";
       return std::nullopt;
@@ -1192,13 +1198,11 @@ class TransactionProvider {
   }
 
   Result<TransactionMetadata> NextTxnMetaForPlain(
-      CoarseTimePoint deadline, bool is_for_release = false) {
+      TransactionFullLocality locality, CoarseTimePoint deadline, bool is_for_release = false) {
     client::internal::InFlightOpsGroupsWithMetadata ops_info;
     if (!next_plain_) {
-      VLOG_WITH_FUNC(1) << "requesting new transaction";
-      // TODO(#28317): We should figure out the correct locality and use it here. With table-level
-      // locks, this will prevent transactions from being started as global or tablespace-local(X).
-      auto txn = Build(TransactionFullLocality::RegionLocal(), deadline, {});
+      VLOG_WITH_FUNC(1) << "requesting new transaction of locality " << locality;
+      auto txn = Build(locality, deadline, {});
       // Don't execute txn->GetMetadata() here since the transaction is not iniatialized with
       // its full metadata yet, like isolation level.
       Synchronizer synchronizer;
@@ -2627,9 +2631,10 @@ class PgClientSession::Impl {
     if (setup_session_result.is_plain && setup_session_result.session_data.transaction) {
       RETURN_NOT_OK(setup_session_result.session_data.transaction->GetMetadata(deadline).get());
     }
+    auto locality = GetTargetTransactionLocality(data->req);
     auto txn_meta_res = setup_session_result.session_data.transaction
         ? setup_session_result.session_data.transaction->GetMetadata(deadline).get()
-        : NextObjectLockingTxnMeta(deadline);
+        : NextObjectLockingTxnMeta(locality, deadline);
     RETURN_NOT_OK(txn_meta_res);
     const auto lock_type = static_cast<TableLockType>(data->req.lock_type());
     VLOG_WITH_PREFIX_AND_FUNC(1)
@@ -2666,12 +2671,12 @@ class PgClientSession::Impl {
         txn_meta_res->status_tablet);
     AcquireObjectLockLocallyWithRetries(
         ts_lock_manager(), std::move(lock_req), deadline, std::move(callback),
-        [session_impl = this, txn = setup_session_result.session_data.transaction]
+        [session_impl = this, txn = setup_session_result.session_data.transaction, locality]
             (CoarseTimePoint deadline) -> Status {
           if (txn) {
             RETURN_NOT_OK(txn->metadata());
           } else {
-            RETURN_NOT_OK(session_impl->NextObjectLockingTxnMeta(deadline));
+            RETURN_NOT_OK(session_impl->NextObjectLockingTxnMeta(locality, deadline));
           }
           return Status::OK();
         });
@@ -2928,7 +2933,9 @@ class PgClientSession::Impl {
     std::tie(data->ops, data->vector_index_query) = VERIFY_RESULT(PrepareOperations(
         &data->req, session, &data->sidecars, tables, vector_index_query_data_,
         data->transaction != nullptr /* has_distributed_txn */,
-        make_lw_function([this, deadline] { return NextObjectLockingTxnMeta(deadline); }),
+        make_lw_function([this, locality, deadline] {
+          return NextObjectLockingTxnMeta(locality, deadline);
+        }),
         IsObjectLockingEnabled()));
     session->FlushAsync([this, data, trace, trace_created_locally,
                          start_time](client::FlushStatus* flush_status) {
@@ -3677,8 +3684,9 @@ class PgClientSession::Impl {
 
   [[nodiscard]] bool IsObjectLockingEnabled() const { return ts_lock_manager() != nullptr; }
 
-  Result<TransactionMetadata> NextObjectLockingTxnMeta(CoarseTimePoint deadline) {
-    auto txn_meta = VERIFY_RESULT(transaction_provider_.NextTxnMetaForPlain(deadline));
+  Result<TransactionMetadata> NextObjectLockingTxnMeta(
+      TransactionFullLocality locality, CoarseTimePoint deadline) {
+    auto txn_meta = VERIFY_RESULT(transaction_provider_.NextTxnMetaForPlain(locality, deadline));
     RegisterLockOwner(txn_meta.transaction_id, txn_meta.status_tablet);
     return txn_meta;
   }
@@ -3691,14 +3699,27 @@ class PgClientSession::Impl {
   }
 
   TransactionFullLocality GetTargetTransactionLocality(const PgPerformRequestPB& request) const {
-    auto tablespace_oids = request.ops() | std::views::transform(GetOpTablespaceOid);
+    return GetTargetTransactionLocality(
+        request.options(), request.ops() | std::views::transform(GetOpTablespaceOid));
+  }
+
+  TransactionFullLocality GetTargetTransactionLocality(
+      const PgAcquireObjectLockRequestPB& request) const {
+    return GetTargetTransactionLocality(
+        request.options(), std::views::single(request.lock_oid().tablespace_oid()));
+  }
+
+  TransactionFullLocality GetTargetTransactionLocality(
+      const PgPerformOptionsPB& options, std::ranges::range auto&& tablespace_oids) const {
+    if (PREDICT_FALSE(FLAGS_TEST_force_initial_region_local)) {
+      return TransactionFullLocality::RegionLocal();
+    }
 
-    if (FLAGS_use_tablespace_based_transaction_placement ||
-        request.options().force_tablespace_locality()) {
-      if (auto oid = request.options().force_tablespace_locality_oid()) {
+    if (FLAGS_use_tablespace_based_transaction_placement || options.force_tablespace_locality()) {
+      if (auto oid = options.force_tablespace_locality_oid()) {
         return TransactionFullLocality::TablespaceLocal(oid);
       }
-      return CalculateTablespaceBasedLocality(tablespace_oids);
+      return CalculateTablespaceBasedLocality(std::move(tablespace_oids));
     }
 
     // TODO: is_all_region_local() handles exactly two cases that tablespace oid check does not:
@@ -3711,10 +3732,10 @@ class PgClientSession::Impl {
     //    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()) {
+    if (!FLAGS_TEST_perform_ignore_pg_is_region_local && options.is_all_region_local()) {
       return TransactionFullLocality::RegionLocal();
     }
-    return CalculateRegionBasedLocality(tablespace_oids);
+    return CalculateRegionBasedLocality(std::move(tablespace_oids));
   }
 
   TransactionFullLocality CalculateRegionBasedLocality(

src/yb/yql/pggate/pg_client.cc

@@ -493,10 +493,12 @@ class PgClient::Impl : public BigDataFetcher {
  public:
   Impl(
       std::reference_wrapper<const WaitEventWatcher> wait_event_watcher,
-      std::atomic<uint64_t>& next_perform_op_serial_no)
+      std::atomic<uint64_t>& next_perform_op_serial_no,
+      const TablespaceMap& tablespace_map)
     : heartbeat_poller_(std::bind(&Impl::Heartbeat, this, false)),
       wait_event_watcher_(wait_event_watcher),
-      next_perform_op_serial_no_(next_perform_op_serial_no) {
+      next_perform_op_serial_no_(next_perform_op_serial_no),
+      tablespace_map_(tablespace_map) {
     tablet_server_count_cache_.fill(0);
   }
 
@@ -932,6 +934,15 @@ class PgClient::Impl : public BigDataFetcher {
     lock_oid->set_relation_oid(lock_id.relation_oid);
     lock_oid->set_object_oid(lock_id.object_oid);
     lock_oid->set_object_sub_oid(lock_id.object_sub_oid);
+    if (lock_id.relation_oid >= kPgFirstNormalObjectId) {
+      auto tablespace_itr = tablespace_map_.find(PgObjectId(lock_id.db_oid, lock_id.relation_oid));
+      if (tablespace_itr == tablespace_map_.end()) {
+        LOG(WARNING) << "Tablespace not found for db_oid=" << lock_id.db_oid
+                     << " relation_oid=" << lock_id.relation_oid;
+      } else {
+        lock_oid->set_tablespace_oid(tablespace_itr->second);
+      }
+    }
     req.set_lock_type(static_cast<tserver::ObjectLockMode>(mode));
 
     auto result_future = PrepareAndSend<AcquireObjectLockData>(
@@ -1744,6 +1755,8 @@ class PgClient::Impl : public BigDataFetcher {
   InterprocessMappedRegion big_mapped_region_;
   ThreadSafeArena object_locks_arena_;
   std::atomic<uint64_t>& next_perform_op_serial_no_;
+
+  const TablespaceMap& tablespace_map_;
 };
 
 std::string DdlMode::ToString() const {
@@ -1761,8 +1774,9 @@ void DdlMode::ToPB(tserver::PgFinishTransactionRequestPB_DdlModePB* dest) const
 
 PgClient::PgClient(
     std::reference_wrapper<const WaitEventWatcher> wait_event_watcher,
-    std::atomic<uint64_t>& seq_number)
-    : impl_(new Impl(wait_event_watcher, seq_number)) {}
+    std::atomic<uint64_t>& seq_number,
+    const TablespaceMap& tablespace_map)
+    : impl_(new Impl(wait_event_watcher, seq_number, tablespace_map)) {}
 
 PgClient::~PgClient() = default;
 

src/yb/yql/pggate/pg_client.h

@@ -167,7 +167,8 @@ class PgClient {
  public:
   PgClient(
       std::reference_wrapper<const WaitEventWatcher> wait_event_watcher,
-      std::atomic<uint64_t>& next_perform_op_serial_no);
+      std::atomic<uint64_t>& next_perform_op_serial_no,
+      const TablespaceMap& tablespace_map);
   ~PgClient();
 
   Status Start(rpc::ProxyCache* proxy_cache,

src/yb/yql/pggate/pg_session.cc

@@ -933,9 +933,6 @@ Result<PerformFuture> PgSession::Perform(BufferableOperations&& ops, PerformOpti
     }
   }
 
-  options.set_force_global_transaction(yb_force_global_transaction);
-  options.set_force_tablespace_locality(yb_force_tablespace_locality);
-  options.set_force_tablespace_locality_oid(yb_force_tablespace_locality_oid);
   options.set_is_all_region_local(std::all_of(
           ops.operations().begin(), ops.operations().end(),
           [](const auto& op) { return op->is_region_local(); }));

src/yb/yql/pggate/pg_txn_manager.cc

@@ -789,6 +789,10 @@ Status PgTxnManager::SetupPerformOptions(
   if (read_time_action && *read_time_action == ReadTimeAction::RESET) {
     options->mutable_read_time()->Clear();
   }
+
+  options->set_force_global_transaction(yb_force_global_transaction);
+  options->set_force_tablespace_locality(yb_force_tablespace_locality);
+  options->set_force_tablespace_locality_oid(yb_force_tablespace_locality_oid);
   return Status::OK();
 }
 

src/yb/yql/pggate/pggate.cc

@@ -666,7 +666,8 @@ PgApiImpl::PgApiImpl(
       }),
       pg_shared_data_(
           *init_postgres_info.shared_data, !init_postgres_info.parallel_leader_session_id),
-      pg_client_(wait_event_watcher_, pg_shared_data_->next_perform_op_serial_no),
+      pg_client_(
+          wait_event_watcher_, pg_shared_data_->next_perform_op_serial_no, tablespace_map_),
       clock_(new server::HybridClock()),
       enable_table_locking_(ShouldEnableTableLocks()),
       pg_txn_manager_(new PgTxnManager(&pg_client_, clock_, pg_callbacks_, enable_table_locking_)),

src/yb/yql/pggate/pggate.h

@@ -952,6 +952,8 @@ class PgApiImpl {
 
   const WaitEventWatcher wait_event_watcher_;
 
+  TablespaceMap tablespace_map_;
+
   PgSharedDataHolder pg_shared_data_;
 
   // TODO Rename to client_ when YBClient is removed.
@@ -973,7 +975,6 @@ class PgApiImpl {
   TupleIdBuilder tuple_id_builder_;
   BufferingSettings buffering_settings_;
   YbctidReaderProvider ybctid_reader_provider_;
-  TablespaceMap tablespace_map_;
   PgFKReferenceCache fk_reference_cache_;
   ExplicitRowLockBuffer explicit_row_lock_buffer_;
 

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

@@ -114,6 +114,10 @@ class GeoTransactionsTest : public GeoTransactionsTestBase {
     }
     auto insert_value = NextInsertValue();
     ASSERT_OK(conn.ExecuteFormat("SET force_global_transaction = $0", ToString(session_var)));
+    if (local_table) {
+      ASSERT_OK(WarmupTablespaceCache(conn, *local_table));
+    }
+    ASSERT_OK(WarmupTablespaceCache(conn, target_table));
     ASSERT_OK(conn.StartTransaction(IsolationLevel::SERIALIZABLE_ISOLATION));
     if (local_table) {
       ASSERT_OK(conn.ExecuteFormat(
@@ -178,6 +182,10 @@ class GeoTransactionsTest : public GeoTransactionsTestBase {
       ASSERT_OK(init_conn(conn));
     }
     ASSERT_OK(conn.ExecuteFormat("SET force_global_transaction = $0", ToString(session_var)));
+    if (local_table) {
+      ASSERT_OK(WarmupTablespaceCache(conn, *local_table));
+    }
+    ASSERT_OK(WarmupTablespaceCache(conn, target_table));
     for (size_t i = 0; i < num_aborts; ++i) {
       ASSERT_OK(conn.StartTransaction(IsolationLevel::SERIALIZABLE_ISOLATION));
       auto insert_value = NextInsertValue();
@@ -329,10 +337,7 @@ class GeoTransactionsTestTableLocksDisabled : public GeoTransactionsTest {
   }
 };
 
-// Fails when table-level locks are enabled due to #28317.
-TEST_F_EX(
-    GeoTransactionsTest, YB_DISABLE_TEST_IN_TSAN(TestTransactionTabletSelection),
-    GeoTransactionsTestTableLocksDisabled) {
+TEST_F(GeoTransactionsTest, YB_DISABLE_TEST_IN_TSAN(TestTransactionTabletSelection)) {
   constexpr int tables_per_region = 1;
 
   ANNOTATE_UNPROTECTED_WRITE(FLAGS_auto_create_local_transaction_tables) = false;
@@ -510,10 +515,7 @@ TEST_F(GeoTransactionsTest, YB_DISABLE_TEST_IN_TSAN(TestMultiRegionTransactionTa
       InsertToLocalFirst::kTrue, ExpectedLocality::kGlobal);
 }
 
-// Fails when table-level locks are enabled due to #28317.
-TEST_F_EX(
-    GeoTransactionsTest, YB_DISABLE_TEST_IN_TSAN(TestAutomaticLocalTransactionTableCreation),
-    GeoTransactionsTestTableLocksDisabled) {
+TEST_F(GeoTransactionsTest, YB_DISABLE_TEST_IN_TSAN(TestAutomaticLocalTransactionTableCreation)) {
   constexpr int tables_per_region = 1;
 
   ANNOTATE_UNPROTECTED_WRITE(FLAGS_auto_create_local_transaction_tables) = true;
@@ -1194,11 +1196,13 @@ TEST_F(GeoTransactionsTablespaceLocalityTest, TestSessionVariableOverride) {
       SetGlobalTransactionsGFlag::kTrue, SetGlobalTransactionSessionVar::kTrue,
       ExpectedLocality::kGlobal, force_tablespace_auto_select);
 
-  auto force_tablespace_bad_select = [](pgwrapper::PGConn& conn) -> Status {
+  auto force_tablespace_bad_select = [this](pgwrapper::PGConn& conn) -> Status {
     RETURN_NOT_OK(conn.Execute("SET yb_force_tablespace_locality = true"));
     // Normal user oids start at 16384, so this should not map to anything, and we can test
     // that it becomes global.
-    return conn.Execute("SET yb_force_tablespace_locality_oid = 1");
+    RETURN_NOT_OK(conn.Execute("SET yb_force_tablespace_locality_oid = 1"));
+    // Discard current transaction.
+    return WarmupTablespaceCache(conn, kTableName);
   };
 
   CheckSuccess(
@@ -1375,8 +1379,7 @@ TEST_F(GeoTransactionsTablespaceLocalityTest, TestAlterSetTablespace) {
       ExpectedLocality::kGlobal);
 }
 
-// Fails when table-level locks are enabled due to #28317.
-class GeoTransactionsWildcardTest : public GeoTransactionsTestTableLocksDisabled {
+class GeoTransactionsWildcardTest : public GeoTransactionsTest {
  protected:
   void SetupTablespaces() override {
     ANNOTATE_UNPROTECTED_WRITE(FLAGS_force_global_transactions) = true;