fix(pubsub): avoid deadlock by ensuring lease refresh timer runs on cq thread

Author: scotthart

google/cloud/pubsub/internal/subscription_lease_management.cc

@@ -144,9 +144,23 @@ void SubscriptionLeaseManagement::StartRefreshTimer(
   shutdown_manager_->StartOperation(__func__, "OnRefreshTimer", [&] {
     using TimerFuture = future<StatusOr<std::chrono::system_clock::time_point>>;
     if (refresh_timer_.valid()) refresh_timer_.cancel();
-    refresh_timer_ =
-        cq_.MakeDeadlineTimer(deadline).then([weak](TimerFuture tp) {
-          if (auto self = weak.lock()) self->OnRefreshTimer(!tp.get());
+    // The caller holds mu_ (the unnamed unique_lock parameter) for the full
+    // duration of this function.  If the timer future is already satisfied
+    // when .then() attaches its continuation -- a deadline in the past
+    // (RefreshMessageLeases can compute extensions as small as 1s, and
+    // kAckDeadlineSlack then puts the deadline before now), or a CQ thread
+    // winning the satisfaction race -- the continuation runs INLINE on this
+    // thread, re-enters OnRefreshTimer -> RefreshMessageLeases, and
+    // self-deadlocks re-acquiring the non-recursive mu_.  Every other path
+    // (AckMessage, OnRead, Shutdown) then blocks behind mu_ and the
+    // subscription freezes until the process restarts.  Dispatch the
+    // continuation through the CompletionQueue so OnRefreshTimer always runs
+    // on a CQ thread with no locks held.
+    refresh_timer_ = cq_.MakeDeadlineTimer(deadline).then(
+        [weak = std::move(weak), cq = cq_](TimerFuture tp) mutable {
+          cq.RunAsync([weak = std::move(weak), cancelled = !tp.get()] {
+            if (auto self = weak.lock()) self->OnRefreshTimer(cancelled);
+          });
         });
   });
 }

google/cloud/pubsub/internal/subscription_lease_management_test.cc

@@ -109,18 +109,21 @@ TEST(SubscriptionLeaseManagementTest, NormalLifecycle) {
   // will verify that only the remaining messages have their lease extended.
   uut->AckMessage("ack-0-1");
   fake_cq->SimulateCompletion(true);
+  // RunAsync, drain the deferred OnRefreshTimer
+  fake_cq->SimulateCompletion(true);
   ASSERT_EQ(1U, fake_cq->size());
 
   // Ack one more message and trigger the new timer.
   uut->NackMessage("ack-0-2");
   fake_cq->SimulateCompletion(true);
+  // RunAsync, drain the deferred OnRefreshTimer
+  fake_cq->SimulateCompletion(true);
   ASSERT_EQ(1U, fake_cq->size());
 
   shutdown_manager->MarkAsShutdown(__func__, Status{});
   uut->Shutdown();
 
-  fake_cq->SimulateCompletion(false);
-  ASSERT_EQ(0U, fake_cq->size());
+  while (!fake_cq->empty()) fake_cq->SimulateCompletion(true);
   EXPECT_THAT(done.get(), IsOk());
 }
 
@@ -154,8 +157,7 @@ TEST(SubscriptionLeaseManagementTest, ShutdownOnError) {
           Status(StatusCode::kPermissionDenied, "uh-oh"))});
   ASSERT_EQ(1U, fake_cq->size());
 
-  fake_cq->SimulateCompletion(false);
-  ASSERT_EQ(0U, fake_cq->size());
+  while (!fake_cq->empty()) fake_cq->SimulateCompletion(true);
   EXPECT_THAT(done.get(), StatusIs(StatusCode::kPermissionDenied));
 }
 
@@ -219,12 +221,13 @@ TEST(SubscriptionLeaseManagementTest, UsesDeadlineExtension) {
 
   // Ignore message and then fire the timer. This will extend the deadline.
   fake_cq->SimulateCompletion(true);
+  // RunAsync, drain the deferred OnRefreshTimer
+  fake_cq->SimulateCompletion(true);
   ASSERT_EQ(1U, fake_cq->size());
 
   shutdown_manager->MarkAsShutdown(__func__, Status{});
   uut->Shutdown();
-  fake_cq->SimulateCompletion(false);
-  ASSERT_EQ(0U, fake_cq->size());
+  while (!fake_cq->empty()) fake_cq->SimulateCompletion(true);
   EXPECT_THAT(done.get(), IsOk());
 }
 
@@ -270,13 +273,78 @@ TEST(SubscriptionLeaseManagementTest, ExpiredMessage) {
   // will verify that only the remaining messages have their lease extended.
   uut->AckMessage("ack-0-1");
   fake_cq->SimulateCompletion(true);
+  // RunAsync, drain the deferred OnRefreshTimer
+  fake_cq->SimulateCompletion(true);
   ASSERT_EQ(1U, fake_cq->size());
 
   shutdown_manager->MarkAsShutdown(__func__, Status{});
   uut->Shutdown();
+  while (!fake_cq->empty()) fake_cq->SimulateCompletion(true);
+  EXPECT_THAT(done.get(), IsOk());
+}
+
+/// @test Regression test for the self-deadlock in StartRefreshTimer: the lease
+/// refresh-timer continuation must be dispatched via `RunAsync` rather than run
+/// inline in the timer callback. Running it inline re-entered the already-held
+/// `mu_` and self-deadlocked the subscription. Here we verify that firing the
+/// timer only *schedules* the refresh (a RunAsync task); the lease extension
+/// runs on a subsequent completion, never synchronously in the timer callback.
+TEST(SubscriptionLeaseManagementTest,
+     RefreshTimerContinuationDispatchedViaRunAsync) {
+  auto mock = std::make_shared<pubsub_testing::MockSubscriptionBatchSource>();
+  std::shared_ptr<BatchCallback> batch_callback;
+  EXPECT_CALL(*mock, Start).WillOnce([&](std::shared_ptr<BatchCallback> cb) {
+    batch_callback = std::move(cb);
+  });
 
-  fake_cq->SimulateCompletion(false);
-  ASSERT_EQ(0U, fake_cq->size());
+  auto mock_batch_callback =
+      std::make_shared<pubsub_testing::MockBatchCallback>();
+  EXPECT_CALL(*mock_batch_callback, callback).Times(1);
+
+  int extend_calls = 0;
+  EXPECT_CALL(*mock, ExtendLeases)
+      .WillRepeatedly(
+          [&](std::vector<std::string> const&, std::chrono::seconds) {
+            ++extend_calls;
+            return make_ready_future(Status{});
+          });
+  EXPECT_CALL(*mock, BulkNack)
+      .WillRepeatedly([](std::vector<std::string> const&) {
+        return make_ready_future(Status{});
+      });
+  EXPECT_CALL(*mock, Shutdown).Times(1);
+
+  auto fake_cq = std::make_shared<FakeCompletionQueueImpl>();
+  CompletionQueue cq(fake_cq);
+  auto shutdown_manager = std::make_shared<SessionShutdownManager>();
+  auto uut = SubscriptionLeaseManagement::Create(cq, shutdown_manager, mock,
+                                                 std::chrono::seconds(345),
+                                                 std::chrono::seconds(600));
+
+  auto done = shutdown_manager->Start({});
+  uut->Start(mock_batch_callback);
+  batch_callback->callback(
+      BatchCallback::StreamingPullResponse{GenerateMessages("0-", 1)});
+  ASSERT_EQ(1U, fake_cq->size());  // the refresh timer is pending
+
+  // Fire the timer. With the fix, the timer callback only *schedules*
+  // OnRefreshTimer via RunAsync; it must NOT extend leases inline (doing so
+  // re-enters the held mutex and deadlocks).
+  fake_cq->SimulateCompletion(true);
+  EXPECT_EQ(0, extend_calls)
+      << "lease refresh ran inline in the timer callback (deadlock path)";
+  EXPECT_EQ(1U, fake_cq->size());  // a RunAsync task is now pending
+
+  // Draining the RunAsync task runs OnRefreshTimer -> ExtendLeases.
+  fake_cq->SimulateCompletion(true);
+  EXPECT_EQ(1, extend_calls);
+
+  // The fix defers OnRefreshTimer to RunAsync, which the fake CQ only executes
+  // on SimulateCompletion(true) (false drops the task). Drain with true so the
+  // deferred refresh operation finishes and the session shutdown completes.
+  shutdown_manager->MarkAsShutdown(__func__, Status{});
+  uut->Shutdown();
+  while (!fake_cq->empty()) fake_cq->SimulateCompletion(true);
   EXPECT_THAT(done.get(), IsOk());
 }