diff --git a/tools/load_tester.cpp b/tools/load_tester.cpp index 72025d0..a1e71ac 100644 --- a/tools/load_tester.cpp +++ b/tools/load_tester.cpp @@ -163,6 +163,8 @@ struct Connection { initRequest(); } + // Match server's connection state management + bool hasMessages() const { return !request.empty(); } bool error = false; ~Connection() { @@ -321,7 +323,7 @@ int main() { continue; } perror("epoll_wait"); - abort(); + abort(); // Keep abort for critical errors like server does } break; } @@ -333,9 +335,9 @@ int main() { events[i].data.ptr = nullptr; const int fd = conn->fd; - if (events[i].events & EPOLLERR) { - // Done with connection - continue; + // Handle connection errors like server + if (events[i].events & (EPOLLERR | EPOLLHUP | EPOLLRDHUP)) { + continue; // Let unique_ptr destructor clean up } if (events[i].events & EPOLLOUT) { bool finished = conn->writeBytes(); @@ -363,17 +365,19 @@ int main() { // Transfer back to epoll instance. This thread or another thread // will wake when fd is ready - if (!conn->request.empty() && !conn->error) { - events[i].events = EPOLLOUT | EPOLLONESHOT; + if (conn->hasMessages()) { + events[i].events = EPOLLOUT | EPOLLONESHOT | EPOLLRDHUP; } else { - events[i].events = EPOLLIN | EPOLLONESHOT; + events[i].events = EPOLLIN | EPOLLONESHOT | EPOLLRDHUP; } conn->tsan_release(); - events[i].data.ptr = conn.release(); + Connection *raw_conn = conn.release(); + events[i].data.ptr = raw_conn; int e = epoll_ctl(epollfd, EPOLL_CTL_MOD, fd, &events[i]); if (e == -1) { - perror("epoll_ctl"); - abort(); + perror("epoll_ctl MOD"); + delete raw_conn; // Clean up on failure like server + continue; } } } @@ -395,17 +399,25 @@ int main() { } // int fd = getConnectFd("127.0.0.1", "4569"); int fd = getConnectFdUnix("weaseldb.sock"); + if (fd == -1) { + continue; // Connection failed, try again + } + + // Create connection with proper ownership like server auto conn = std::make_unique( fd, connectionId.fetch_add(1, std::memory_order_relaxed)); - // Post to epoll instance + + // Add to epoll with proper events matching server pattern struct epoll_event event{}; - event.events = EPOLLOUT | EPOLLONESHOT; + event.events = EPOLLOUT | EPOLLONESHOT | EPOLLRDHUP; conn->tsan_release(); - event.data.ptr = conn.release(); + Connection *raw_conn = conn.release(); + event.data.ptr = raw_conn; e = epoll_ctl(epollfd, EPOLL_CTL_ADD, fd, &event); if (e == -1) { - perror("epoll_ctl"); - abort(); + perror("epoll_ctl ADD"); + delete raw_conn; // Clean up on failure like server + continue; } } });