chore: force shutdown connections on takeover

kostasrim · kostasrim · commit 5f3de98d5e53 · 2025-12-01T15:07:05.000+02:00
Signed-off-by: Kostas Kyrimis &lt;kostas@dragonflydb.io&gt;
diff --git a/src/facade/dragonfly_connection.cc b/src/facade/dragonfly_connection.cc
@@ -272,7 +272,7 @@ void LogTraffic(uint32_t id, bool has_more, absl::Span<RespExpr> resp,
 
   // Write the data itself.
   array<iovec, 16> blobs;
-  unsigned index = 0;
+  uint64_t index = 0;
   if (next != stack_buf) {
     blobs[index++] = iovec{.iov_base = stack_buf, .iov_len = size_t(next - stack_buf)};
   }
@@ -538,7 +538,7 @@ void Connection::AsyncOperations::operator()(const PubMessage& pub_msg) {
     return;
   }
 
-  unsigned i = 0;
+  uint64_t i = 0;
   array<string_view, 4> arr;
   if (pub_msg.pattern.empty()) {
     arr[i++] = pub_msg.is_sharded ? "smessage" : "message";
@@ -609,11 +609,11 @@ void UpdateLibNameVerMap(const string& name, const string& ver, int delta) {
 }
 }  // namespace
 
-void Connection::Init(unsigned io_threads) {
+void Connection::Init(uint64_t io_threads) {
   CHECK(thread_queue_backpressure == nullptr);
   thread_queue_backpressure = new QueueBackpressure[io_threads];
 
-  for (unsigned i = 0; i < io_threads; ++i) {
+  for (uint64_t i = 0; i < io_threads; ++i) {
     auto& qbp = thread_queue_backpressure[i];
     qbp.publish_buffer_limit = GetFlag(FLAGS_publish_buffer_limit);
     qbp.pipeline_cache_limit = GetFlag(FLAGS_request_cache_limit);
@@ -861,7 +861,7 @@ void Connection::HandleRequests() {
   }
 }
 
-unsigned Connection::GetSendWaitTimeSec() const {
+uint64_t Connection::GetSendWaitTimeSec() const {
   if (reply_builder_ && reply_builder_->IsSendActive()) {
     return (util::fb2::ProactorBase::GetMonotonicTimeNs() - reply_builder_->GetLastSendTimeNs()) /
            1'000'000'000;
@@ -877,11 +877,11 @@ void Connection::RegisterBreakHook(BreakerCb breaker_cb) {
 pair<string, string> Connection::GetClientInfoBeforeAfterTid() const {
   if (!socket_) {
     LOG(DFATAL) << "unexpected null socket_ "
-                << " phase " << unsigned(phase_) << ", is_http: " << unsigned(is_http_);
+                << " phase " << uint64_t(phase_) << ", is_http: " << unsigned(is_http_);
     return {};
   }
 
-  CHECK_LT(unsigned(phase_), NUM_PHASES);
+  CHECK_LT(uint64_t(phase_), NUM_PHASES);
 
   string before;
   auto le = LocalBindStr();
@@ -947,7 +947,7 @@ pair<string, string> Connection::GetClientInfoBeforeAfterTid() const {
   return {std::move(before), std::move(after)};
 }
 
-string Connection::GetClientInfo(unsigned thread_id) const {
+string Connection::GetClientInfo(uint64_t thread_id) const {
   auto [before, after] = GetClientInfoBeforeAfterTid();
   absl::StrAppend(&before, " tid=", thread_id);
   absl::StrAppend(&before, after);
@@ -1213,7 +1213,7 @@ void Connection::DispatchSingle(bool has_more, absl::FunctionRef<void()> invoke_
   }
 }
 
-Connection::ParserStatus Connection::ParseRedis(unsigned max_busy_cycles) {
+Connection::ParserStatus Connection::ParseRedis(uint64_t max_busy_cycles) {
   uint32_t consumed = 0;
   RedisParser::Result result = RedisParser::OK;
 
@@ -1827,7 +1827,7 @@ bool Connection::Migrate(util::fb2::ProactorBase* dest) {
 Connection::WeakRef Connection::Borrow() {
   DCHECK(self_);
 
-  return {self_, unsigned(socket_->proactor()->GetPoolIndex()), id_};
+  return {self_, uint64_t(socket_->proactor()->GetPoolIndex()), id_};
 }
 
 void Connection::ShutdownThreadLocal() {
@@ -2095,7 +2095,7 @@ bool Connection::IsReplySizeOverLimit() const {
 }
 
 void Connection::UpdateFromFlags() {
-  unsigned tid = fb2::ProactorBase::me()->GetPoolIndex();
+  uint64_t tid = fb2::ProactorBase::me()->GetPoolIndex();
   thread_queue_backpressure[tid].pipeline_queue_max_len = GetFlag(FLAGS_pipeline_queue_limit);
   thread_queue_backpressure[tid].pipeline_buffer_limit = GetFlag(FLAGS_pipeline_buffer_limit);
   thread_queue_backpressure[tid].pipeline_cnd.notify_all();
@@ -2126,11 +2126,11 @@ void Connection::TrackRequestSize(bool enable) {
   }
 }
 
-void Connection::EnsureMemoryBudget(unsigned tid) {
+void Connection::EnsureMemoryBudget(uint64_t tid) {
   thread_queue_backpressure[tid].EnsureBelowLimit();
 }
 
-Connection::WeakRef::WeakRef(const std::shared_ptr<Connection>& ptr, unsigned thread_id,
+Connection::WeakRef::WeakRef(const std::shared_ptr<Connection>& ptr, uint64_t thread_id,
                              uint32_t client_id)
     : ptr_{ptr}, last_known_thread_id_{thread_id}, client_id_{client_id} {
 }
diff --git a/src/facade/dragonfly_connection.h b/src/facade/dragonfly_connection.h
@@ -53,7 +53,7 @@ class SinkReplyBuilder;
 // a separate dispatch queue that is processed on a separate fiber.
 class Connection : public util::Connection {
  public:
-  static void Init(unsigned io_threads);
+  static void Init(uint64_t io_threads);
   static void Shutdown();
   static void ShutdownThreadLocal();
 
@@ -186,7 +186,7 @@ class Connection : public util::Connection {
   struct WeakRef {
    public:
     // Get residing thread of connection. Thread-safe.
-    unsigned LastKnownThreadId() const {
+    uint64_t LastKnownThreadId() const {
       return last_known_thread_id_;
     }
     // Get pointer to connection if still valid, nullptr if expired.
@@ -206,10 +206,10 @@ class Connection : public util::Connection {
    private:
     friend class Connection;
 
-    WeakRef(const std::shared_ptr<Connection>& ptr, unsigned thread_id, uint32_t client_id);
+    WeakRef(const std::shared_ptr<Connection>& ptr, uint64_t thread_id, uint32_t client_id);
 
     std::weak_ptr<Connection> ptr_;
-    unsigned last_known_thread_id_;
+    uint64_t last_known_thread_id_;
     uint32_t client_id_;
   };
 
@@ -247,7 +247,7 @@ class Connection : public util::Connection {
 
   bool IsCurrentlyDispatching() const;
 
-  std::string GetClientInfo(unsigned thread_id) const;
+  std::string GetClientInfo(uint64_t thread_id) const;
   std::string GetClientInfo() const;
 
   virtual std::string RemoteEndpointStr() const;  // virtual because overwritten in test_utils
@@ -313,14 +313,14 @@ class Connection : public util::Connection {
   static std::vector<std::string> GetMutableFlagNames();  // Triggers UpdateFromFlags
 
   static void TrackRequestSize(bool enable);
-  static void EnsureMemoryBudget(unsigned tid);
+  static void EnsureMemoryBudget(uint64_t tid);
   static void GetRequestSizeHistogramThreadLocal(std::string* hist);
 
-  unsigned idle_time() const {
+  uint64_t idle_time() const {
     return time(nullptr) - last_interaction_;
   }
 
-  unsigned GetSendWaitTimeSec() const;
+  uint64_t GetSendWaitTimeSec() const;
 
   Phase phase() const {
     return phase_;
@@ -370,7 +370,7 @@ class Connection : public util::Connection {
   // Create new pipeline request, re-use from pool when possible.
   PipelineMessagePtr FromArgs(const RespVec& args);
 
-  ParserStatus ParseRedis(unsigned max_busy_cycles);
+  ParserStatus ParseRedis(uint64_t max_busy_cycles);
   ParserStatus ParseMemcache();
 
   void OnBreakCb(int32_t mask);
@@ -456,7 +456,7 @@ class Connection : public util::Connection {
   std::string lib_name_;
   std::string lib_ver_;
 
-  unsigned parser_error_ = 0;
+  uint64_t parser_error_ = 0;
 
   BreakerCb breaker_cb_;
 
diff --git a/src/server/dflycmd.cc b/src/server/dflycmd.cc
@@ -40,6 +40,8 @@ ABSL_RETIRED_FLAG(uint32_t, allow_partial_sync_with_lsn_diff, 0,
 ABSL_DECLARE_FLAG(bool, info_replication_valkey_compatible);
 ABSL_DECLARE_FLAG(uint32_t, replication_timeout);
 ABSL_DECLARE_FLAG(uint32_t, shard_repl_backlog_len);
+ABSL_FLAG(bool, experimental_force_takeover, false,
+          "Attempts to force takeover in case of stuck connections");
 
 namespace dfly {
 
@@ -466,6 +468,45 @@ std::optional<LSN> DflyCmd::ParseLsnVec(std::string_view last_master_lsn,
   return {lsn_vec[flow_id]};
 }
 
+void DflyCmd::ForceShutdownStuckConnections(uint64_t timeout) {
+  vector<facade::Connection::WeakRef> conn_refs;
+  auto cb = [&](unsigned thread_index, util::Connection* conn) {
+    facade::Connection* dcon = static_cast<facade::Connection*>(conn);
+    LOG(INFO) << dcon->DebugInfo();
+    // Kill Connection here
+    facade::Connection* dfly_conn = static_cast<facade::Connection*>(conn);
+    using Phase = facade::Connection::Phase;
+    auto phase = dfly_conn->phase();
+    if (dfly_conn->cntx() && dfly_conn->cntx()->replica_conn) {
+      return;
+    }
+
+    bool idle_read = phase == Phase::READ_SOCKET && dfly_conn->idle_time() > timeout;
+
+    bool stuck_sending = dfly_conn->IsSending() && dfly_conn->GetSendWaitTimeSec() > timeout;
+
+    if (idle_read || stuck_sending) {
+      LOG(INFO) << "Connection check: " << dfly_conn->GetClientInfo()
+                << ", phase=" << static_cast<int>(phase) << ", idle_time=" << dfly_conn->idle_time()
+                << ", is_sending=" << dfly_conn->IsSending() << ", idle_read=" << idle_read
+                << ", stuck_sending=" << stuck_sending;
+    }
+    conn_refs.push_back(dfly_conn->Borrow());
+  };
+
+  for (auto* listener : sf_->GetListeners()) {
+    listener->TraverseConnections(cb);
+  }
+
+  VLOG(1) << "Found " << conn_refs.size() << " stucked connections ";
+  for (auto& ref : conn_refs) {
+    facade::Connection* conn = ref.Get();
+    if (conn) {
+      conn->ShutdownSelfBlocking();
+    }
+  }
+}
+
 // DFLY TAKEOVER <timeout_sec> [SAVE] <sync_id>
 // timeout_sec - number of seconds to wait for TAKEOVER to converge.
 // SAVE option is used only by tests.
@@ -506,7 +547,7 @@ void DflyCmd::TakeOver(CmdArgList args, RedisReplyBuilder* rb, ConnectionContext
   LOG(INFO) << "Takeover initiated, locking down the database.";
   absl::Duration timeout_dur = absl::Seconds(timeout);
   absl::Time end_time = absl::Now() + timeout_dur;
-  AggregateStatus status;
+  OpStatus status = OpStatus::OK;
 
   // We need to await for all dispatches to finish: Otherwise a transaction might be scheduled
   // after this function exits but before the actual shutdown.
@@ -520,13 +561,22 @@ void DflyCmd::TakeOver(CmdArgList args, RedisReplyBuilder* rb, ConnectionContext
     LOG(WARNING) << "Couldn't wait for commands to finish dispatching. " << timeout_dur;
     status = OpStatus::TIMED_OUT;
 
-    auto cb = [&](unsigned thread_index, util::Connection* conn) {
-      facade::Connection* dcon = static_cast<facade::Connection*>(conn);
-      LOG(INFO) << dcon->DebugInfo();
-    };
-
-    for (auto* listener : sf_->GetListeners()) {
-      listener->TraverseConnections(cb);
+    // Force takeover on the same duration if flag is set
+    if (absl::GetFlag(FLAGS_experimental_force_takeover)) {
+      ForceShutdownStuckConnections(uint64_t(timeout));
+
+      // Safety net.
+      facade::DispatchTracker tracker{sf_->GetNonPriviligedListeners(), cntx->conn(), false, false};
+      shard_set->pool()->AwaitFiberOnAll([&](unsigned index, auto* pb) {
+        sf_->CancelBlockingOnThread();
+        tracker.TrackOnThread();
+      });
+
+      status = OpStatus::OK;
+      if (!tracker.Wait(timeout_dur)) {
+        LOG(ERROR) << "Could not force execute takeover";
+        status = OpStatus::TIMED_OUT;
+      }
     }
   }
 
@@ -540,10 +590,11 @@ void DflyCmd::TakeOver(CmdArgList args, RedisReplyBuilder* rb, ConnectionContext
   });
 
   atomic_bool catchup_success = true;
-  if (*status == OpStatus::OK) {
+  if (status == OpStatus::OK) {
     dfly::SharedLock lk{replica_ptr->shared_mu};
-    auto cb = [replica_ptr = replica_ptr, end_time, &catchup_success](EngineShard* shard) {
-      if (!WaitReplicaFlowToCatchup(end_time, replica_ptr.get(), shard)) {
+    auto time = end_time + timeout_dur;
+    auto cb = [replica_ptr = replica_ptr, time, &catchup_success](EngineShard* shard) {
+      if (!WaitReplicaFlowToCatchup(time, replica_ptr.get(), shard)) {
         catchup_success.store(false);
       }
     };
@@ -552,8 +603,9 @@ void DflyCmd::TakeOver(CmdArgList args, RedisReplyBuilder* rb, ConnectionContext
 
   VLOG(1) << "WaitReplicaFlowToCatchup done";
 
-  if (*status != OpStatus::OK || !catchup_success.load()) {
+  if (status != OpStatus::OK || !catchup_success.load()) {
     sf_->service().SwitchState(GlobalState::TAKEN_OVER, GlobalState::ACTIVE);
+    LOG(INFO) << status << " " << catchup_success.load() << " " << &status;
     return rb->SendError("Takeover failed!");
   }
 
diff --git a/src/server/dflycmd.h b/src/server/dflycmd.h
@@ -183,6 +183,11 @@ class DflyCmd {
   // Switch to stable state replication.
   void StartStable(CmdArgList args, Transaction* tx, RedisReplyBuilder* rb);
 
+  // Helper for takeover flow. Sometimes connections get stuck on send (because of pipelines)
+  // and this causes the takeover flow to fail because checkpoint messages are not processed.
+  // This function force shuts down those connection and allows the node to complete the takeover.
+  void ForceShutdownStuckConnections(uint64_t timeout);
+
   // TAKEOVER <syncid>
   // Shut this master down atomically with replica promotion.
   void TakeOver(CmdArgList args, RedisReplyBuilder* rb, ConnectionContext* cntx);
diff --git a/tests/dragonfly/replication_test.py b/tests/dragonfly/replication_test.py
@@ -3668,7 +3668,7 @@ async def test_replica_of_self(async_client):
         await async_client.execute_command(f"replicaof 127.0.0.1 {port}")
 
 
-@dfly_args({"proactor_threads": 2})
+@dfly_args({"proactor_threads": 2, "experimental_force_takeover": True})
 async def test_takeover_with_stuck_connections(df_factory: DflyInstanceFactory):
     master = df_factory.create()
     master.start()
@@ -3686,8 +3686,13 @@ async def test_takeover_with_stuck_connections(df_factory: DflyInstanceFactory):
 
     async def get_task():
         while True:
-            writer.write(f"GET a\n".encode())
-            await writer.drain()
+            # Will get killed by takeover because it's stucked
+            try:
+                writer.write(f"GET a\n".encode())
+                await writer.drain()
+            except:
+                return
+
             await asyncio.sleep(0.1)
 
     get = asyncio.create_task(get_task())
@@ -3715,5 +3720,7 @@ async def wait_for_stuck_on_send():
     async with async_timeout.timeout(240):
         await wait_for_replicas_state(replica_cl)
 
-    with pytest.raises(redis.exceptions.ResponseError) as e:
-        await replica_cl.execute_command("REPLTAKEOVER 5")
+    res = await replica_cl.execute_command("REPLTAKEOVER 5")
+    assert res == "OK"
+
+    await get
