diff --git a/tools/load_tester.cpp b/tools/load_tester.cpp index bffcc08..785d184 100644 --- a/tools/load_tester.cpp +++ b/tools/load_tester.cpp @@ -723,10 +723,30 @@ int main(int argc, char *argv[]) { auto conn = std::make_unique( fd, connectionId.fetch_add(1, std::memory_order_relaxed)); - // Add to epoll with proper events matching server pattern - struct epoll_event event{}; + // Try to write once in the connect thread before handing off to network + // threads assert(conn->hasMessages()); - event.events = EPOLLOUT | EPOLLONESHOT | EPOLLRDHUP; + bool writeFinished = conn->writeBytes(); + if (conn->error) { + continue; // Connection failed, destructor will clean up + } + + // Determine the appropriate epoll events based on write result + struct epoll_event event{}; + if (writeFinished) { + // All data was written, wait for response + int shutdown_result = shutdown(conn->fd, SHUT_WR); + if (shutdown_result == -1) { + perror("shutdown"); + continue; + } + event.events = EPOLLIN | EPOLLONESHOT | EPOLLRDHUP; + } else { + // Partial write, need to continue writing + event.events = EPOLLOUT | EPOLLONESHOT | EPOLLRDHUP; + } + + // Add to epoll for network threads to handle remaining I/O conn->tsan_release(); Connection *raw_conn = conn.release(); event.data.ptr = raw_conn;