Merge pull request #31 from Ataba29/fragmentation_integration_testing

Fragmentation integration testing

Author: Ataba29

.github/workflows/ci.yml

@@ -42,3 +42,42 @@ jobs:
 
       - name: Test
         run: ctest --test-dir build --output-on-failure
+
+  integration-test:
+    name: Integration Tests (Docker)
+    runs-on: ubuntu-24.04
+    needs: build-linux
+    steps:
+      - uses: actions/checkout@v7
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install test dependencies
+        run: pip install -r src/Tests/integration/requirements.txt
+
+      - name: Build Docker image
+        run: docker build -t byteforge:ci .
+
+      - name: Run ByteForge container
+        run: docker run -d --name byteforge -p 6625:6625 byteforge:ci
+
+      - name: Wait for server to be ready
+        run: |
+          for i in {1..30}; do
+            (echo > /dev/tcp/127.0.0.1/6625) 2>/dev/null && break
+            sleep 1
+          done
+
+      - name: Run integration tests
+        run: pytest src/Tests/integration -v
+
+      - name: Dump container logs on failure
+        if: failure()
+        run: docker logs byteforge
+
+      - name: Stop container
+        if: always()
+        run: docker rm -f byteforge

.gitignore

@@ -8,6 +8,8 @@ CMakeFiles/
 CMakeScripts/
 cmake_install.cmake
 Makefile
+.venv/
+__pycache__/
 
 # ==========================================
 # Compiled Executables and Libraries and Logs

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

@@ -125,7 +125,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";
     }
@@ -182,27 +182,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);
             }
@@ -217,21 +204,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)
     {
@@ -247,19 +235,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;
@@ -277,6 +283,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;
         }
 
@@ -317,6 +324,18 @@ 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);
+    }
 }
 
 void Server::onSessionExpired(SocketType sock)

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 */
     std::mutex expiredMutex;                                    /** Mutex to guard the expiredSockets vector */
     std::vector<SocketType> expiredSockets;                     /** Notifys the event loop of connections to remove */
 
@@ -89,16 +87,29 @@ 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);
 
     /**
      * @brief Adds sockets that are expired into the expired vector to be removed later by eventloop
      * @param sock which is the client socket that is to be removed
      */
     void onSessionExpired(SocketType sock);
+
+    /**
+    * 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/Tests/integration/client.py

@@ -0,0 +1,66 @@
+import socket
+import time
+
+
+class ByteForgeClient:
+    def __init__(self, host="127.0.0.1", port=6625, timeout=5.0,
+                 max_retries=5, retry_delay=0.25):
+        self.host = host
+        self.port = port
+        self.timeout = timeout
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+        self.sock = self._connect_with_retry()
+
+    def _connect_with_retry(self):
+        last_error = None
+        for _ in range(self.max_retries):
+            try:
+                return socket.create_connection((self.host, self.port), timeout=self.timeout)
+            except (ConnectionRefusedError, ConnectionResetError,
+                    ConnectionAbortedError, OSError) as e:
+                last_error = e
+                time.sleep(self.retry_delay)
+        raise ConnectionError(
+            f"Failed to connect to {self.host}:{self.port} after {self.max_retries} attempts"
+        ) from last_error
+
+    def _send(self, command: str) -> str:
+        last_error = None
+        for _ in range(self.max_retries):
+            try:
+                self.sock.sendall((command + "\n").encode())
+                return self._recv_line()
+            except (ConnectionResetError, ConnectionAbortedError,
+                    BrokenPipeError, OSError) as e:
+                last_error = e
+                try:
+                    self.sock.close()
+                except OSError:
+                    pass
+                time.sleep(self.retry_delay)
+                self.sock = self._connect_with_retry()
+        raise ConnectionError(
+            f"Failed to send command after {self.max_retries} attempts"
+        ) from last_error
+
+    def _recv_line(self) -> str:
+        data = b""
+        while not data.endswith(b"\n"):
+            chunk = self.sock.recv(1024)
+            if not chunk:
+                break
+            data += chunk
+        return data.decode().strip()
+
+    def insert(self, key: str, value: str) -> str:
+        return self._send(f"INSERT {key} {value}")
+
+    def get(self, key: str) -> str:
+        return self._send(f"GET {key}")
+
+    def delete(self, key: str) -> str:
+        return self._send(f"DELETE {key}")
+
+    def close(self):
+        self.sock.close()