Fix Fragmentaion Issue

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<EventLoopEntry> &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

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<EventLoopEntry> &out) override;
+    bool rearm(SocketType sock) override;
 
 private:
     /// File descriptor for the epoll instance itself.

src/Networking/EventLoop.h

@@ -65,6 +65,9 @@ class IEventLoop
      *         nothing ready, or -1 on error.
      */
     virtual int wait(std::vector<EventLoopEntry> &out) = 0;
+
+    // Explained inside the Server.h
+    virtual bool rearm(SocketType sock) = 0;
 };
 
 #endif // KV_DATABASE_EVENTLOOP_H

src/Networking/IocpEventLoop.cpp

@@ -82,10 +82,17 @@ int IocpEventLoop::wait(std::vector<EventLoopEntry> &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

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<EventLoopEntry> &out) override;
+    bool rearm(SocketType sock) override;
 
 private:
     /// Posts (or re-posts) the zero-byte WSARecv that arms readiness notification for a socket.

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<Connection>(AcceptSocket, active_key);
         eventLoop->add(AcceptSocket);
         std::cout << "[SERVER] Client connected and registered!\n";
     }
@@ -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<std::mutex> 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);
-                }
-
-                Connection conn = it->second; // small struct, cheap to copy
-                tpool.acceptJob([this, conn]()
-                                {
-                                    messageHandler(conn.socket, conn.sessionKey);
-                                    std::lock_guard<std::mutex> lock(busyMutex);
-                                    busySockets.erase(conn.socket); });
-            }
-            else // HangUp or Error
+                if (it == connections.end()) continue;
+
+                std::shared_ptr<Connection> user_connection = it->second; // No Copy
+                // wont fire again until we re-arm it.
+                tpool.acceptJob([this, conn = std::move(user_connection)] {
+                    messageHandler(conn);
+                });
+            } else // HangUp or Error
             {
                 closeConnection(entry.socket);
             }
@@ -198,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<Connection> 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)
     {
@@ -228,19 +216,37 @@ 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";
         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;
@@ -258,6 +264,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 +305,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);
+    }
 }

src/Server/Server.h

@@ -37,11 +37,9 @@ 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<SocketType, Connection> connections;     /** Client connections, keyed by socket */
+    std::unordered_map<SocketType, std::shared_ptr<Connection>> connections;     /** Client connections, keyed by socket */
     std::unique_ptr<IEventLoop> eventLoop;                      /** Watches all client sockets for readiness */
     std::thread eventLoopThread;                                /** Thread that runs runEventLoop() */
-    std::mutex busyMutex;                                       /** Guards busySockets */
-    std::unordered_set<SocketType> busySockets;                 /** Sockets with a recv job already queued/running */
 
     /**
      * @brief Runs continuously on eventLoopThread: waits for socket readiness
@@ -87,10 +85,23 @@ 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<Connection> userConnection);
+
+    /**
+    * 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);
+
+
+    uint16_t MAX_COMMAND_BUFFER_LENGTH = 2024;
 };
 
 #endif

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