feat: add client max connection

src/kiwi.cc

@@ -188,6 +188,8 @@ bool KiwiDB::Init() {
   auto num = g_config.worker_threads_num + g_config.slave_threads_num;
   options_.SetThreadNum(num);
 
+  options_.SetMaxClients(g_config.max_clients);
+
   // now we only use fast cmd thread pool
   auto status = cmd_threads_.Init(g_config.fast_cmd_threads_num, 1, "kiwi-cmd");
   if (!status.ok()) {

src/net/event_server.h

@@ -95,6 +95,7 @@ class EventServer final {
   std::vector<std::unique_ptr<ThreadManager<T>>> threadsManager_;
 
   std::mutex mtx_;
+
   std::condition_variable cv_;
 
   std::shared_ptr<Timer> timer_;

src/net/net_options.h

@@ -22,6 +22,9 @@ class NetOptions {
 
   bool GetRwSeparation() const { return rwSeparation_; }
 
+  void SetMaxClients(uint32_t maxClients) { maxClients_ = maxClients; }
+
+  uint32_t GetMaxClients() const { return maxClients_; }
   void SetOpTcpKeepAlive(uint32_t tcpKeepAlive) { tcpKeepAlive_ = tcpKeepAlive; }
 
   uint32_t GetOpTcpKeepAlive() const { return tcpKeepAlive_; }
@@ -31,6 +34,7 @@ class NetOptions {
 
   int8_t threadNum_ = 1;  // The number of threads
 
+  uint32_t maxClients_ = 1;      // The maximum number of connections(default 40000)
   uint32_t tcpKeepAlive_ = 300;  // The timeout of the keepalive connection in seconds
 };
 

src/net/thread_manager.h

@@ -8,6 +8,7 @@
 #pragma once
 
 #include <atomic>
+#include <cstdint>
 #include <memory>
 #include <mutex>
 #include <shared_mutex>
@@ -18,6 +19,8 @@
 #include "io_thread.h"
 #include "net_options.h"
 
+#include "log.h"
+
 #if defined(HAVE_EPOLL)
 
 #  include "epoll_event.h"
@@ -90,18 +93,24 @@ class ThreadManager {
 
   uint64_t DoTCPConnect(T &t, int fd, const std::shared_ptr<Connection> &conn);
 
+  uint32_t get_client_count() const { return clientCount_.load(); }
+
+  void client_count_decrement() { clientCount_.fetch_sub(1, std::memory_order_relaxed); }
+
  private:
   const int8_t index_ = 0;            // The index of the thread
   uint32_t tcpKeepAlive_ = 300;       // The timeout of the keepalive connection in seconds
   std::atomic<bool> running_ = true;  // Whether the thread is running
 
   NetOptions netOptions_;
 
+  inline static std::atomic<uint32_t> clientCount_{0};
+
   std::unique_ptr<IOThread> readThread_;   // Read thread
   std::unique_ptr<IOThread> writeThread_;  // Write thread
 
-  // All connections for the current thread
-  std::unordered_map<uint64_t, std::pair<T, std::shared_ptr<Connection>>> connections_;
+  std::unordered_map<uint64_t, std::pair<T, std::shared_ptr<Connection>>>
+      connections_;  // All connections for the current thread
 
   std::shared_mutex mutex_;
 
@@ -149,6 +158,22 @@ void ThreadManager<T>::Stop() {
 template <typename T>
 requires HasSetFdFunction<T>
 void ThreadManager<T>::OnNetEventCreate(int fd, const std::shared_ptr<Connection> &conn) {
+  uint32_t expected = get_client_count();
+  if (!clientCount_.compare_exchange_strong(expected, expected + 1, std::memory_order_seq_cst,
+                                            std::memory_order_seq_cst) ||
+      expected >= netOptions_.GetMaxClients()) {
+    INFO("Max client connections, refuse new connection fd: %d", fd);
+    std::string response = "-ERR max clients reached\r\n";
+    ssize_t sent = ::send(fd, response.c_str(), response.size(), 0);
+    if (sent < 0) {
+      ERROR("Failed to send error response to fd: %d, errno: %d", fd, errno);
+    }
+    if (::close(fd) < 0) {
+      ERROR("Failed to close fd: %d, errno: %d", fd, errno);
+    }
+    return;
+  }
+
   T t;
   onInit_(&t);
   auto connId = getConnId();
@@ -206,6 +231,7 @@ void ThreadManager<T>::OnNetEventClose(uint64_t connId, std::string &&err) {
   iter->second.second->netEvent_->Close();  // close socket
   onClose_(iter->second.first, std::move(err));
   connections_.erase(iter);
+  client_count_decrement();
 }
 
 template <typename T>