feat: adds optional timeout for client connections (TCPServer) (#190)

* feat: adds optional timeout for client connections (TCPServer)

* feat: add support for handling client timeouts in TCP server

Author: M0r13n

CHANGELOG.txt

@@ -1,6 +1,10 @@
 ====================
  pyais CHANGELOG
 ====================
+-------------------------------------------------------------------------------
+ Version 2.13.1 21 Aug 2025
+-------------------------------------------------------------------------------
+* add support for handling client timeouts in TCP server
 -------------------------------------------------------------------------------
  Version 2.13.0 11 Aug 2025
 -------------------------------------------------------------------------------

pyais/__init__.py

@@ -5,7 +5,7 @@
 from pyais.tracker import AISTracker, AISTrack
 
 __license__ = 'MIT'
-__version__ = '2.13.0'
+__version__ = '2.13.1'
 __author__ = 'Leon Morten Richter'
 
 __all__ = (

pyais/stream.py

@@ -1,4 +1,5 @@
 import selectors
+import time
 import typing
 import pathlib
 import queue
@@ -422,14 +423,25 @@ class TCPServer(SocketStream):
     new connections, processing client data with proper message boundary detection,
     and cleaning up closed connections, making it suitable for applications requiring
     concurrent TCP communication with multiple clients.
+
+    Args:
+        host:           The hostname or IP address to bind to.
+        port:           The port number to bind to. Defaults to 80.
+        preprocessor:   Optional preprocessor for handling incoming data.
+                        If None, no preprocessing will be applied.
+        tbq:            Optional tag block queue for managing tagged data blocks.
+                        If None, no tag block queuing will be used.
+        timeout:        Connection timeout in seconds. Use -1 for no timeout.
+                        Defaults to -1.
     """
 
     def __init__(
         self,
         host: str,
         port: int = 80,
         preprocessor: typing.Optional[PreprocessorProtocol] = None,
-        tbq: typing.Optional[TagBlockQueue] = None
+        tbq: typing.Optional[TagBlockQueue] = None,
+        timeout: float = -1
     ) -> None:
         # Create a socket and bind it
         server_sock = socket(AF_INET, SOCK_STREAM)
@@ -447,6 +459,8 @@ def __init__(
 
         # Keep track of client sockets to clean them up later
         self._client_sockets: typing.Set[socket] = set()
+        self._client_last_activity: typing.Dict[socket, float] = {}
+        self.timeout = timeout
 
         super().__init__(server_sock, preprocessor=preprocessor, tbq=tbq)
 
@@ -456,13 +470,9 @@ def recv(self) -> bytes:
         if self._message_queue:
             return self._message_queue.popleft()
 
+        last_cleanup = time.time()
         while True:
-            events = self.sel.select()
-            if not events:
-                # No data available
-                return b''
-
-            for key, mask in events:
+            for key, mask in self.sel.select(timeout=1.0):
                 if key.data is None:
                     self.accept(key.fileobj)  # type: ignore
                 else:
@@ -471,6 +481,26 @@ def recv(self) -> bytes:
                     if self._message_queue:
                         return self._message_queue.popleft()
 
+            # Periodic cleanup for connection timeouts
+            if self.timeout > 0 and time.time() - last_cleanup >= self.timeout:
+                self._cleanup_idle_connections()
+                last_cleanup = time.time()
+
+    def _cleanup_idle_connections(self) -> None:
+        current_time = time.time()
+        for sock, last_activity in self._client_last_activity.copy().items():
+            if current_time - last_activity > self.timeout:
+                self._close_client(sock)
+
+    def _close_client(self, client_sock: socket) -> None:
+        try:
+            self.sel.unregister(client_sock)
+        except (KeyError, ValueError):
+            pass
+        self._client_last_activity.pop(client_sock, None)
+        self._client_sockets.discard(client_sock)
+        client_sock.close()
+
     def read(self) -> typing.Generator[bytes, None, None]:
         """Use custom implementation that handles multiple clients properly"""
         try:
@@ -482,34 +512,35 @@ def read(self) -> typing.Generator[bytes, None, None]:
             self.sel.close()
 
     def accept(self, sock: socket) -> None:
-        conn, addr = sock.accept()
-        conn.setblocking(False)
-        self._client_sockets.add(sock)
+        client_sock, addr = sock.accept()
+        client_sock.setblocking(False)
+        self._client_sockets.add(client_sock)
 
         data = ClientConnection(
             addr=addr,
             partial_buffer=b''
         )
-        self.sel.register(conn, selectors.EVENT_READ, data=data)
+        self.sel.register(client_sock, selectors.EVENT_READ, data=data)
+        self._client_last_activity[client_sock] = time.time()
 
     def service(self, key: selectors.SelectorKey, mask: int) -> None:
         """Handle client data with per-connection buffering"""
-        sock: socket = typing.cast(socket, key.fileobj)
+        client_sock: socket = typing.cast(socket, key.fileobj)
         data = key.data
 
         if mask & selectors.EVENT_READ:
             try:
-                recv_data = sock.recv(self.BUF_SIZE)
+                recv_data = client_sock.recv(self.BUF_SIZE)
                 if recv_data:
                     # Process the received data with the connection's buffer
                     self._process_client_data(data, recv_data)
                 else:
                     # close
-                    self.sel.unregister(sock)
-                    sock.close()
+                    self._close_client(client_sock)
             except ConnectionResetError:
-                self.sel.unregister(sock)
-                sock.close()
+                self._close_client(client_sock)
+
+        self._client_last_activity[client_sock] = time.time()
 
     def _process_client_data(self, client_data: ClientConnection, new_data: bytes) -> None:
         """Process data from a specific client, handling partial messages"""
@@ -536,12 +567,7 @@ def _process_client_data(self, client_data: ClientConnection, new_data: bytes) -
 
     def close(self) -> None:
         """Properly close all connections and selector"""
-        for sock in self._client_sockets.copy():
-            try:
-                self.sel.unregister(sock)
-            except (KeyError, ValueError):
-                pass
-            sock.close()
-            self._client_sockets.discard(sock)
+        for client_sock in self._client_sockets.copy():
+            self._close_client(client_sock)
         self.sel.close()
         super().close()

tests/test_tcp_server.py

@@ -13,8 +13,7 @@ class TestMessageProcessing(unittest.TestCase):
 
     def setUp(self):
         """Set up server for message processing tests"""
-        with patch('socket.socket'), \
-                patch('selectors.DefaultSelector'):
+        with patch('socket.socket'), patch('selectors.DefaultSelector'):
 
             self.server = TCPServer('127.0.0.1', 8080)
             self.server._message_queue = deque()