Add coordinator connection reusing (#6179)

* feat: add coordinator connections reusing

Author: BorysTheDev

src/server/CMakeLists.txt

@@ -82,7 +82,7 @@ add_library(dragonfly_lib
             server_family.cc string_family.cc list_family.cc generic_family.cc
             ${DF_FAMILY_SRCS}
             main_service.cc memory_cmd.cc rdb_load.cc rdb_save.cc replica.cc http_api.cc
-            protocol_client.cc
+            protocol_client.cc owned_resp_expr.cc
             snapshot.cc script_mgr.cc
             detail/compressor.cc detail/decompress.cc detail/save_stages_controller.cc detail/snapshot_storage.cc
             version.cc container_utils.cc

src/server/cluster/cluster_family.cc

@@ -16,6 +16,7 @@
 #include "facade/error.h"
 #include "server/acl/acl_commands_def.h"
 #include "server/channel_store.h"
+#include "server/cluster/coordinator.h"
 #include "server/command_registry.h"
 #include "server/conn_context.h"
 #include "server/dflycmd.h"
@@ -80,6 +81,7 @@ ClusterFamily::ClusterFamily(ServerFamily* server_family) : server_family_(serve
 }
 
 void ClusterFamily::Shutdown() {
+  Coordinator::Current().Shutdown();
   shard_set->pool()->at(0)->Await([this]() ABSL_LOCKS_EXCLUDED(set_config_mu) {
     PreparedToRemoveOutgoingMigrations outgoing_migrations;  // should be removed without mutex lock
     {

src/server/cluster/coordinator.cc

@@ -14,79 +14,189 @@ using namespace facade;
 
 namespace dfly::cluster {
 
-class Coordinator::CrossShardClient : private ProtocolClient {
+class Coordinator::CrossShardRequest {
+ public:
+  CrossShardRequest(std::string cmd, Coordinator::RespCB cb, uint32_t total_shards)
+      : command_(std::move(cmd)), cb_(std::move(cb)), shard_processed_(total_shards) {
+  }
+
+  const std::string& GetCommand() const {
+    return command_;
+  }
+
+  template <class... Args> void Exec(Args&&... args) {
+    cb_(std::forward<Args>(args)...);
+    if (shard_processed_.fetch_sub(1, std::memory_order_relaxed) == 1) {
+      future_.Resolve(GenericError{});
+    }
+  }
+
+  util::fb2::Future<GenericError>& GetFuture() {
+    return future_;
+  }
+
+ private:
+  std::string command_;
+  Coordinator::RespCB cb_;
+  util::fb2::Future<GenericError> future_;
+  std::atomic_uint32_t shard_processed_;
+};
+
+class Coordinator::CrossShardClient : public ProtocolClient {
  public:
   CrossShardClient(std::string host, uint16_t port) : ProtocolClient(std::move(host), port) {
   }
 
   using ProtocolClient::CloseSocket;
   ~CrossShardClient() {
+    exec_st_.Cancel();
+    waker_.notifyAll();
     CloseSocket();
+    send_fb_.Join();
+    resp_fb_.Join();
   }
 
-  void Init() {
+  [[nodiscard]] bool Init() {
     VLOG(1) << "Resolving host DNS to " << server().Description();
     if (error_code ec = ResolveHostDns(); ec) {
       LOG(WARNING) << "Could not resolve host DNS to " << server().Description() << ": "
                    << ec.message();
       exec_st_.ReportError(GenericError(ec, "Could not resolve host dns."));
-      return;
+      return false;
     }
     VLOG(1) << "Start coordinator connection to " << server().Description();
     auto timeout = 3000ms;  // TODO add flag;
     if (auto ec = ConnectAndAuth(timeout, &exec_st_); ec) {
       LOG(WARNING) << "Couldn't connect to " << server().Description() << ": " << ec.message()
                    << ", socket state: " << GetSocketInfo(Sock()->native_handle());
       exec_st_.ReportError(GenericError(ec, "Couldn't connect to source."));
-      return;
+      return false;
     }
 
     ResetParser(RedisParser::Mode::CLIENT);
+    send_fb_ = util::fb2::Fiber("CSS_SendFb", &CrossShardClient::SendFb, this);
+    resp_fb_ = util::fb2::Fiber("CSS_RespFb", &CrossShardClient::RespFb, this);
+    return true;
   }
 
   void Cancel() {
+    exec_st_.Cancel();
     ShutdownSocket();
   }
 
-  void SendCommand(std::string_view cmd, const RespCB& cb) {
-    if (auto ec = ProtocolClient::SendCommand(cmd); ec) {
-      LOG(WARNING) << "Coordinator could not send command to : " << server().Description() << ": "
-                   << ec.message() << ", socket state: " << GetSocketInfo(Sock()->native_handle());
-      exec_st_.ReportError(GenericError(ec, "Could not send command."));
+  void EnqueueCommand(CrossShardRequestPtr req) {
+    std::lock_guard lk(mu_);
+    send_queue_.push(req);
+    resp_queue_.push(req);
+    ready_to_send_ = true;
+    waker_.notifyAll();
+  }
+
+  void SendFb() {
+    while (!exec_st_.IsCancelled()) {
+      waker_.await([this] { return exec_st_.IsCancelled() || ready_to_send_; });
+      if (exec_st_.IsCancelled())
+        return;
+      std::lock_guard lk(mu_);
+      while (!send_queue_.empty()) {
+        if (auto ec = ProtocolClient::SendCommand(send_queue_.front()->GetCommand()); ec) {
+          exec_st_.ReportError(GenericError(
+              ec, absl::StrCat("Coordinator could not send command to : ", server().Description(),
+                               "socket state: ", GetSocketInfo(Sock()->native_handle()))));
+          // TODO reinit connection.
+          break;
+        }
+        send_queue_.pop();
+      }
+      ready_to_send_ = false;
     }
-    auto timeout = 30000;  // TODO add flag;
-    if (auto resp = ReadRespReply(timeout); !resp) {
-      LOG(WARNING) << "Error reading response from " << server().Description() << ": "
-                   << resp.error() << ", socket state: " + GetSocketInfo(Sock()->native_handle());
+  }
+
+  void RespFb() {
+    while (!exec_st_.IsCancelled()) {
+      waker_.await([this] { return exec_st_.IsCancelled() || ready_to_send_; });
+      if (exec_st_.IsCancelled())
+        return;
+      std::lock_guard lk(mu_);
+      constexpr auto timeout = 3000;  // TODO add flag and add usage in ReadRespReply.
+      while (!resp_queue_.empty()) {
+        auto resp = TakeRespReply(timeout);
+        if (!resp) {
+          LOG(WARNING) << "Error reading response from " << server().Description() << ": "
+                       << resp.error()
+                       << ", socket state: " + GetSocketInfo(Sock()->native_handle());
+
+          // TODO make all requests fail in this case.
+          // TODO reinit connection.
+          LOG(FATAL) << "Coordinator RespFb read error, not implemented recovery yet.";
+          break;
+        }
+        resp_queue_.front()->Exec(*resp);
+        resp_queue_.pop();
+      }
     }
-    cb(LastResponseArgs());
-    // add response processing
   }
+
+ private:
+  std::queue<std::shared_ptr<CrossShardRequest>> send_queue_;
+  std::queue<std::shared_ptr<CrossShardRequest>> resp_queue_;
+
+  util::fb2::Fiber send_fb_;
+  util::fb2::Fiber resp_fb_;
+  util::fb2::EventCount waker_;
+
+  mutable util::fb2::Mutex mu_;
+  std::atomic_bool ready_to_send_ = false;
 };
 
 Coordinator& Coordinator::Current() {
   static Coordinator instance;
   return instance;
 }
 
-void Coordinator::DispatchAll(std::string_view command, RespCB cb) {
+std::shared_ptr<Coordinator::CrossShardClient> Coordinator::GetClient(const std::string& host,
+                                                                      uint16_t port) {
+  for (const auto& client : clients_) {
+    if (client->GetHost() == host && client->GetPort() == port) {
+      return client;
+    }
+  }
+  auto new_client = std::make_shared<CrossShardClient>(host, port);
+  if (new_client->Init()) {
+    clients_.emplace_back(new_client);
+    return new_client;
+  }
+  return nullptr;
+}
+
+util::fb2::Future<GenericError> Coordinator::DispatchAll(std::string command, RespCB cb) {
   auto cluster_config = ClusterConfig::Current();
   if (!cluster_config) {
     VLOG(2) << "No cluster config found for coordinator plan creation.";
-    return;
+    LOG(FATAL) << "No cluster config, not implemented logic yet.";
+    return {};
   }
   VLOG(2) << "Dispatching command to all shards: " << command;
   auto shards_config = cluster_config->GetConfig();
 
-  std::vector<std::unique_ptr<CrossShardClient>> clients;
+  auto shard_request = std::make_shared<CrossShardRequest>(std::move(command), std::move(cb),
+                                                           shards_config.size() - 1);
+
   for (const auto& shard : shards_config) {
     if (shard.master.id == cluster_config->MyId()) {
       continue;
     }
-    clients.emplace_back(std::make_unique<CrossShardClient>(shard.master.ip, shard.master.port));
-    clients.back()->Init();
-    clients.back()->SendCommand(std::string(command), cb);
+    const auto& client = GetClient(shard.master.ip, shard.master.port);
+    if (!client) {
+      VLOG(1) << "Could not get coordinator client for " << shard.master.ip << ":"
+              << shard.master.port;
+      cb({});  // TODO add error propagation.
+      LOG(FATAL) << "No error processing, not implemented logic yet.";
+      return {};
+    }
+    client->EnqueueCommand(shard_request);
   }
+  return shard_request->GetFuture();
 }
 
 }  // namespace dfly::cluster

src/server/cluster/coordinator.h

@@ -5,6 +5,7 @@
 
 #include "server/cluster/cluster_defs.h"
 #include "server/protocol_client.h"
+#include "util/fibers/future.h"
 
 namespace dfly::cluster {
 
@@ -13,14 +14,24 @@ namespace dfly::cluster {
 // It can be used to exeute commands on all shards or specific shards.
 class Coordinator {
  public:
-  using RespCB = std::function<void(const facade::RespVec&)>;
+  using RespCB = std::function<void(const OwnedRespExpr::Vec&)>;  // TODO add error.
 
   static Coordinator& Current();
-  void DispatchAll(std::string_view command, RespCB cb);
+  [[nodiscard]] util::fb2::Future<GenericError> DispatchAll(std::string command, RespCB cb);
+
+  void Shutdown() {
+    // TODO add proper shutdown logic. We need to prevent new clients creation. Maybe we need to
+    // wait destroying of existing clients.
+    clients_.clear();
+  }
 
  private:
   Coordinator() = default;
   class CrossShardClient;
+  class CrossShardRequest;
+  using CrossShardRequestPtr = std::shared_ptr<Coordinator::CrossShardRequest>;
+  std::shared_ptr<CrossShardClient> GetClient(const std::string& host, uint16_t port);
+  std::vector<std::shared_ptr<CrossShardClient>> clients_;
 };
 
 }  // namespace dfly::cluster

src/server/owned_resp_expr.cc

@@ -0,0 +1,79 @@
+// Copyright 2022, DragonflyDB authors.  All rights reserved.
+// See LICENSE for licensing terms.
+//
+
+#include "server/owned_resp_expr.h"
+
+#include <core/overloaded.h>
+
+namespace dfly {
+
+OwnedRespExpr::OwnedRespExpr(const facade::RespExpr& expr) {
+  type = expr.type;
+  std::visit(dfly::Overloaded{
+                 [&](facade::RespExpr::Buffer buf) { u = Buffer(buf.begin(), buf.end()); },
+                 [&](const facade::RespExpr::Vec* vec) {
+                   Vec new_vec;
+                   for (const auto& item : *vec) {
+                     new_vec.emplace_back(item);
+                   }
+                   u = new_vec;
+                 },
+                 [&](int64_t i) { u = i; },
+                 [&](double d) { u = d; },
+             },
+             expr.u);
+}
+
+std::string_view ToSV(OwnedRespExpr::Buffer buf) {
+  return std::string_view{reinterpret_cast<const char*>(buf.data()), buf.size()};
+}
+
+}  // namespace dfly
+
+namespace std {
+ostream& operator<<(ostream& os, const dfly::OwnedRespExpr& e) {
+  using dfly::OwnedRespExpr;
+  using dfly::ToSV;
+
+  switch (e.type) {
+    case OwnedRespExpr::Type::INT64:
+      os << "i" << get<int64_t>(e.u);
+      break;
+    case OwnedRespExpr::Type::DOUBLE:
+      os << "d" << get<double>(e.u);
+      break;
+    case OwnedRespExpr::Type::STRING:
+      os << "'" << ToSV(get<OwnedRespExpr::Buffer>(e.u)) << "'";
+      break;
+    case OwnedRespExpr::Type::NIL:
+      os << "nil";
+      break;
+    case OwnedRespExpr::Type::NIL_ARRAY:
+      os << "[]";
+      break;
+    case OwnedRespExpr::Type::ARRAY:
+      os << dfly::TakeRespSpan{std::get<OwnedRespExpr::Vec>(e.u)};
+      break;
+    case OwnedRespExpr::Type::ERROR:
+      os << "e(" << ToSV(get<OwnedRespExpr::Buffer>(e.u)) << ")";
+      break;
+  }
+
+  return os;
+}
+
+ostream& operator<<(ostream& os, dfly::TakeRespSpan ras) {
+  os << "[";
+  if (!ras.empty()) {
+    for (size_t i = 0; i < ras.size() - 1; ++i) {
+      os << ras[i] << ",";
+    }
+    os << ras.back();
+  }
+  os << "]";
+
+  return os;
+}
+
+}  // namespace std