feat: allow migrating client connections via CLIENT MIGRATE (#5551)

* feat: allow migrating client connections via CLIENT MIGRATE

---------

Signed-off-by: Roman Gershman <[email protected]>

Author: romange

src/facade/dragonfly_connection.cc

@@ -707,6 +707,9 @@ void Connection::OnPreMigrateThread() {
   // is marked beforehand.
   migration_in_process_ = true;
 
+  // Mark as not owned by any thread as it going through the dark hole
+  self_.reset();
+
   socket_->CancelOnErrorCb();
   DCHECK(!async_fb_.IsJoinable()) << GetClientId();
 }
@@ -719,6 +722,7 @@ void Connection::OnPostMigrateThread() {
     socket_->RegisterOnErrorCb([this](int32_t mask) { this->OnBreakCb(mask); });
   }
   migration_in_process_ = false;
+  self_ = {make_shared<std::monostate>(), this};  // Recreate shared_ptr to self.
   DCHECK(!async_fb_.IsJoinable());
 
   // If someone had sent Async during the migration, we must create async_fb_.
@@ -1820,7 +1824,7 @@ bool Connection::Migrate(util::fb2::ProactorBase* dest) {
 Connection::WeakRef Connection::Borrow() {
   DCHECK(self_);
 
-  return WeakRef(self_, socket_->proactor()->GetPoolIndex(), id_);
+  return {self_, unsigned(socket_->proactor()->GetPoolIndex()), id_};
 }
 
 void Connection::ShutdownThreadLocal() {
@@ -2003,8 +2007,8 @@ facade::ConnectionContext* Connection::cntx() {
   return cc_.get();
 }
 
-void Connection::RequestAsyncMigration(util::fb2::ProactorBase* dest) {
-  if (!migration_enabled_ || cc_ == nullptr) {
+void Connection::RequestAsyncMigration(util::fb2::ProactorBase* dest, bool force) {
+  if ((!force && !migration_enabled_) || cc_ == nullptr) {
     return;
   }
 
@@ -2189,22 +2193,17 @@ void Connection::SetPipelineLowBoundStats(unsigned limit) {
 
 Connection::WeakRef::WeakRef(std::shared_ptr<Connection> ptr, unsigned thread_id,
                              uint32_t client_id)
-    : ptr_{std::move(ptr)}, thread_id_{thread_id}, client_id_{client_id} {
-}
-
-unsigned Connection::WeakRef::Thread() const {
-  return thread_id_;
+    : ptr_{std::move(ptr)}, last_known_thread_id_{thread_id}, client_id_{client_id} {
 }
 
 Connection* Connection::WeakRef::Get() const {
-  // We should never access the connection object from other threads.
-  DCHECK_EQ(ProactorBase::me()->GetPoolIndex(), int(thread_id_));
+  auto sptr = ptr_.lock();
 
   //  The connection can only be deleted on this thread, so
   //  this pointer is valid until the next suspension.
   //  Note: keeping a shared_ptr doesn't prolong the lifetime because
   //  it doesn't manage the underlying connection. See definition of `self_`.
-  return ptr_.lock().get();
+  return sptr.get();
 }
 
 bool Connection::WeakRef::IsExpired() const {

src/facade/dragonfly_connection.h

@@ -70,7 +70,7 @@ class Connection : public util::Connection {
 
   // PubSub message, either incoming message for active subscription or reply for new subscription.
   struct PubMessage {
-    std::string pattern{};              // non-empty for pattern subscriber
+    std::string pattern;                // non-empty for pattern subscriber
     std::shared_ptr<char[]> buf;        // stores channel name and message
     std::string_view channel, message;  // channel and message parts from buf
     bool should_unsubscribe = false;    // unsubscribe from channel after sending the message
@@ -176,15 +176,16 @@ class Connection : public util::Connection {
   static_assert(sizeof(MessageHandle) <= 80,
                 "Big structs should use indirection to avoid wasting deque space!");
 
-  enum Phase { SETUP, READ_SOCKET, PROCESS, SHUTTING_DOWN, PRECLOSE, NUM_PHASES };
+  enum Phase : uint8_t { SETUP, READ_SOCKET, PROCESS, SHUTTING_DOWN, PRECLOSE, NUM_PHASES };
 
   // Weak reference to a connection, invalidated upon connection close.
   // Used to dispatch async operations for the connection without worrying about pointer lifetime.
   struct WeakRef {
    public:
     // Get residing thread of connection. Thread-safe.
-    unsigned Thread() const;
-
+    unsigned LastKnownThreadId() const {
+      return last_known_thread_id_;
+    }
     // Get pointer to connection if still valid, nullptr if expired.
     // Can only be called from connection's thread. Validity is guaranteed
     // only until the next suspension point.
@@ -205,7 +206,7 @@ class Connection : public util::Connection {
     WeakRef(std::shared_ptr<Connection> ptr, unsigned thread_id, uint32_t client_id);
 
     std::weak_ptr<Connection> ptr_;
-    unsigned thread_id_;
+    unsigned last_known_thread_id_;
     uint32_t client_id_;
   };
 
@@ -288,8 +289,9 @@ class Connection : public util::Connection {
   ConnectionContext* cntx();
 
   // Requests that at some point, this connection will be migrated to `dest` thread.
-  // Connections will migrate at most once, and only when the flag --migrate_connections is true.
-  void RequestAsyncMigration(util::fb2::ProactorBase* dest);
+  // If force is false, the connection will migrate at most once,
+  // and only when the flag --migrate_connections is true.
+  void RequestAsyncMigration(util::fb2::ProactorBase* dest, bool force);
 
   // Starts traffic logging in the calling thread. Must be a proactor thread.
   // Each thread creates its own log file combining requests from all the connections in
@@ -335,7 +337,7 @@ class Connection : public util::Connection {
   std::unique_ptr<ConnectionContext> cc_;  // Null for http connections
 
  private:
-  enum ParserStatus { OK, NEED_MORE, ERROR };
+  enum ParserStatus : uint8_t { OK, NEED_MORE, ERROR };
 
   struct AsyncOperations;
 
@@ -443,6 +445,7 @@ class Connection : public util::Connection {
 
   uint32_t id_;
   Protocol protocol_;
+  Phase phase_ = SETUP;
 
   struct {
     size_t read_cnt = 0;                // total number of read calls
@@ -459,7 +462,6 @@ class Connection : public util::Connection {
   ServiceInterface* service_;
 
   time_t creation_time_, last_interaction_;
-  Phase phase_ = SETUP;
   std::string name_;
 
   std::string lib_name_;

src/server/channel_store.cc

@@ -48,11 +48,11 @@ auto BuildSender(string_view channel, facade::ArgRange messages, bool unsubscrib
 }  // namespace
 
 bool ChannelStore::Subscriber::ByThread(const Subscriber& lhs, const Subscriber& rhs) {
-  return ByThreadId(lhs, rhs.Thread());
+  return ByThreadId(lhs, rhs.LastKnownThreadId());
 }
 
 bool ChannelStore::Subscriber::ByThreadId(const Subscriber& lhs, const unsigned thread) {
-  return lhs.Thread() < thread;
+  return lhs.LastKnownThreadId() < thread;
 }
 
 ChannelStore::UpdatablePointer::UpdatablePointer(const UpdatablePointer& other) {
@@ -128,7 +128,7 @@ unsigned ChannelStore::SendMessages(std::string_view channel, facade::ArgRange m
   int32_t last_thread = -1;
 
   for (auto& sub : subscribers) {
-    int sub_thread = sub.Thread();
+    int sub_thread = sub.LastKnownThreadId();
     DCHECK_LE(last_thread, sub_thread);
     if (last_thread == sub_thread)  // skip same thread
       continue;
@@ -139,15 +139,15 @@ unsigned ChannelStore::SendMessages(std::string_view channel, facade::ArgRange m
     // Make sure the connection thread has enough memory budget to accept the message.
     // This is a heuristic and not entirely hermetic since the connection memory might
     // get filled again.
-    facade::Connection::EnsureMemoryBudget(sub.Thread());
+    facade::Connection::EnsureMemoryBudget(sub_thread);
     last_thread = sub_thread;
   }
 
   auto subscribers_ptr = make_shared<decltype(subscribers)>(std::move(subscribers));
   auto cb = [subscribers_ptr, send = BuildSender(channel, messages)](unsigned idx, auto*) {
     auto it = lower_bound(subscribers_ptr->begin(), subscribers_ptr->end(), idx,
                           ChannelStore::Subscriber::ByThreadId);
-    while (it != subscribers_ptr->end() && it->Thread() == idx) {
+    while (it != subscribers_ptr->end() && it->LastKnownThreadId() == idx) {
       if (auto* ptr = it->Get(); ptr && ptr->cntx() != nullptr)
         send(ptr, it->pattern);
       it++;
@@ -227,7 +227,7 @@ void ChannelStore::UnsubscribeConnectionsFromDeletedSlots(const ChannelsSubMap&
 
     auto it = lower_bound(subscribers.begin(), subscribers.end(), idx,
                           ChannelStore::Subscriber::ByThreadId);
-    while (it != subscribers.end() && it->Thread() == idx) {
+    while (it != subscribers.end() && it->LastKnownThreadId() == idx) {
       // if ptr->cntx() is null, a connection might have closed or be in the process of closing
       if (auto* ptr = it->Get(); ptr && ptr->cntx() != nullptr) {
         DCHECK(it->pattern.empty());

src/server/db_slice.cc

@@ -1569,19 +1569,20 @@ void DbSlice::QueueInvalidationTrackingMessageAtomic(std::string_view key) {
   auto [pend_it, inserted] = pending_send_map_.emplace(key, std::move(moved_set));
   if (!inserted) {
     ConnectionHashSet& client_set = pend_it->second;
-    for (auto& client : moved_set) {
-      client_set.insert(client);
+    for (auto& weak_ref : moved_set) {
+      client_set.insert(weak_ref);
     }
   }
 }
 
-void DbSlice::SendQueuedInvalidationMessagesCb(const TrackingMap& track_map, unsigned idx) const {
+void DbSlice::SendQueuedInvalidationMessagesCb(const TrackingMap& track_map,
+                                               unsigned calling_thread_id) const {
   for (auto& [key, client_list] : track_map) {
-    for (auto& client : client_list) {
-      if (client.IsExpired() || (client.Thread() != idx)) {
-        continue;
+    for (auto& weak_ref : client_list) {
+      if (weak_ref.IsExpired() || (weak_ref.LastKnownThreadId() != calling_thread_id)) {
+        continue;  // Expired or migrated.
       }
-      auto* conn = client.Get();
+      auto* conn = weak_ref.Get();
       auto* cntx = static_cast<ConnectionContext*>(conn->cntx());
       if (cntx && cntx->conn_state.tracking_info_.IsTrackingOn()) {
         conn->SendInvalidationMessageAsync({key});
@@ -1597,8 +1598,8 @@ void DbSlice::SendQueuedInvalidationMessages() {
     // Notify all the clients. this function is not efficient,
     // because it broadcasts to all threads unrelated to the subscribers for the key.
     auto local_map = std::move(pending_send_map_);
-    auto cb = [&](unsigned idx, util::ProactorBase*) {
-      SendQueuedInvalidationMessagesCb(local_map, idx);
+    auto cb = [&](unsigned thread_id, util::ProactorBase*) {
+      SendQueuedInvalidationMessagesCb(local_map, thread_id);
     };
 
     shard_set->pool()->AwaitBrief(std::move(cb));

src/server/main_service.cc

@@ -1378,8 +1378,7 @@ OpResult<void> OpTrackKeys(const OpArgs slice_args, const facade::Connection::We
     return OpStatus::OK;
   }
 
-  DVLOG(2) << "Start tracking keys for client ID: " << conn_ref.GetClientId()
-           << " with thread ID: " << conn_ref.Thread();
+  DVLOG(2) << "Start tracking keys for client ID: " << conn_ref.GetClientId();
 
   auto& db_slice = slice_args.GetDbSlice();
   // TODO: There is a bug here that we track all arguments instead of tracking only keys.
@@ -2125,7 +2124,7 @@ void Service::EvalInternal(CmdArgList args, const EvalArgs& eval_args, Interpret
     if (*sid != ServerState::tlocal()->thread_index()) {
       VLOG(2) << "Migrating connection " << cntx->conn() << " from "
               << ProactorBase::me()->GetPoolIndex() << " to " << *sid;
-      cntx->conn()->RequestAsyncMigration(shard_set->pool()->at(*sid));
+      cntx->conn()->RequestAsyncMigration(shard_set->pool()->at(*sid), false);
     }
   } else {
     Transaction::MultiMode script_mode = DetermineMultiMode(*params);

src/server/server_family.cc

@@ -596,6 +596,49 @@ void ClientKill(CmdArgList args, absl::Span<facade::Listener*> listeners, SinkRe
   }
 }
 
+void ClientMigrate(CmdArgList args, absl::Span<facade::Listener*> listeners,
+                   SinkReplyBuilder* builder, ConnectionContext* cntx) {
+  if (args.size() != 2) {
+    return builder->SendError(kSyntaxErr);
+  }
+
+  uint32_t id;
+  if (!absl::SimpleAtoi(args[0], &id)) {
+    return builder->SendError("Invalid client id");
+  }
+
+  uint32_t tid = 0;
+  if (!absl::SimpleAtoi(args[1], &tid) || tid >= shard_set->pool()->size()) {
+    return builder->SendError("Invalid thread id");
+  }
+
+  unsigned migrated = 0;
+  auto cb_brief = [&](unsigned current_tid, ProactorBase* p) {
+    if (current_tid == tid) {
+      return;  // we should not migrate to the same thread
+    }
+
+    auto traverse_cb = [&](unsigned, util::Connection* conn) {
+      facade::Connection* dconn = static_cast<facade::Connection*>(conn);
+      if (dconn->GetClientId() == id) {
+        ++migrated;
+        dconn->RequestAsyncMigration(shard_set->pool()->at(tid), true /* force */);
+      }
+    };
+
+    for (auto* listener : listeners) {
+      if (listener->IsPrivilegedInterface())
+        continue;  // skip privileged interfaces
+
+      listener->TraverseConnectionsOnThread(traverse_cb, UINT32_MAX, nullptr);
+    }
+  };
+
+  shard_set->pool()->AwaitBrief(cb_brief);
+
+  return builder->SendLong(migrated);
+}
+
 std::string_view GetOSString() {
   // Call uname() only once since it can be expensive. Cache the final result in a static string.
   static string os_string = []() {
@@ -2096,6 +2139,8 @@ void ClientHelp(SinkReplyBuilder* builder) {
       "    * LIB-VER: the client lib version.",
       "TRACKING (ON|OFF) [OPTIN] [OPTOUT] [NOLOOP]",
       "    Control server assisted client side caching.",
+      "MIGRATE <client-id> <tid>",
+      "    Migrates connection specified by client-id to the specified thread id.",
       "HELP",
       "    Print this help."};
   auto* rb = static_cast<RedisReplyBuilder*>(builder);
@@ -2128,6 +2173,8 @@ void ServerFamily::Client(CmdArgList args, const CommandContext& cmd_cntx) {
     return ClientSetInfo(sub_args, builder, cntx);
   } else if (sub_cmd == "ID") {
     return ClientId(sub_args, builder, cntx);
+  } else if (sub_cmd == "MIGRATE") {
+    return ClientMigrate(sub_args, absl::MakeSpan(listeners_), builder, cntx);
   } else if (sub_cmd == "HELP") {
     return ClientHelp(builder);
   }