From fddf4a231c80f943bb114be56163f188dbdc777f Mon Sep 17 00:00:00 2001 From: razimograbi Date: Sun, 26 Jul 2026 19:21:14 +0300 Subject: [PATCH 1/2] remove busy_socket and busy_mutex, replaced with EPOLLONESHOT --- src/Networking/EpollEventLoop.cpp | 11 ++++++++- src/Networking/EpollEventLoop.h | 1 + src/Networking/EventLoop.h | 3 +++ src/Networking/IocpEventLoop.cpp | 9 ++++++- src/Networking/IocpEventLoop.h | 1 + src/Server/Server.cpp | 39 ++++++++++++++++--------------- src/Server/Server.h | 12 ++++++++-- 7 files changed, 53 insertions(+), 23 deletions(-) diff --git a/src/Networking/EpollEventLoop.cpp b/src/Networking/EpollEventLoop.cpp index 35e7695..9fb8a0c 100644 --- a/src/Networking/EpollEventLoop.cpp +++ b/src/Networking/EpollEventLoop.cpp @@ -18,7 +18,7 @@ EpollEventLoop::~EpollEventLoop() void EpollEventLoop::add(SocketType sock) { epoll_event ev{}; - ev.events = EPOLLIN; // watch for "readable" (level-triggered by default) + ev.events = EPOLLIN | EPOLLONESHOT; // watch for "readable", (intentionally "mutes" the socket after the first notification) ev.data.fd = sock; epoll_ctl(epollFd, EPOLL_CTL_ADD, sock, &ev); @@ -67,4 +67,13 @@ int EpollEventLoop::wait(std::vector &out) return numReady; } +bool EpollEventLoop::rearm(SocketType sock) +{ + epoll_event ev{}; + ev.events = EPOLLIN | EPOLLONESHOT; + ev.data.fd = sock; + + return epoll_ctl(epollFd, EPOLL_CTL_MOD, sock, &ev) == 0; +} + #endif //_WIN32 \ No newline at end of file diff --git a/src/Networking/EpollEventLoop.h b/src/Networking/EpollEventLoop.h index 5bbc572..f43dc54 100644 --- a/src/Networking/EpollEventLoop.h +++ b/src/Networking/EpollEventLoop.h @@ -33,6 +33,7 @@ class EpollEventLoop : public IEventLoop void add(SocketType sock) override; void remove(SocketType sock) override; int wait(std::vector &out) override; + bool rearm(SocketType sock) override; private: /// File descriptor for the epoll instance itself. diff --git a/src/Networking/EventLoop.h b/src/Networking/EventLoop.h index 0f306e5..1c9524f 100644 --- a/src/Networking/EventLoop.h +++ b/src/Networking/EventLoop.h @@ -65,6 +65,9 @@ class IEventLoop * nothing ready, or -1 on error. */ virtual int wait(std::vector &out) = 0; + + // Explained inside the Server.h + virtual bool rearm(SocketType sock) = 0; }; #endif // KV_DATABASE_EVENTLOOP_H \ No newline at end of file diff --git a/src/Networking/IocpEventLoop.cpp b/src/Networking/IocpEventLoop.cpp index 42d6625..d4a0ff9 100644 --- a/src/Networking/IocpEventLoop.cpp +++ b/src/Networking/IocpEventLoop.cpp @@ -82,10 +82,17 @@ int IocpEventLoop::wait(std::vector &out) // recv() next; if THAT returns 0, that's how a graceful close is // detected - same pattern as epoll's EPOLLIN + recv()==0 on Linux. out.push_back(EventLoopEntry{sock, IOEvent::Readable}); - armRead(sock); // re-arm so we're notified again for the next batch of data + // We no longer auto-rearm here. The worker thread is now responsible + // for calling rearm() once it finishes processing fragmentation. } return 1; } +bool IocpEventLoop::rearm(SocketType sock) +{ + armRead(sock); + return true; // WSARecv failures handle themselves asynchronously in wait() +} + #endif //_WIN32 \ No newline at end of file diff --git a/src/Networking/IocpEventLoop.h b/src/Networking/IocpEventLoop.h index a5b2b76..517684f 100644 --- a/src/Networking/IocpEventLoop.h +++ b/src/Networking/IocpEventLoop.h @@ -49,6 +49,7 @@ class IocpEventLoop : public IEventLoop void add(SocketType sock) override; void remove(SocketType sock) override; int wait(std::vector &out) override; + bool rearm(SocketType sock) override; private: /// Posts (or re-posts) the zero-byte WSARecv that arms readiness notification for a socket. diff --git a/src/Server/Server.cpp b/src/Server/Server.cpp index 4bb311b..1fb70f4 100644 --- a/src/Server/Server.cpp +++ b/src/Server/Server.cpp @@ -163,27 +163,14 @@ void Server::runEventLoop() if (entry.event == IOEvent::Readable) { auto it = connections.find(entry.socket); - if (it == connections.end()) - continue; // already cleaned up - - { - std::lock_guard lock(busyMutex); - // A job for this socket is already queued or running - - // skip this notification. Level-triggered epoll will - // notify us again next wait() if data is still unread. - if (busySockets.count(entry.socket)) - continue; - busySockets.insert(entry.socket); - } + if (it == connections.end()) continue; Connection conn = it->second; // small struct, cheap to copy - tpool.acceptJob([this, conn]() - { - messageHandler(conn.socket, conn.sessionKey); - std::lock_guard lock(busyMutex); - busySockets.erase(conn.socket); }); - } - else // HangUp or Error + // wont fire again until we re-arm it. + tpool.acceptJob([this, conn] { + messageHandler(conn.socket, conn.sessionKey); + }); + } else // HangUp or Error { closeConnection(entry.socket); } @@ -228,6 +215,7 @@ void Server::messageHandler(SocketType clientSocket, const SessionKey &sessionKe // Nothing to read right now - a different job for this socket // already drained it, or epoll notified us before this job // got scheduled. Not an error, just nothing to do. + this->rearmSocket(clientSocket); return; } std::cout << "[CLIENT] recv error, disconnecting\n"; @@ -258,6 +246,7 @@ void Server::messageHandler(SocketType clientSocket, const SessionKey &sessionKe { std::string response = "Empty Value Recieved, Try again\n"; send(clientSocket, response.c_str(), response.length(), 0); + this->rearmSocket(clientSocket); return; } @@ -298,4 +287,16 @@ void Server::messageHandler(SocketType clientSocket, const SessionKey &sessionKe std::string response = "No command was received\n"; send(clientSocket, response.c_str(), response.length(), 0); } + + //Re-enable epoll notifications for the next command from this client + this->rearmSocket(clientSocket); +} + +void Server::rearmSocket(SocketType clientSocket) +{ + if (!eventLoop->rearm(clientSocket)) + { + std::cout << "[SERVER] Failed to re-arm socket " << clientSocket << ", closing.\n"; + closeConnection(clientSocket); + } } \ No newline at end of file diff --git a/src/Server/Server.h b/src/Server/Server.h index 5605944..1c88afc 100644 --- a/src/Server/Server.h +++ b/src/Server/Server.h @@ -40,8 +40,6 @@ class Server std::unordered_map connections; /** Client connections, keyed by socket */ std::unique_ptr eventLoop; /** Watches all client sockets for readiness */ std::thread eventLoopThread; /** Thread that runs runEventLoop() */ - std::mutex busyMutex; /** Guards busySockets */ - std::unordered_set busySockets; /** Sockets with a recv job already queued/running */ /** * @brief Runs continuously on eventLoopThread: waits for socket readiness @@ -91,6 +89,16 @@ class Server * @param sessionKey The session tied to this client. */ void messageHandler(SocketType clientSocket, const SessionKey &sessionKey); + + /** + * What it is: An explicit call (epoll_ctl with EPOLL_CTL_MOD) executed when a thread finishes its work. + * What it does: Unmutes the socket so epoll can start listening for network activity again. + * Why we use it: Because EPOLLONESHOT completely mutes the socket, it will stay dead forever unless re-armed. + * Every exit path in your thread—whether it finished a full message or is waiting for more bytes (fragmentation)—must call rearm(). + * + * @param clientSocket + */ + void rearmSocket(SocketType clientSocket); }; #endif \ No newline at end of file From 7d7d6dfa65a1ab75ac03650f58d9412b2ea5ffaa Mon Sep 17 00:00:00 2001 From: razimograbi Date: Sun, 26 Jul 2026 21:14:46 +0300 Subject: [PATCH 2/2] add_fragmentation_support --- src/Server/Server.cpp | 38 ++++++++++++++++++++++++++---------- src/Server/Server.h | 11 +++++++---- src/UserSession/Connection.h | 6 ++++++ 3 files changed, 41 insertions(+), 14 deletions(-) diff --git a/src/Server/Server.cpp b/src/Server/Server.cpp index 1fb70f4..e11bcc2 100644 --- a/src/Server/Server.cpp +++ b/src/Server/Server.cpp @@ -121,7 +121,7 @@ void Server::acceptClients() CloseSocket(AcceptSocket); continue; } - connections[AcceptSocket] = Connection{AcceptSocket, active_key}; + connections[AcceptSocket] = std::make_shared(AcceptSocket, active_key); eventLoop->add(AcceptSocket); std::cout << "[SERVER] Client connected and registered!\n"; } @@ -165,10 +165,10 @@ void Server::runEventLoop() auto it = connections.find(entry.socket); if (it == connections.end()) continue; - Connection conn = it->second; // small struct, cheap to copy + std::shared_ptr user_connection = it->second; // No Copy // wont fire again until we re-arm it. - tpool.acceptJob([this, conn] { - messageHandler(conn.socket, conn.sessionKey); + tpool.acceptJob([this, conn = std::move(user_connection)] { + messageHandler(conn); }); } else // HangUp or Error { @@ -185,21 +185,22 @@ void Server::closeConnection(SocketType sock) auto it = connections.find(sock); if (it != connections.end()) { - userSessionManager.remove_session(it->second.sessionKey); + userSessionManager.remove_session(it->second->sessionKey); connections.erase(it); } CloseSocket(sock); } -void Server::messageHandler(SocketType clientSocket, const SessionKey &sessionKey) +void Server::messageHandler(std::shared_ptr clientConnection) { std::cout << "[CLIENT] Handling client message\n"; - char buffer[1024]; + char tempBuffer[1024]; + SocketType clientSocket = clientConnection->socket; // On Linux, the buffer is safely passed to standard recv - int bytesReceived = recv(clientSocket, buffer, sizeof(buffer) - 1, 0); + int bytesReceived = recv(clientSocket, tempBuffer, sizeof(tempBuffer), 0); if (bytesReceived == 0) { @@ -222,13 +223,30 @@ void Server::messageHandler(SocketType clientSocket, const SessionKey &sessionKe closeConnection(clientSocket); return; } - userSessionManager.update_activity(sessionKey); + + clientConnection->commandBuffer.append(tempBuffer, bytesReceived); + + if (clientConnection->commandBuffer.length() > this->MAX_COMMAND_BUFFER_LENGTH) { + std::cout << "[CLIENT] Client is abusing the command buffer disconnecting them\n"; + closeConnection(clientSocket); + return; + } + + if (clientConnection->commandBuffer.find('\n') == std::string::npos) { + this->rearmSocket(clientSocket); + return; + } + + + + userSessionManager.update_activity(clientConnection->sessionKey); std::cout << "[CLIENT] Received " << bytesReceived << " bytes\n"; - std::string message(buffer, bytesReceived); + std::string message = clientConnection->commandBuffer; std::cout << "[CLIENT] Message: " << message << "\n"; std::istringstream iss(message); std::string command, key, value; + clientConnection->commandBuffer.clear(); iss >> command; iss >> key; diff --git a/src/Server/Server.h b/src/Server/Server.h index 1c88afc..8b29416 100644 --- a/src/Server/Server.h +++ b/src/Server/Server.h @@ -37,7 +37,7 @@ class Server RateLimiter rt; /** Server owns an instance of RateLimter class */ UserSessionManager userSessionManager; /** Managing User Sessions */ UserSessionBackgroundWorker user_session_background_worker; /** Background worker that sweeps expired sessions*/ - std::unordered_map connections; /** Client connections, keyed by socket */ + std::unordered_map> connections; /** Client connections, keyed by socket */ std::unique_ptr eventLoop; /** Watches all client sockets for readiness */ std::thread eventLoopThread; /** Thread that runs runEventLoop() */ @@ -85,10 +85,10 @@ class Server /** * @brief Handles one ready-to-read event for a client: one recv() call, * command parsing, and response. - * @param clientSocket The socket that has data available. - * @param sessionKey The session tied to this client. + * @param userConnection The user connection + * */ - void messageHandler(SocketType clientSocket, const SessionKey &sessionKey); + void messageHandler(std::shared_ptr userConnection); /** * What it is: An explicit call (epoll_ctl with EPOLL_CTL_MOD) executed when a thread finishes its work. @@ -99,6 +99,9 @@ class Server * @param clientSocket */ void rearmSocket(SocketType clientSocket); + + + uint16_t MAX_COMMAND_BUFFER_LENGTH = 2024; }; #endif \ No newline at end of file diff --git a/src/UserSession/Connection.h b/src/UserSession/Connection.h index bd0caaa..32daa03 100644 --- a/src/UserSession/Connection.h +++ b/src/UserSession/Connection.h @@ -22,6 +22,12 @@ struct Connection /// The session key identifying this client's UserSession entry. SessionKey sessionKey; + + /// This is the command of the user (Used for fragmentation management) + std::string commandBuffer; + + Connection(SocketType sock, SessionKey key) + : socket(sock), sessionKey(key) {} }; #endif // CONNECTION_H