Max connection limit

src/main.cpp

@@ -20,6 +20,7 @@
 #include <vector>
 
 std::atomic<bool> shutdown_requested{false};
+std::atomic<int> activeConnections{0};
 
 #ifndef __has_feature
 #define __has_feature(x) 0
@@ -124,9 +125,12 @@ struct Connection {
   struct sockaddr_storage addr; // sockaddr_storage handles IPv4/IPv6
 
   Connection(struct sockaddr_storage addr, int fd, int64_t id)
-      : fd(fd), id(id), addr(addr) {}
+      : fd(fd), id(id), addr(addr) {
+    activeConnections.fetch_add(1, std::memory_order_relaxed);
+  }
 
   ~Connection() {
+    activeConnections.fetch_sub(1, std::memory_order_relaxed);
     int e = close(fd);
     if (e == -1) {
       perror("close");
@@ -237,6 +241,8 @@ int main(int argc, char *argv[]) {
             << std::endl;
   std::cout << "Event batch size: " << config->server.event_batch_size
             << std::endl;
+  std::cout << "Max connections: " << config->server.max_connections
+            << std::endl;
   std::cout << "Min request ID length: " << config->commit.min_request_id_length
             << std::endl;
   std::cout << "Request ID retention: "
@@ -262,6 +268,8 @@ int main(int argc, char *argv[]) {
     abort();
   }
 
+  std::atomic<int64_t> connectionId{0};
+
   // Network threads from configuration
   int networkThreads = config->server.network_threads;
 
@@ -340,12 +348,11 @@ int main(int argc, char *argv[]) {
         });
   }
 
-  std::atomic<int64_t> connectionId{0};
-
   // Accept threads from configuration
   int acceptThreads = config->server.accept_threads;
   for (int i = 0; i < acceptThreads; ++i) {
-    threads.emplace_back([epollfd, i, sockfd, &connectionId]() {
+    threads.emplace_back([epollfd, i, sockfd, &connectionId,
+                          max_connections = config->server.max_connections]() {
       pthread_setname_np(pthread_self(),
                          ("accept-" + std::to_string(i)).c_str());
       // Call accept in a loop
@@ -361,6 +368,15 @@ int main(int argc, char *argv[]) {
           perror("accept4");
           continue;
         }
+
+        // Check connection limit (0 means unlimited)
+        if (max_connections > 0 &&
+            activeConnections.load(std::memory_order_relaxed) >=
+                max_connections) {
+          // Reject connection by immediately closing it
+          close(fd);
+          continue;
+        }
         auto conn = std::make_unique<Connection>(
             addr, fd, connectionId.fetch_add(1, std::memory_order_relaxed));
         // Transfer new connection to epoll ownership