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