Fix some issues with transferring conn back to server

src/server.cpp

@@ -100,16 +100,19 @@ void Server::releaseBackToServer(std::unique_ptr<Connection> connection) {
 
   // Try to get the server from the connection's weak_ptr
   if (auto server = connection->server_.lock()) {
-    // Server still exists - release raw pointer and let server take over
-    Connection *raw_conn = connection.release();
-    server->receiveConnectionBack(raw_conn);
+    // Server still exists - pass unique_ptr directly
+    server->receiveConnectionBack(std::move(connection));
   }
 
   // If server is gone, connection will be automatically cleaned up when
   // unique_ptr destructs
 }
 
-void Server::receiveConnectionBack(Connection *connection) {
+void Server::receiveConnectionBack(std::unique_ptr<Connection> connection) {
+  if (!connection) {
+    return; // Nothing to process
+  }
+
   // Re-add the connection to epoll for continued processing
   struct epoll_event event{};
 
@@ -120,17 +123,19 @@ void Server::receiveConnectionBack(Connection *connection) {
   }
 
   connection->tsan_release();
-  event.data.fd = connection->getFd();
+  int fd = connection->getFd();
+  event.data.fd = fd;
 
-  // Distribute connections round-robin across epoll instances
-  size_t epoll_index =
-      connection_distribution_counter_.fetch_add(1, std::memory_order_relaxed) %
-      epoll_fds_.size();
+  // Store connection in registry before adding to epoll
+  // This mirrors the pattern used in process_connection_batch
+  size_t epoll_index = connection->getEpollIndex();
   int epollfd = epoll_fds_[epoll_index];
+  connection_registry_.store(fd, std::move(connection));
 
-  if (epoll_ctl(epollfd, EPOLL_CTL_ADD, connection->getFd(), &event) == -1) {
-    perror("epoll_ctl ADD in receiveConnectionBack");
-    delete connection; // Clean up on failure
+  if (epoll_ctl(epollfd, EPOLL_CTL_MOD, fd, &event) == -1) {
+    perror("epoll_ctl MOD in receiveConnectionBack");
+    // Remove from registry and clean up on failure
+    (void)connection_registry_.remove(fd);
   }
 }
 
@@ -417,10 +422,11 @@ void Server::start_io_threads(std::vector<std::thread> &threads) {
             }
 
             // Transfer ownership from registry to batch processing
+            size_t epoll_index = thread_id % epoll_fds_.size();
             batch[batch_count] = std::unique_ptr<Connection>(new Connection(
                 addr, fd,
                 connection_id_.fetch_add(1, std::memory_order_relaxed),
-                &handler_, weak_from_this()));
+                epoll_index, &handler_, weak_from_this()));
             batch_events[batch_count] =
                 EPOLLIN; // New connections always start with read
             batch_count++;