[#26672] docdb: Fixed deadlock on cleaning expired PG sessions when pg_client_use_shared_memory is on

Summary:
Faced on a 2024.2.2.1-b7 cluster with running YSQL workload and `--pg_client_use_shared_memory=true`.
All 4 threads of `iotp_TabletServer` threadpool are waiting inside `ScheduledTaskTracker::CompleteShutdown()`.
All of them are waiting for `big_shared_mem_expiration_task_.CompleteShutdown()`.

`big_shared_mem_expiration_task_` was scheduled but not launched due to `iotp_TabletServer` pool is busy with processing `PgClientServiceImpl::Impl::PgCheckExpiredSessions` and all `PgCheckExpiredSessions` are waiting inside `ScheduledTaskTracker::CompleteShutdown()` for `big_shared_mem_expiration_task_ to` complete.

Added unit-test to reproduce this issue and implemented fix.

Also updated `ScheduledTaskTracker` logs for easier debugging in future.
Logs before update:
```
I0401 17:00:43.968747 3586401 scheduler.cc:230] Waiting 1 tasks to complete
```
After update:
```
I0404 19:42:28.436432 1233831 scheduler.cc:231] big_shared_mem_expiration_task(0x000073c0bed8dcd8): Waiting 1 tasks to complete
```
Jira: DB-16052

Test Plan: PgConnTest.SessionExpiration

Reviewers: sergei, amitanand

Reviewed By: sergei, amitanand

Subscribers: yql, ybase

Tags: #jenkins-ready

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

Author: ttyusupov

src/yb/consensus/raft_consensus.cc

@@ -426,7 +426,8 @@ RaftConsensus::RaftConsensus(
       queue_(std::move(queue)),
       rng_(GetRandomSeed32()),
       withhold_votes_until_(MonoTime::Min()),
-      step_down_check_tracker_(&peer_proxy_factory_->messenger()->scheduler()),
+      step_down_check_tracker_(
+          "step_down_check_tracker", &peer_proxy_factory_->messenger()->scheduler()),
       mark_dirty_clbk_(std::move(mark_dirty_clbk)),
       shutdown_(false),
       follower_memory_pressure_rejections_(tablet_metric_entity->FindOrCreateMetric<Counter>(

src/yb/rpc/scheduler.cc

@@ -194,14 +194,15 @@ IoService& Scheduler::io_service() {
   return impl_->io_service();
 }
 
-ScheduledTaskTracker::ScheduledTaskTracker(Scheduler* scheduler)
-    : scheduler_(DCHECK_NOTNULL(scheduler)) {}
+ScheduledTaskTracker::ScheduledTaskTracker(const std::string& name, Scheduler* scheduler)
+    : log_prefix_(Format("$0($1): ", name, static_cast<void*>(this))),
+      scheduler_(DCHECK_NOTNULL(scheduler)) {}
 
 ScheduledTaskTracker::~ScheduledTaskTracker() {
   auto last_scheduled_task_id = last_scheduled_task_id_.load(std::memory_order_acquire);
   if (last_scheduled_task_id != rpc::kInvalidTaskId) {
     auto num_scheduled = num_scheduled_.load(std::memory_order_acquire);
-    LOG_IF(DFATAL, num_scheduled != kShutdownMark)
+    LOG_IF_WITH_PREFIX(DFATAL, num_scheduled != kShutdownMark)
         << "Shutdown did not complete on ScheduledTaskTracker";
   }
 }
@@ -220,14 +221,18 @@ void ScheduledTaskTracker::StartShutdown() {
   }
 }
 
+bool ScheduledTaskTracker::ReadyToShutdown() const {
+  return num_scheduled_.load(std::memory_order_acquire) == kShutdownMark;
+}
+
 void ScheduledTaskTracker::CompleteShutdown() {
   for (;;) {
     auto left = num_scheduled_.load(std::memory_order_acquire) - kShutdownMark;
     if (left <= 0) {
-      LOG_IF(DFATAL, left < 0) << "Negative number of tasks left: " << left;
+      LOG_IF_WITH_PREFIX(DFATAL, left < 0) << "Negative number of tasks left: " << left;
       break;
     }
-    YB_LOG_EVERY_N_SECS(INFO, 1) << "Waiting " << left << " tasks to complete";
+    YB_LOG_WITH_PREFIX_EVERY_N_SECS(INFO, 1) << "Waiting " << left << " tasks to complete";
     Abort();
     std::this_thread::sleep_for(1ms);
   }

src/yb/rpc/scheduler.h

@@ -118,7 +118,7 @@ class ScheduledTaskTracker {
   ScheduledTaskTracker() = default;
   ~ScheduledTaskTracker();
 
-  explicit ScheduledTaskTracker(Scheduler* scheduler);
+  ScheduledTaskTracker(const std::string& name, Scheduler* scheduler);
 
   void Bind(Scheduler* scheduler) {
     scheduler_ = scheduler;
@@ -147,14 +147,20 @@ class ScheduledTaskTracker {
   void Abort();
 
   void StartShutdown();
+  bool ReadyToShutdown() const;
   void CompleteShutdown();
 
   void Shutdown() {
     StartShutdown();
     CompleteShutdown();
   }
 
+  const std::string& LogPrefix() {
+    return log_prefix_;
+  }
+
  private:
+  std::string log_prefix_;
   Scheduler* scheduler_ = nullptr;
   std::atomic<int64_t> num_scheduled_{0};
   std::atomic<rpc::ScheduledTaskId> last_scheduled_task_id_{rpc::kInvalidTaskId};

src/yb/tserver/pg_client_service.cc

@@ -133,6 +133,9 @@ DEFINE_test_flag(uint64, ysql_oid_prefetch_adjustment, 0,
                  "production environment. In unit test we use this flag to force allocation of "
                  "large Postgres OIDs.");
 
+DEFINE_test_flag(uint64, delay_before_complete_expired_pg_sessions_shutdown_ms, 0,
+                 "Inject delay before completing shutdown of expired PG sessions.");
+
 DEFINE_RUNTIME_uint64(ysql_cdc_active_replication_slot_window_ms, 60000,
                       "Determines the window in milliseconds in which if a client has consumed the "
                       "changes of a ReplicationSlot across any tablet, then it is considered to be "
@@ -244,7 +247,8 @@ class LockablePgClientSession {
   }
 
   bool ReadyToShutdown() const {
-    return !exchange_runnable_ || exchange_runnable_->ReadyToShutdown();
+    return (!exchange_runnable_ || exchange_runnable_->ReadyToShutdown()) &&
+           session_.ReadyToShutdown();
   }
 
   void CompleteShutdown() {
@@ -452,8 +456,8 @@ class PgClientServiceImpl::Impl {
         transaction_pool_provider_(std::move(transaction_pool_provider)),
         messenger_(*messenger),
         table_cache_(client_future_),
-        check_expired_sessions_(&messenger->scheduler()),
-        check_object_id_allocators_(&messenger->scheduler()),
+        check_expired_sessions_("check_expired_sessions", &messenger->scheduler()),
+        check_object_id_allocators_("check_object_id_allocators", &messenger->scheduler()),
         response_cache_(parent_mem_tracker, metric_entity),
         instance_id_(permanent_uuid),
         shared_mem_pool_(parent_mem_tracker, instance_id_),
@@ -1833,6 +1837,7 @@ class PgClientServiceImpl::Impl {
       session->session().StartShutdown();
       txn_snapshot_manager_.UnregisterAll(session->id());
     }
+    AtomicFlagSleepMs(&FLAGS_TEST_delay_before_complete_expired_pg_sessions_shutdown_ms);
     for (const auto& session : expired_sessions) {
       if (session->session().ReadyToShutdown()) {
         session->session().CompleteShutdown();

src/yb/tserver/pg_client_session.cc

@@ -1126,7 +1126,7 @@ class PgClientSession::Impl {
         lease_epoch_(lease_epoch),
         ts_lock_manager_(std::move(lock_manager)),
         transaction_provider_(std::move(transaction_builder)),
-        big_shared_mem_expiration_task_(&scheduler),
+        big_shared_mem_expiration_task_("big_shared_mem_expiration_task", &scheduler),
         read_point_history_(PrefixLogger(id_)) {}
 
   [[nodiscard]] auto id() const {return id_; }
@@ -2131,6 +2131,10 @@ class PgClientSession::Impl {
     big_shared_mem_expiration_task_.StartShutdown();
   }
 
+  bool ReadyToShutdown() {
+    return big_shared_mem_expiration_task_.ReadyToShutdown();
+  }
+
   void CompleteShutdown() {
     big_shared_mem_expiration_task_.CompleteShutdown();
   }
@@ -3123,6 +3127,10 @@ void PgClientSession::StartShutdown() {
   return impl_->StartShutdown();
 }
 
+bool PgClientSession::ReadyToShutdown() const {
+  return impl_->ReadyToShutdown();
+}
+
 void PgClientSession::CompleteShutdown() {
   impl_->CompleteShutdown();
 }

src/yb/tserver/pg_client_session.h

@@ -122,6 +122,7 @@ class PgClientSession final {
   std::pair<uint64_t, std::byte*> ObtainBigSharedMemorySegment(size_t size);
 
   void StartShutdown();
+  bool ReadyToShutdown() const;
   void CompleteShutdown();
 
   Result<ReadHybridTime> GetTxnSnapshotReadTime(

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

@@ -10,7 +10,9 @@
 // or implied.  See the License for the specific language governing permissions and limitations
 // under the License.
 
+#include "yb/util/backoff_waiter.h"
 #include "yb/util/logging.h"
+#include "yb/util/test_thread_holder.h"
 
 #include "yb/yql/pgwrapper/libpq_test_base.h"
 
@@ -20,6 +22,11 @@ namespace yb {
 namespace pgwrapper {
 
 class PgConnTest : public LibPqTestBase {
+ public:
+  Result<PGConn> Connect(bool simple_query_protocol = false) {
+    return LibPqTestBase::Connect(simple_query_protocol);
+  }
+
  protected:
   void TestUriAuth();
 
@@ -183,5 +190,117 @@ TEST_F_EX(PgConnTest, ConnectionLimit, PgConnTestLimit) {
   ASSERT_STR_CONTAINS(s.message().ToBuffer(), "sorry, too many clients already");
 }
 
+class PgSessionExpirationTest : public PgConnTest {
+ public:
+  int GetNumTabletServers() const override {
+    return 1;
+  }
+
+  void UpdateMiniClusterOptions(ExternalMiniClusterOptions* options) override {
+    PgConnTest::UpdateMiniClusterOptions(options);
+    for (const auto& tserver_flag : std::initializer_list<std::string>{
+             "--pg_client_use_shared_memory=true",
+             "--max_big_shared_memory_segment_size=1048576",
+             Format("--io_thread_pool_size=$0", kIoThreadPoolSize),
+             // Should be big enough for other PgClientServiceImpl::Impl::CheckExpiredSessions
+             // to be called by scheduler.
+             Format("--TEST_delay_before_complete_expired_pg_sessions_shutdown_ms=1000"),
+             Format(
+                 "--pg_client_heartbeat_interval_ms=$0",
+                 ToMilliseconds(kPgClientHeartbeatInterval)),
+             Format(
+                 "--pg_client_session_expiration_ms=$0",
+                 ToMilliseconds(kPgClientSessionExpiration)),
+             Format(
+                 "--big_shared_memory_segment_session_expiration_time_ms=$0",
+                 ToMilliseconds(kBigSharedMemorySegmentSesionExpiration)),
+         }) {
+      options->extra_tserver_flags.push_back(tserver_flag);
+    }
+  }
+
+  static constexpr auto kPgClientHeartbeatInterval =
+      RegularBuildVsDebugVsSanitizers(1500ms, 2000ms, 2000ms);
+  static constexpr auto kPgClientSessionExpiration = kPgClientHeartbeatInterval * 2;
+  static constexpr auto kBigSharedMemorySegmentSesionExpiration = kPgClientSessionExpiration * 2;
+  static constexpr auto kIoThreadPoolSize = 2;
+};
+
+TEST_F_EX(PgConnTest, SessionExpiration, PgSessionExpirationTest) {
+  constexpr auto kGroupSize = 5;
+  constexpr auto kNumGroups = kIoThreadPoolSize;
+  constexpr auto kDelayBetweenGroupConnectionsClose = 100ms;
+
+  class ClientJob {
+   public:
+    ClientJob(const std::string& name, PgConnTest& test, CountDownLatch& latch)
+        : log_prefix_(name + ": "), test_(test), latch_(latch) {}
+
+    void operator()() const {
+      LOG_WITH_PREFIX(INFO) << "Connecting...";
+      PGConn conn = ASSERT_RESULT(test_.Connect());
+      LOG_WITH_PREFIX(INFO) << "Connected";
+      LOG_WITH_PREFIX(INFO) << "Got response: " << ASSERT_RESULT(conn.FetchRow<int32_t>("SELECT 1"))
+                            << " latch count: " << latch_.count();
+      latch_.CountDown();
+
+      // Close connections for the group at once.
+      LOG_WITH_PREFIX(INFO) << "Waiting for latch";
+      latch_.Wait();
+      LOG_WITH_PREFIX(INFO) << "Closing connection...";
+    }
+
+    const std::string& LogPrefix() const { return log_prefix_; }
+
+   private:
+    const std::string log_prefix_;
+    PgConnTest& test_;
+    CountDownLatch& latch_;
+  };
+
+  TestThreadHolder thread_holder;
+  std::vector<std::unique_ptr<CountDownLatch>> latches;
+  for (int i = 0; i < kNumGroups; ++i) {
+    latches.push_back(std::make_unique<CountDownLatch>(kGroupSize + 1));
+    for (int j = 0; j < kGroupSize; ++j) {
+      thread_holder.AddThreadFunctor(
+          ClientJob(Format("Group #$0, thread #$1", i, j), *this, *latches.back()));
+    }
+  }
+
+  ASSERT_OK(LoggedWaitFor(
+      [&latches] {
+        for (auto& latch : latches) {
+          if (latch->count() > 1) {
+            return false;
+          }
+        }
+        return true;
+      },
+      RegularBuildVsDebugVsSanitizers(60s, 120s, 180s),
+      Format("Wait for all queries to return.")));
+
+  // Unblock threads with delay between groups.
+  for (auto& latch : latches) {
+    LOG(INFO) << "latch->count(): " << latch->count();
+    latch->CountDown();
+    std::this_thread::sleep_for(kDelayBetweenGroupConnectionsClose);
+  }
+
+  LOG(INFO) << "Joining test threads";
+  thread_holder.JoinAll();
+
+  // Wait for sessions to expire.
+  std::this_thread::sleep_for(kPgClientSessionExpiration + 500ms);
+
+  LOG(INFO) << "Shutting down cluster";
+  cluster_->Shutdown();
+  for (const auto& server : cluster_->daemons()) {
+    if (server) {
+      ASSERT_FALSE(server->WasUnsafeShutdown());
+    }
+  }
+}
+
 } // namespace pgwrapper
 } // namespace yb