feat: add client command

src/base_cmd.h

@@ -76,6 +76,13 @@ const std::string kCmdNameUnwatch = "unwatch";
 const std::string kCmdNameDiscard = "discard";
 
 // admin
+const std::string kCmdNameClient = "client";
+const std::string kSubCmdNameClientGetname = "getname";
+const std::string kSubCmdNameClientSetname = "setname";
+const std::string kSubCmdNameClientId = "id";
+const std::string kSubCmdNameClientList = "list";
+const std::string kSubCmdNameClientKill = "kill";
+
 const std::string kCmdNameConfig = "config";
 const std::string kSubCmdNameConfigGet = "get";
 const std::string kSubCmdNameConfigSet = "set";

src/client.cc

@@ -21,6 +21,8 @@
 
 namespace pikiwidb {
 
+const ClientInfo ClientInfo::invalidClientInfo = {-1, "", -1, -1};
+
 void CmdRes::RedisAppendLen(std::string& str, int64_t ori, const std::string& prefix) {
   str.append(prefix);
   str.append(pstd::Int2string(ori));
@@ -487,6 +489,13 @@ int PClient::PeerPort() const {
   return -1;
 }
 
+const int PClient::GetFd() const {
+  if (auto c = getTcpConnection(); c) {
+    return c->Fd();
+  }
+  return -1;
+}
+
 bool PClient::SendPacket(const std::string& buf) {
   if (auto c = getTcpConnection(); c) {
     return c->SendPacket(buf);
@@ -549,27 +558,34 @@ bool PClient::isClusterCmdTarget() const {
   return PRAFT.GetClusterCmdCtx().GetPeerIp() == PeerIP() && PRAFT.GetClusterCmdCtx().GetPort() == PeerPort();
 }
 
-int PClient::uniqueID() const {
+int PClient::GetUniqueId() const {
   if (auto c = getTcpConnection(); c) {
     return c->GetUniqueId();
   }
 
   return -1;
 }
 
+ClientInfo PClient::GetClientInfo() const {
+  if (auto c = getTcpConnection(); c) {
+    return {GetUniqueId(), PeerIP().c_str(), PeerPort(), GetFd()};
+  }
+  return ClientInfo::invalidClientInfo;
+}
+
 bool PClient::Watch(int dbno, const std::string& key) {
   DEBUG("Client {} watch {}, db {}", name_, key, dbno);
   return watch_keys_[dbno].insert(key).second;
 }
 
 bool PClient::NotifyDirty(int dbno, const std::string& key) {
   if (IsFlagOn(kClientFlagDirty)) {
-    INFO("client is already dirty {}", uniqueID());
+    INFO("client is already dirty {}", GetUniqueId());
     return true;
   }
 
   if (watch_keys_[dbno].contains(key)) {
-    INFO("{} client become dirty because key {} in db {}", uniqueID(), key, dbno);
+    INFO("{} client become dirty because key {} in db {}", GetUniqueId(), key, dbno);
     SetFlag(kClientFlagDirty);
     return true;
   } else {

src/client.h

@@ -19,7 +19,6 @@
 #include "storage/storage.h"
 
 namespace pikiwidb {
-
 class CmdRes {
  public:
   enum CmdRet {
@@ -111,6 +110,14 @@ enum class ClientState {
 class DB;
 struct PSlaveInfo;
 
+struct ClientInfo {
+  int client_id;
+  std::string ip;
+  int port;
+  int fd;
+  static const ClientInfo invalidClientInfo;
+  bool operator==(ClientInfo& ci) const { return client_id == ci.client_id; }
+};
 class PClient : public std::enable_shared_from_this<PClient>, public CmdRes {
  public:
   PClient() = delete;
@@ -122,6 +129,9 @@ class PClient : public std::enable_shared_from_this<PClient>, public CmdRes {
 
   const std::string& PeerIP() const;
   int PeerPort() const;
+  const int GetFd() const;
+  int GetUniqueId() const;
+  ClientInfo GetClientInfo() const;
 
   bool SendPacket(const std::string& buf);
   bool SendPacket(const void* data, size_t size);
@@ -223,7 +233,6 @@ class PClient : public std::enable_shared_from_this<PClient>, public CmdRes {
   int processInlineCmd(const char*, size_t, std::vector<std::string>&);
   void reset();
   bool isPeerMaster() const;
-  int uniqueID() const;
 
   bool isClusterCmdTarget() const;
 

src/client_map.cc

@@ -0,0 +1,103 @@
+#include "client_map.h"
+#include "log.h"
+
+namespace pikiwidb {
+
+uint32_t ClientMap::GetAllClientInfos(std::vector<ClientInfo>& results) {
+  // client info string type: ip, port, fd.
+  std::shared_lock<std::shared_mutex> client_map_lock(client_map_mutex_);
+  for (auto& [id, client_weak] : clients_) {
+    if (auto client = client_weak.lock()) {
+      results.emplace_back(client->GetClientInfo());
+    }
+  }
+  return results.size();
+}
+
+bool ClientMap::AddClient(int id, std::weak_ptr<PClient> client) {
+  std::unique_lock client_map_lock(client_map_mutex_);
+  if (clients_.find(id) == clients_.end()) {
+    clients_.insert({id, client});
+    return true;
+  }
+  return false;
+}
+
+ClientInfo ClientMap::GetClientsInfoById(int id) {
+  std::shared_lock client_map_lock(client_map_mutex_);
+  if (auto it = clients_.find(id); it != clients_.end()) {
+    if (auto client = it->second.lock(); client) {
+      return client->GetClientInfo();
+    }
+  }
+  ERROR("Client with ID {} not found in GetClientsInfoById", id);
+  return ClientInfo::invalidClientInfo;
+}
+
+bool ClientMap::RemoveClientById(int id) {
+  std::unique_lock client_map_lock(client_map_mutex_);
+  if (auto it = clients_.find(id); it != clients_.end()) {
+    clients_.erase(it);
+    INFO("Removed client with ID {}", id);
+    return true;
+  }
+  return false;
+}
+
+bool ClientMap::KillAllClients() {
+  std::vector<std::shared_ptr<PClient>> clients_to_close;
+  {
+    std::shared_lock<std::shared_mutex> client_map_lock(client_map_mutex_);
+    for (auto& [id, client_weak] : clients_) {
+      if (auto client = client_weak.lock()) {
+        clients_to_close.push_back(client);
+      }
+    }
+  }
+  for (auto& client : clients_to_close) {
+    client->Close();
+  }
+  return true;
+}
+
+bool ClientMap::KillClientByAddrPort(const std::string& addr_port) {
+  std::shared_ptr<PClient> client_to_close;
+  {
+    std::shared_lock<std::shared_mutex> client_map_lock(client_map_mutex_);
+    for (auto& [id, client_weak] : clients_) {
+      if (auto client = client_weak.lock()) {
+        std::string client_ip_port = client->PeerIP() + ":" + std::to_string(client->PeerPort());
+        if (client_ip_port == addr_port) {
+          client_to_close = client;
+          break;
+        }
+      }
+    }
+  }
+  if (client_to_close) {
+    client_to_close->Close();
+    return true;
+  }
+  return false;
+}
+
+bool ClientMap::KillClientById(int client_id) {
+  std::shared_ptr<PClient> client_to_close;
+  {
+    std::shared_lock<std::shared_mutex> client_map_lock(client_map_mutex_);
+    if (auto it = clients_.find(client_id); it != clients_.end()) {
+      if (auto client = it->second.lock()) {
+        client_to_close = client;
+      }
+    }
+  }
+  if (client_to_close) {
+    INFO("Closing client with ID {}", client_id);
+    client_to_close->Close();
+    INFO("Client with ID {} closed", client_id);
+    return true;
+  }
+  return false;
+}
+
+}  // namespace pikiwidb

src/client_map.h

@@ -0,0 +1,41 @@
+#pragma once
+
+#include <map>
+#include <memory>
+#include <shared_mutex>
+#include <string>
+#include "client.h"
+
+namespace pikiwidb {
+class ClientMap {
+ private:
+  ClientMap() = default;
+  // 禁用复制构造函数和赋值运算符
+
+ private:
+  std::map<int, std::weak_ptr<PClient>> clients_;
+  std::shared_mutex client_map_mutex_;
+
+ public:
+  static ClientMap& getInstance() {
+    static ClientMap instance;
+    return instance;
+  }
+
+  ClientMap(const ClientMap&) = delete;
+  ClientMap& operator=(const ClientMap&) = delete;
+
+  // client info function
+  pikiwidb::ClientInfo GetClientsInfoById(int id);
+  uint32_t GetAllClientInfos(std::vector<ClientInfo>& results);
+
+  bool AddClient(int id, std::weak_ptr<PClient>);
+
+  bool RemoveClientById(int id);
+
+  bool KillAllClients();
+  bool KillClientById(int client_id);
+  bool KillClientByAddrPort(const std::string& addr_port);
+};
+
+}  // namespace pikiwidb