Skip to content

Commit afe4a7f

Browse files
fix: catch BaseException in connect() to clean up on CancelledError
The `except Exception:` handler in connect() did not catch CancelledError (a BaseException since Python 3.9). If a task was cancelled during handshake() or open_database(), the protocol's close() was never called, leaking the TCP socket. The connection also appeared as "connected" despite being half-initialized. Change to `except BaseException:` so cancellation properly closes the transport and resets the protocol to None. Closes #083 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent b87f3e0 commit afe4a7f

2 files changed

Lines changed: 40 additions & 1 deletion

File tree

src/dqliteclient/connection.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ async def connect(self) -> None:
101101
try:
102102
await self._protocol.handshake()
103103
self._db_id = await self._protocol.open_database(self._database)
104-
except Exception:
104+
except BaseException:
105105
self._protocol.close()
106106
self._protocol = None
107107
raise

tests/test_connection.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -493,6 +493,45 @@ async def cancelled_transaction():
493493
# _in_transaction must be cleaned up
494494
assert not conn._in_transaction
495495

496+
async def test_connect_cancellation_cleans_up_protocol(self) -> None:
497+
"""Cancelling connect() during handshake must close the transport."""
498+
import asyncio
499+
500+
conn = DqliteConnection("localhost:9001")
501+
502+
mock_reader = AsyncMock()
503+
mock_writer = MagicMock()
504+
mock_writer.drain = AsyncMock()
505+
mock_writer.close = MagicMock()
506+
mock_writer.wait_closed = AsyncMock()
507+
508+
handshake_entered = asyncio.Event()
509+
510+
async def slow_handshake(*args, **kwargs):
511+
handshake_entered.set()
512+
await asyncio.sleep(10) # Will be cancelled
513+
514+
with (
515+
patch("asyncio.open_connection", return_value=(mock_reader, mock_writer)),
516+
patch("dqliteclient.connection.DqliteProtocol") as MockProto,
517+
):
518+
proto_instance = MagicMock()
519+
proto_instance.handshake = AsyncMock(side_effect=slow_handshake)
520+
proto_instance.close = MagicMock()
521+
MockProto.return_value = proto_instance
522+
523+
task = asyncio.create_task(conn.connect())
524+
await handshake_entered.wait()
525+
task.cancel()
526+
527+
with pytest.raises(asyncio.CancelledError):
528+
await task
529+
530+
# The protocol must have been closed to avoid socket leak
531+
proto_instance.close.assert_called()
532+
# The connection must not appear as connected
533+
assert not conn.is_connected
534+
496535
async def test_not_leader_error_invalidates_connection(self) -> None:
497536
"""OperationalError with 'not leader' code should invalidate the connection."""
498537
conn = DqliteConnection("localhost:9001")

0 commit comments

Comments
 (0)