[#27832] docdb: TServers try to relinquish their ysql lease on receiving SIGTERM

Summary:
On the shutdown path in tservers, make a best-effort attempt to release the ysql lease if one is currently held. This allows DDLs and DMLs on other tservers to proceed without waiting for the shut down TServer's ysql lease to expire normally.

**Upgrade/Rollback safety:**
Proto changes:

  - adds `lease_relinquished` to `SysObjectLockEntryPB::lease_info`. This field defaults to `false`, we expect upgraded servers to interpret its absence as `false`.
  - adds a new RPC, `RelinquishYsqlLease`, to the `MasterDdlService`

Test Plan:
```
./yb_build.sh release --cxx-test master-test --gtest_filter 'MasterTest.RelinquishLease*'
./yb_build.sh release --cxx-test object_lock-test --gtest_filter 'ExternalObjectLockTestLongLeaseTTL.SigTerm'
```

Reviewers: amitanand, bkolagani

Reviewed By: amitanand

Subscribers: slingam, rthallam, ybase

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

Author: druzac

src/yb/integration-tests/object_lock-test.cc

@@ -1135,6 +1135,29 @@ TEST_F(ExternalObjectLockTest, RefreshYsqlLease) {
   ASSERT_FALSE(info.has_ddl_lock_entries());
 }
 
+class ExternalObjectLockTestLongLeaseTTL : public ExternalObjectLockTest {
+  ExternalMiniClusterOptions MakeExternalMiniClusterOptions() override;
+  int ReplicationFactor() override;
+  size_t NumberOfTabletServers() override;
+};
+
+TEST_F(ExternalObjectLockTestLongLeaseTTL, SigTerm) {
+  auto kLeaseTimeoutDeadline = MonoDelta::FromSeconds(10);
+  ASSERT_LT(
+      kLeaseTimeoutDeadline.ToMilliseconds(),
+      std::stoll(ASSERT_RESULT(cluster_->GetFlag(
+          ASSERT_NOTNULL(cluster_->GetLeaderMaster()), "master_ysql_operation_lease_ttl_ms"))));
+  auto ts = tablet_server(0);
+  auto ts_uuid = ts->uuid();
+  ts->Shutdown(SafeShutdown::kTrue);
+  // TServer should lose its lease before the lease TTL.
+  // Use EXPECT here so we restart the TS in case of failure.
+  EXPECT_OK(WaitForTServerLeaseToExpire(ts_uuid, kLeaseTimeoutDeadline));
+  ASSERT_OK(ts->Restart());
+  // TServer should be able to re-acquire the lease after it restarts.
+  ASSERT_OK(WaitForTServerLease(ts_uuid, 10s));
+}
+
 class MultiMasterObjectLockTest : public ObjectLockTest {
  protected:
   int num_masters() override { return 3; }
@@ -1731,4 +1754,24 @@ ExternalObjectLockTestOneTSWithoutLease::MakeExternalMiniClusterOptions() {
   return opts;
 }
 
+ExternalMiniClusterOptions ExternalObjectLockTestLongLeaseTTL::MakeExternalMiniClusterOptions() {
+  ExternalMiniClusterOptions opts;
+  opts.num_tablet_servers = NumberOfTabletServers();
+  opts.replication_factor = ReplicationFactor();
+  opts.enable_ysql = true;
+  opts.extra_master_flags = {
+      Format("--master_ysql_operation_lease_ttl_ms=$0", 20 * 1000),
+      "--enable_load_balancing=false"};
+  opts.extra_tserver_flags = {
+      Format("--allowed_preview_flags_csv=$0,$1",
+             "enable_object_locking_for_table_locks", "ysql_yb_ddl_transaction_block_enabled"),
+      "--enable_object_locking_for_table_locks",
+      "--ysql_yb_ddl_transaction_block_enabled",
+      Format("--ysql_lease_refresher_interval_ms=$0", kDefaultYSQLLeaseRefreshIntervalMilli)};
+  return opts;
+}
+
+int ExternalObjectLockTestLongLeaseTTL::ReplicationFactor() { return 1; }
+size_t ExternalObjectLockTestLongLeaseTTL::NumberOfTabletServers() { return 1; }
+
 }  // namespace yb

src/yb/master/catalog_entity_info.cc

@@ -1360,9 +1360,24 @@ std::string DdlLogEntry::id() const {
 // ObjectLockInfo
 // ================================================================================================
 
-std::variant<ObjectLockInfo::WriteLock, SysObjectLockEntryPB::LeaseInfoPB>
+Result<std::variant<ObjectLockInfo::WriteLock, SysObjectLockEntryPB::LeaseInfoPB>>
 ObjectLockInfo::RefreshYsqlOperationLease(const NodeInstancePB& instance, MonoDelta lease_ttl) {
   auto l = LockForWrite();
+  auto& current_lease_info = l->pb.lease_info();
+  if (instance.instance_seqno() < current_lease_info.instance_seqno()) {
+    return STATUS_FORMAT(
+        IllegalState,
+        "Cannot grant lease, instance seqno of requestor $0 is lower than instance seqno of a "
+        "previously granted lease $1",
+        instance.instance_seqno(), current_lease_info.instance_seqno());
+  }
+  if (current_lease_info.lease_relinquished() &&
+      instance.instance_seqno() <= current_lease_info.instance_seqno()) {
+    return STATUS_FORMAT(
+        IllegalState,
+        "Cannot grant lease, lease has been relinquished by instance_seqno $0 already",
+        current_lease_info.instance_seqno());
+  }
   {
     std::lock_guard l(mutex_);
     // When doing this mutation we cannot be sure the tserver receives the response.
@@ -1377,6 +1392,7 @@ ObjectLockInfo::RefreshYsqlOperationLease(const NodeInstancePB& instance, MonoDe
   lease_info.set_live_lease(true);
   lease_info.set_lease_epoch(lease_info.lease_epoch() + 1);
   lease_info.set_instance_seqno(instance.instance_seqno());
+  lease_info.set_lease_relinquished(false);
   return std::move(l);
 }
 

src/yb/master/catalog_entity_info.h

@@ -1102,7 +1102,7 @@ class ObjectLockInfo : public MetadataCowWrapper<PersistentObjectLockInfo> {
   // Return the user defined type's ID. Does not require synchronization.
   virtual const std::string& id() const override { return ts_uuid_; }
 
-  std::variant<ObjectLockInfo::WriteLock, SysObjectLockEntryPB::LeaseInfoPB>
+  Result<std::variant<ObjectLockInfo::WriteLock, SysObjectLockEntryPB::LeaseInfoPB>>
   RefreshYsqlOperationLease(const NodeInstancePB& instance, MonoDelta lease_ttl) EXCLUDES(mutex_);
 
   virtual void Load(const SysObjectLockEntryPB& metadata) override;

src/yb/master/catalog_entity_info.proto

@@ -452,6 +452,8 @@ message SysObjectLockEntryPB {
 
     // The instance_seqno of the tserver process to which this lease was granted.
     optional int64 instance_seqno = 3;
+
+    optional bool lease_relinquished = 4;
   }
   optional LeaseInfoPB lease_info = 2;
   map<uint64, tserver.ReleaseObjectLockRequestPB> in_progress_release_request = 3;

src/yb/master/catalog_manager.cc

@@ -5957,6 +5957,13 @@ Status CatalogManager::RefreshYsqlLease(const RefreshYsqlLeaseRequestPB* req,
   return object_lock_info_manager_->RefreshYsqlLease(*req, *resp, *rpc, epoch);
 }
 
+Status CatalogManager::RelinquishYsqlLease(const RelinquishYsqlLeaseRequestPB* req,
+                                           RelinquishYsqlLeaseResponsePB* resp,
+                                           rpc::RpcContext* rpc,
+                                           const LeaderEpoch& epoch) {
+  return object_lock_info_manager_->RelinquishYsqlLease(*req, *resp, *rpc, epoch);
+}
+
 Status CatalogManager::TruncateTable(const TableId& table_id,
                                      TruncateTableResponsePB* resp,
                                      rpc::RpcContext* rpc,

src/yb/master/catalog_manager.h

@@ -382,14 +382,19 @@ class CatalogManager : public CatalogManagerIf, public SnapshotCoordinatorContex
                        rpc::RpcContext* rpc,
                        const LeaderEpoch& epoch);
 
+  // Get the information about an in-progress truncate operation.
+  Status IsTruncateTableDone(const IsTruncateTableDoneRequestPB* req,
+                             IsTruncateTableDoneResponsePB* resp);
+
   Status RefreshYsqlLease(const RefreshYsqlLeaseRequestPB* req,
                           RefreshYsqlLeaseResponsePB* resp,
                           rpc::RpcContext* rpc,
                           const LeaderEpoch& epoch);
 
-  // Get the information about an in-progress truncate operation.
-  Status IsTruncateTableDone(const IsTruncateTableDoneRequestPB* req,
-                             IsTruncateTableDoneResponsePB* resp);
+  Status RelinquishYsqlLease(const RelinquishYsqlLeaseRequestPB* req,
+                             RelinquishYsqlLeaseResponsePB* resp,
+                             rpc::RpcContext* rpc,
+                             const LeaderEpoch& epoch);
 
   // Backfill the specified index.  Currently only supported for YSQL.  YCQL does not need this as
   // master automatically runs backfill according to the DocDB permissions.

src/yb/master/master-test.cc

@@ -127,7 +127,8 @@ class MasterTest : public MasterTestBase {
   Result<scoped_refptr<NamespaceInfo>> FindNamespaceByName(
       YQLDatabase db_type, const std::string& name);
 
-  Result<TSHeartbeatResponsePB> SendNewTSRegistrationHeartbeat(const std::string& uuid);
+  Result<TSHeartbeatResponsePB> SendNewTSRegistrationHeartbeat(
+      const std::string& uuid, int64_t instance_seqno);
 
  private:
   // Used by SendNewTSRegistrationHeartbeat to avoid host port collisions.
@@ -2775,16 +2776,16 @@ TEST_F(MasterTest, RefreshYsqlLeaseWithoutRegistration) {
 
 TEST_F(MasterTest, RefreshYsqlLease) {
   ANNOTATE_UNPROTECTED_WRITE(FLAGS_enable_ysql_operation_lease) = true;
-  const std::string kTsUUID1 = "ts-uuid1";
-  const std::string kTsUUID2 = "ts-uuid2";
+  const std::string kTsUUID = "ts-uuid1";
+  constexpr uint64_t kSeqno = 1;
 
-  auto reg_resp1 = ASSERT_RESULT(SendNewTSRegistrationHeartbeat(kTsUUID1));
+  auto reg_resp1 = ASSERT_RESULT(SendNewTSRegistrationHeartbeat(kTsUUID, kSeqno));
   ASSERT_FALSE(reg_resp1.needs_reregister());
 
   auto ddl_client = MasterDDLClient{std::move(*proxy_ddl_)};
   auto lease_refresh_send_time_ms = MonoTime::Now().GetDeltaSinceMin().ToMilliseconds();
-  auto info = ASSERT_RESULT(ddl_client.RefreshYsqlLease(
-      kTsUUID1, /* instance_seqno */ 1, lease_refresh_send_time_ms, {}));
+  auto info =
+      ASSERT_RESULT(ddl_client.RefreshYsqlLease(kTsUUID, kSeqno, lease_refresh_send_time_ms, {}));
   ASSERT_TRUE(info.new_lease());
   ASSERT_EQ(info.lease_epoch(), 1);
   ASSERT_GT(
@@ -2794,27 +2795,82 @@ TEST_F(MasterTest, RefreshYsqlLease) {
   // todo(zdrudi): but we need to do this and check the bootstrap entries...
   // Refresh lease again. Since we omitted current lease epoch, master leader should still say this
   // is a new lease.
-  info = ASSERT_RESULT(ddl_client.RefreshYsqlLease(
-      kTsUUID1, /* instance_seqno */ 1, lease_refresh_send_time_ms, {}));
+  info =
+      ASSERT_RESULT(ddl_client.RefreshYsqlLease(kTsUUID, kSeqno, lease_refresh_send_time_ms, {}));
   ASSERT_TRUE(info.new_lease());
   ASSERT_EQ(info.lease_epoch(), 1);
   ASSERT_GT(info.lease_expiry_time_ms(), lease_refresh_send_time_ms);
 
   // Refresh lease again. We included current lease epoch but it's incorrect.
-  info = ASSERT_RESULT(
-      ddl_client.RefreshYsqlLease(kTsUUID1, /* instance_seqno */ 1, lease_refresh_send_time_ms, 0));
+  info = ASSERT_RESULT(ddl_client.RefreshYsqlLease(kTsUUID, kSeqno, lease_refresh_send_time_ms, 0));
   ASSERT_TRUE(info.new_lease());
   ASSERT_EQ(info.lease_epoch(), 1);
   ASSERT_GT(info.lease_expiry_time_ms(), lease_refresh_send_time_ms);
 
   // Refresh lease again. Current lease epoch is correct so master leader should not set new lease
   // bit.
-  info = ASSERT_RESULT(
-      ddl_client.RefreshYsqlLease(kTsUUID1, /* instance_seqno */ 1, lease_refresh_send_time_ms, 1));
+  info = ASSERT_RESULT(ddl_client.RefreshYsqlLease(kTsUUID, kSeqno, lease_refresh_send_time_ms, 1));
   ASSERT_FALSE(info.new_lease());
   ASSERT_GT(info.lease_expiry_time_ms(), lease_refresh_send_time_ms);
 }
 
+TEST_F(MasterTest, RelinquishLease) {
+  ANNOTATE_UNPROTECTED_WRITE(FLAGS_enable_ysql_operation_lease) = true;
+  const std::string kTsUUID = "ts-uuid1";
+  constexpr uint64_t kSeqno1 = 1;
+  auto reg_resp1 = ASSERT_RESULT(SendNewTSRegistrationHeartbeat(kTsUUID, kSeqno1));
+  ASSERT_FALSE(reg_resp1.needs_reregister());
+
+  auto ddl_client = MasterDDLClient{std::move(*proxy_ddl_)};
+  auto info1 = ASSERT_RESULT(ddl_client.RefreshYsqlLease(
+      kTsUUID, kSeqno1, MonoTime::Now().GetDeltaSinceMin().ToMilliseconds(), {}));
+  ASSERT_TRUE(info1.new_lease());
+  ASSERT_EQ(info1.lease_epoch(), 1);
+
+  ASSERT_OK(ddl_client.RelinquishYsqlLease(kTsUUID, kSeqno1));
+  auto list_ts = ASSERT_RESULT(cluster_client_->ListTabletServers());
+  ASSERT_EQ(list_ts.servers_size(), 1);
+  ASSERT_FALSE(list_ts.servers(0).lease_info().is_live());
+
+  // We should be able to get a new lease for a new instance of the ts.
+  constexpr uint64_t kSeqno2 = 2;
+  auto reg_resp2 = ASSERT_RESULT(SendNewTSRegistrationHeartbeat(kTsUUID, kSeqno2));
+  ASSERT_FALSE(reg_resp2.needs_reregister());
+  auto info2 = ASSERT_RESULT(ddl_client.RefreshYsqlLease(
+      kTsUUID, kSeqno2, MonoTime::Now().GetDeltaSinceMin().ToMilliseconds(), {}));
+  ASSERT_TRUE(info2.new_lease());
+  ASSERT_EQ(info2.lease_epoch(), 2);
+}
+
+TEST_F(MasterTest, RelinquishLeaseOfReplacedTS) {
+  ANNOTATE_UNPROTECTED_WRITE(FLAGS_enable_ysql_operation_lease) = true;
+  const std::string kTsUUID = "ts-uuid1";
+  constexpr uint64_t kSeqno1 = 1;
+  auto reg_resp1 = ASSERT_RESULT(SendNewTSRegistrationHeartbeat(kTsUUID, kSeqno1));
+  ASSERT_FALSE(reg_resp1.needs_reregister());
+
+  auto ddl_client = MasterDDLClient{std::move(*proxy_ddl_)};
+  auto info1 = ASSERT_RESULT(ddl_client.RefreshYsqlLease(
+      kTsUUID, kSeqno1, MonoTime::Now().GetDeltaSinceMin().ToMilliseconds(), {}));
+  ASSERT_TRUE(info1.new_lease());
+  ASSERT_EQ(info1.lease_epoch(), 1);
+
+  // Simulate a second ts process replacing the first.
+  constexpr uint64_t kSeqno2 = 2;
+  auto reg_resp2 = ASSERT_RESULT(SendNewTSRegistrationHeartbeat(kTsUUID, kSeqno2));
+  ASSERT_FALSE(reg_resp2.needs_reregister());
+  auto info2 = ASSERT_RESULT(ddl_client.RefreshYsqlLease(
+      kTsUUID, kSeqno2, MonoTime::Now().GetDeltaSinceMin().ToMilliseconds(), {}));
+  ASSERT_TRUE(info2.new_lease());
+  ASSERT_EQ(info2.lease_epoch(), 2);
+
+  // Send relinquish lease request from the first ts instance.
+  auto status = ddl_client.RelinquishYsqlLease(kTsUUID, kSeqno1);
+  ASSERT_NOK(status);
+  ASSERT_STR_CONTAINS(
+      status.ToString(), "Relinquish lease request for a replaced tserver instance");
+}
+
 Result<TSHeartbeatResponsePB> MasterTest::SendHeartbeat(
     TSToMasterCommonPB common, std::optional<TSRegistrationPB> registration,
     std::optional<TabletReportPB> report) {
@@ -2839,7 +2895,8 @@ Result<TSHeartbeatResponsePB> MasterTest::SendHeartbeat(
   return resp;
 }
 
-Result<TSHeartbeatResponsePB> MasterTest::SendNewTSRegistrationHeartbeat(const std::string& uuid) {
+Result<TSHeartbeatResponsePB> MasterTest::SendNewTSRegistrationHeartbeat(
+    const std::string& uuid, int64_t instance_seqno) {
   TSRegistrationPB reg;
   *reg.mutable_common()->add_private_rpc_addresses() =
       MakeHostPortPB("localhost", 1000 + registered_ts_count_);
@@ -2849,7 +2906,7 @@ Result<TSHeartbeatResponsePB> MasterTest::SendNewTSRegistrationHeartbeat(const s
 
   TSToMasterCommonPB common;
   common.mutable_ts_instance()->set_permanent_uuid(uuid);
-  common.mutable_ts_instance()->set_instance_seqno(1);
+  common.mutable_ts_instance()->set_instance_seqno(instance_seqno);
   auto result = SendHeartbeat(common, reg);
   if (result.ok()) {
     registered_ts_count_++;

src/yb/master/master_ddl.proto

@@ -864,6 +864,14 @@ message RefreshYsqlLeaseResponsePB {
   }
 }
 
+message RelinquishYsqlLeaseRequestPB {
+  optional NodeInstancePB instance = 1;
+}
+
+message RelinquishYsqlLeaseResponsePB {
+  optional MasterErrorPB error = 1;
+}
+
 service MasterDdl {
   option (yb.rpc.custom_service_name) = "yb.master.MasterService";
 
@@ -928,8 +936,9 @@ service MasterDdl {
     option (yb.rpc.send_metadata) = true;
   }
   rpc RefreshYsqlLease(RefreshYsqlLeaseRequestPB)
-    returns (RefreshYsqlLeaseResponsePB);
-
+      returns (RefreshYsqlLeaseResponsePB);
+  rpc RelinquishYsqlLease(RelinquishYsqlLeaseRequestPB)
+      returns (RelinquishYsqlLeaseResponsePB);
   rpc GetObjectLockStatus(GetObjectLockStatusRequestPB)
       returns (GetObjectLockStatusResponsePB);
 }

src/yb/master/master_ddl_client.cc

@@ -93,6 +93,17 @@ Result<RefreshYsqlLeaseInfoPB> MasterDDLClient::RefreshYsqlLease(
   return resp.info();
 }
 
+Status MasterDDLClient::RelinquishYsqlLease(
+    const std::string& permanent_uuid, int64_t instance_seqno) {
+  RelinquishYsqlLeaseRequestPB req;
+  req.mutable_instance()->set_permanent_uuid(permanent_uuid);
+  req.mutable_instance()->set_instance_seqno(instance_seqno);
+  RelinquishYsqlLeaseResponsePB resp;
+  rpc::RpcController rpc;
+  RETURN_NOT_OK(proxy_.RelinquishYsqlLease(req, &resp, &rpc));
+  return ResponseStatus(resp);
+}
+
 Status MasterDDLClient::DeleteTable(const TableId& id, MonoDelta timeout) {
   DeleteTableRequestPB req;
   req.mutable_table()->set_table_id(id);

src/yb/master/master_ddl_client.h

@@ -40,6 +40,8 @@ class MasterDDLClient {
       const std::string& permanent_uuid, int64_t instance_seqno, uint64_t time_ms,
       std::optional<uint64_t> current_lease_epoch);
 
+  Status RelinquishYsqlLease(const std::string& permanent_uuid, int64_t instance_seqno);
+
   Status DeleteTable(const TableId& id, MonoDelta timeout);
 
  private: