fix: add read timeouts to protocol to prevent indefinite hangs

Protocol reads now use asyncio.wait_for with a configurable timeout
(default 15s). After handshake, the timeout is increased to at least
the server's heartbeat timeout. Previously, a hung server with an
open TCP connection would block the client forever.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: antoineleclair

src/dqliteclient/cluster.py

@@ -71,7 +71,7 @@ async def _query_leader(self, address: str) -> str | None:
         except (TimeoutError, OSError):
             return None
 
-        protocol = DqliteProtocol(reader, writer)
+        protocol = DqliteProtocol(reader, writer, timeout=self._timeout)
 
         try:
             await protocol.handshake()

src/dqliteclient/connection.py

@@ -61,7 +61,7 @@ async def connect(self) -> None:
         except OSError as e:
             raise DqliteConnectionError(f"Failed to connect to {self._address}: {e}") from e
 
-        self._protocol = DqliteProtocol(reader, writer)
+        self._protocol = DqliteProtocol(reader, writer, timeout=self._timeout)
 
         try:
             await self._protocol.handshake()

src/dqliteclient/protocol.py

@@ -32,6 +32,7 @@ def __init__(
         self,
         reader: asyncio.StreamReader,
         writer: asyncio.StreamWriter,
+        timeout: float = 15.0,
     ) -> None:
         self._reader = reader
         self._writer = writer
@@ -40,6 +41,7 @@ def __init__(
         self._buffer = ReadBuffer()
         self._client_id = 0
         self._heartbeat_timeout = 0
+        self._timeout = timeout
 
     async def handshake(self, client_id: int = 0) -> int:
         """Perform protocol handshake.
@@ -66,6 +68,10 @@ async def handshake(self, client_id: int = 0) -> int:
 
         self._client_id = client_id
         self._heartbeat_timeout = response.heartbeat_timeout
+        # Use heartbeat timeout for subsequent reads if larger than default
+        if response.heartbeat_timeout > 0:
+            heartbeat_seconds = response.heartbeat_timeout / 1000.0
+            self._timeout = max(self._timeout, heartbeat_seconds)
         return response.heartbeat_timeout
 
     async def get_leader(self) -> tuple[int, str]:
@@ -191,23 +197,33 @@ async def query_sql(
 
         return column_names, all_rows
 
+    async def _read_data(self) -> bytes:
+        """Read data from the stream with timeout."""
+        try:
+            data = await asyncio.wait_for(
+                self._reader.read(4096), timeout=self._timeout
+            )
+        except TimeoutError:
+            raise DqliteConnectionError(
+                f"Server read timed out after {self._timeout}s"
+            ) from None
+        if not data:
+            raise DqliteConnectionError("Connection closed by server")
+        return data
+
     async def _read_continuation(self) -> RowsResponse:
         """Read and decode a ROWS continuation frame."""
         while True:
             result = self._decoder.decode_continuation()
             if result is not None:
                 return result
-            data = await self._reader.read(4096)
-            if not data:
-                raise DqliteConnectionError("Connection closed by server")
+            data = await self._read_data()
             self._decoder.feed(data)
 
     async def _read_response(self) -> Message:
         """Read and decode the next response message."""
         while not self._decoder.has_message():
-            data = await self._reader.read(4096)
-            if not data:
-                raise DqliteConnectionError("Connection closed by server")
+            data = await self._read_data()
             self._decoder.feed(data)
 
         message = self._decoder.decode()

tests/test_protocol.py

@@ -138,6 +138,28 @@ async def test_query_sql_multipart(
         assert rows[0] == [1, "alice"]
         assert rows[1] == [2, "bob"]
 
+    async def test_read_timeout(
+        self,
+        mock_reader: AsyncMock,
+        mock_writer: MagicMock,
+    ) -> None:
+        """Protocol reads should time out instead of blocking forever."""
+        import asyncio
+
+        protocol = DqliteProtocol(mock_reader, mock_writer, timeout=0.1)
+
+        # Simulate a server that hangs (never returns data)
+        async def hang_forever(*args, **kwargs):
+            await asyncio.sleep(100)
+            return b""
+
+        mock_reader.read.side_effect = hang_forever
+
+        from dqliteclient.exceptions import DqliteConnectionError
+
+        with pytest.raises(DqliteConnectionError, match="timed out"):
+            await protocol.exec_sql(1, "SELECT 1")
+
     async def test_close(
         self,
         protocol: DqliteProtocol,