[#27677] docdb: Support lock_timeout for object locks

Summary:
Added support for  lock_timeout for object locks. When taking table-level locks (like LOCK TABLE), the command will now respect the timeout setting and return an error if it can’t get the lock in time.

Note: Right now, only table-level locks support lock_timeout. Row-level locks  are not yet covered.
Jira: DB-17276

Test Plan: ./yb_build.sh --cxx-test pgwrapper_pg_object_locks-test --gtest_filter PgObjectLocksTest.VerifyLockTimeout

Reviewers: bkolagani

Reviewed By: bkolagani

Subscribers: rthallam, ybase, yql

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

Author: yusong-yan

src/yb/yql/pggate/pg_client.cc

@@ -225,6 +225,7 @@ ExchangeFuture<Data>::Result ExchangeFuture<Data>::get() {
 
 template <class LWReqPB, class LWRespPB, class ResTp, tserver::PgSharedExchangeReqType ShExcReqType>
 struct PgClientData : public FetchBigDataCallback {
+  using RequestType = LWReqPB;
   using ResultType = ResTp;
   static constexpr tserver::PgSharedExchangeReqType kSharedExchangeRequestType = ShExcReqType;
 
@@ -874,9 +875,10 @@ class PgClient::Impl : public BigDataFetcher {
       auto& exchange = session_shared_mem_->exchange();
       auto out = exchange.Obtain(kHeaderSize + kMetadataSize + data->req.SerializedSize());
       if (out) {
+        auto timeout = GetTimeout<typename Data::RequestType>();
         *reinterpret_cast<uint8_t *>(out) = Data::kSharedExchangeRequestType;
         out += sizeof(uint8_t);
-        LittleEndian::Store64(out, timeout_.ToMilliseconds());
+        LittleEndian::Store64(out, timeout.ToMilliseconds());
         out += sizeof(uint64_t);
         out = pointer_cast<std::byte*>(metadata.SerializeToArray(to_uchar_ptr(out)));
         const auto size = data->req.SerializedSize();
@@ -893,13 +895,14 @@ class PgClient::Impl : public BigDataFetcher {
           data->promise.set_value(MakeExchangeResult(*data, status));
           return data->promise.get_future();
         }
-        data->SetupExchange(&exchange, this, timeout_);
+        data->SetupExchange(&exchange, this, timeout);
         return ExchangeFuture<Data>(std::move(data));
       }
     }
     data->controller.set_invoke_callback_mode(rpc::InvokeCallbackMode::kReactorThread);
     (proxy_.get()->*method)(
-        data->req, &data->resp, SetupController(&data->controller), [data] {
+        data->req, &data->resp, SetupController<typename Data::RequestType>(&data->controller),
+        [data] {
           data->promise.set_value(MakeExchangeResult(*data, data->controller.CheckedResponse()));
         });
     return data->promise.get_future();
@@ -1633,16 +1636,23 @@ class PgClient::Impl : public BigDataFetcher {
     return Format("Session id $0: ", session_id_);
   }
 
+  template <typename T = void>
+  MonoDelta GetTimeout() {
+    if constexpr (std::is_same_v<T, tserver::PgAcquireAdvisoryLockRequestPB> ||
+                  std::is_same_v<T, tserver::LWPgAcquireObjectLockRequestPB>) {
+      return std::min(timeout_, lock_timeout_);
+    }
+    return timeout_;
+  }
+
   template <typename T = void>
   rpc::RpcController* SetupController(
       rpc::RpcController* controller,
       CoarseTimePoint deadline = CoarseTimePoint()) {
     if (deadline != CoarseTimePoint()) {
       controller->set_deadline(deadline);
-    } else if constexpr (std::is_same_v<T, tserver::PgAcquireAdvisoryLockRequestPB>) {
-      controller->set_timeout(std::min(timeout_, lock_timeout_));
     } else {
-      controller->set_timeout(timeout_);
+      controller->set_timeout(GetTimeout<T>());
     }
     return controller;
   }

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

@@ -59,6 +59,7 @@ DECLARE_bool(TEST_tserver_enable_ysql_lease_refresh);
 DECLARE_int64(olm_poll_interval_ms);
 DECLARE_string(vmodule);
 DECLARE_bool(ysql_enable_auto_analyze);
+DECLARE_int32(pg_client_extra_timeout_ms);
 
 using namespace std::literals;
 
@@ -817,6 +818,45 @@ TEST_F(PgObjectLocksTest, RetryExclusiveLockOnTserverLeaseRefresh) {
   ASSERT_OK(conn1.CommitTransaction());
 }
 
+TEST_F(PgObjectLocksTest, VerifyLockTimeout) {
+  constexpr int kLockTimeoutDurationMs = 3000;
+  constexpr int kExtraTimeoutMs = 1000;
+
+  auto setup_conn = ASSERT_RESULT(Connect());
+  ASSERT_OK(setup_conn.Execute("CREATE TABLE t (k INT PRIMARY KEY, v INT)"));
+  ASSERT_OK(setup_conn.Execute("INSERT INTO t VALUES (1, 1)"));
+
+  auto conn1 = ASSERT_RESULT(Connect());
+  auto conn2 = ASSERT_RESULT(Connect());
+
+  ASSERT_OK(conn1.StartTransaction(IsolationLevel::SNAPSHOT_ISOLATION));
+  ASSERT_OK(conn1.Execute("LOCK TABLE t IN ACCESS EXCLUSIVE MODE;"));
+
+  ASSERT_OK(conn2.ExecuteFormat("SET lock_timeout='$0ms';", kLockTimeoutDurationMs));
+  ASSERT_OK(conn2.StartTransaction(IsolationLevel::SNAPSHOT_ISOLATION));
+
+  std::future<Status> conn2_lock_future = std::async(std::launch::async, [&]() -> Status {
+    return conn2.Execute("LOCK TABLE t IN ACCESS EXCLUSIVE MODE;");
+  });
+
+  // Wait briefly to confirm that conn2 is actually blocking on the lock.
+  ASSERT_EQ(conn2_lock_future.wait_for(
+      std::chrono::seconds(1)), std::future_status::timeout);
+
+  auto expected_timeout_ms =
+      kLockTimeoutDurationMs + FLAGS_pg_client_extra_timeout_ms + kExtraTimeoutMs;
+  ASSERT_EQ(conn2_lock_future.wait_for(
+      std::chrono::milliseconds(expected_timeout_ms * kTimeMultiplier)), std::future_status::ready);
+
+  // Verify the lock attempt fails with lock timeout error.
+  Status result = conn2_lock_future.get();
+  ASSERT_NOK(result);
+  ASSERT_STR_CONTAINS(result.ToString(), "Timed out");
+
+  // Conn1 should still be able to commit
+  ASSERT_OK(conn1.CommitTransaction());
+}
+
 YB_STRONGLY_TYPED_BOOL(DoMasterFailover);
 YB_STRONGLY_TYPED_BOOL(UseExplicitLocksInsteadOfDdl);
 class PgObjecLocksTestOutOfOrderMessageHandling