Connection registry

Now we can use leak sanitizer. Yay!

src/server.cpp

@@ -1,10 +1,12 @@
 #include "server.hpp"
 #include "connection.hpp"
+#include "connection_registry.hpp"
 #include <csignal>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <fcntl.h>
+#include <memory>
 #include <netdb.h>
 #include <netinet/tcp.h>
 #include <pthread.h>
@@ -25,9 +27,18 @@ std::shared_ptr<Server> Server::create(const weaseldb::Config &config,
 }
 
 Server::Server(const weaseldb::Config &config, ConnectionHandler &handler)
-    : config_(config), handler_(handler) {}
+    : config_(config), handler_(handler), connection_registry_() {}
 
-Server::~Server() { cleanup_resources(); }
+Server::~Server() {
+  // CRITICAL: All I/O threads are guaranteed to be joined before the destructor
+  // is called because they are owned by the run() method's call frame.
+  // This eliminates any possibility of race conditions during connection
+  // cleanup.
+
+  // Clean up any remaining connections using proper ordering
+  connection_registry_.shutdown_cleanup();
+  cleanup_resources();
+}
 
 void Server::run() {
   setup_shutdown_pipe();
@@ -36,12 +47,21 @@ void Server::run() {
 
   create_epoll_instances();
 
-  start_io_threads();
+  // Create I/O threads locally in this call frame
+  // CRITICAL: By owning threads in run()'s call frame, we guarantee they are
+  // joined before run() returns, eliminating any race conditions in ~Server()
+  std::vector<std::thread> threads;
+  start_io_threads(threads);
 
-  // Wait for all threads to complete
-  for (auto &thread : threads_) {
+  // Wait for all threads to complete before returning
+  // This ensures all I/O threads are fully stopped before the Server
+  // destructor can be called, preventing race conditions during connection
+  // cleanup
+  for (auto &thread : threads) {
     thread.join();
   }
+
+  // At this point, all threads are joined and it's safe to destroy the Server
 }
 
 void Server::shutdown() {
@@ -285,11 +305,11 @@ int Server::get_epoll_for_thread(int thread_id) const {
   return epoll_fds_[thread_id % epoll_fds_.size()];
 }
 
-void Server::start_io_threads() {
+void Server::start_io_threads(std::vector<std::thread> &threads) {
   int io_threads = config_.server.io_threads;
 
   for (int thread_id = 0; thread_id < io_threads; ++thread_id) {
-    threads_.emplace_back([this, thread_id]() {
+    threads.emplace_back([this, thread_id]() {
       pthread_setname_np(pthread_self(),
                          ("io-" + std::to_string(thread_id)).c_str());
 
@@ -325,16 +345,20 @@ void Server::start_io_threads() {
           }
 
           // Handle existing connection events
-          std::unique_ptr<Connection> conn{
-              static_cast<Connection *>(events[i].data.ptr)};
-          conn->tsan_acquire();
-          events[i].data.ptr = nullptr;
-
-          if (events[i].events & EPOLLERR) {
-            continue; // Connection closed - unique_ptr destructor cleans up
+          std::unique_ptr<Connection> conn;
+          {
+            // borrowed
+            Connection *conn_ = static_cast<Connection *>(events[i].data.ptr);
+            conn_->tsan_acquire();
+            conn = connection_registry_.remove(conn_->getFd());
           }
 
-          // Add to regular batch - I/O will be processed in batch
+          if (events[i].events & (EPOLLERR | EPOLLHUP)) {
+            // unique_ptr will automatically delete on scope exit
+            continue;
+          }
+
+          // Transfer ownership from registry to batch processing
           batch[batch_count] = std::move(conn);
           batch_events[batch_count] = events[i].events;
           batch_count++;
@@ -378,11 +402,11 @@ void Server::start_io_threads() {
               perror("setsockopt SO_KEEPALIVE");
             }
 
-            // Add to batch - I/O will be processed in batch
-            batch[batch_count] = Connection::createForServer(
+            // Transfer ownership from registry to batch processing
+            batch[batch_count] = std::unique_ptr<Connection>(new Connection(
                 addr, fd,
                 connection_id_.fetch_add(1, std::memory_order_relaxed),
-                &handler_, weak_from_this());
+                &handler_, weak_from_this()));
             batch_events[batch_count] =
                 EPOLLIN; // New connections always start with read
             batch_count++;
@@ -473,28 +497,30 @@ void Server::process_connection_batch(
     }
   }
 
-  // Call post-batch handler
+  // Call post-batch handler - handlers can take ownership here
   handler_.on_post_batch(batch);
 
   // Transfer all remaining connections back to epoll
-  for (auto &conn : batch) {
-    if (conn) {
+  for (auto &conn_ptr : batch) {
+    if (conn_ptr) {
+      int fd = conn_ptr->getFd();
+
       struct epoll_event event{};
-      if (!conn->hasMessages()) {
+      if (!conn_ptr->hasMessages()) {
         event.events = EPOLLIN | EPOLLONESHOT;
       } else {
         event.events = EPOLLOUT | EPOLLONESHOT;
       }
 
-      int fd = conn->getFd();
-      conn->tsan_release();
-      Connection *raw_conn = conn.release();
-      event.data.ptr = raw_conn;
-
+      conn_ptr->tsan_release();
+      event.data.ptr = conn_ptr.get(); // Use raw pointer for epoll
+      // Put connection back in registry since handler didn't take ownership.
+      // Must happen before epoll_ctl
+      connection_registry_.store(fd, std::move(conn_ptr));
       int epoll_op = is_new ? EPOLL_CTL_ADD : EPOLL_CTL_MOD;
       if (epoll_ctl(epollfd, epoll_op, fd, &event) == -1) {
         perror(is_new ? "epoll_ctl ADD" : "epoll_ctl MOD");
-        delete raw_conn;
+        (void)connection_registry_.remove(fd);
       }
     }
   }