Pre-ping idle connections on acquire to detect peer-FIN zombies

The protocol-level is_connected flag is just self._protocol is not
None — it does not detect a half-closed TCP socket. An idle pool
connection that saw a clean peer FIN (graceful leader flip,
intentional close from the server side) passed the gate and got
handed to the caller, whose first query failed on the wire. The
pool already had _socket_looks_dead — a cheap non-blocking transport
peek (transport.is_closing() / reader.at_eof() / poisoned-decoder
short-circuit) — but only used it on the release path, not on
acquire.

Extend the acquire-side liveness gate to OR in _socket_looks_dead
and explicitly close the dead-but-protocol-still-set connection
before draining the rest of the idle queue. Make _socket_looks_dead
tolerate partial mocks via getattr on _protocol so tests that don't
construct a full protocol slice still work.

Update the existing _BrokenConn test fake to ship a healthy default
protocol mock so the pre-ping does not preempt the broken-cleanup
branch the test is actually exercising.

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

Author: antoineleclair

src/dqliteclient/pool.py

@@ -71,7 +71,11 @@ def _socket_looks_dead(conn: DqliteConnection) -> bool:
     reader. Mocked / missing attributes default to False (assume alive) so
     the check never produces a false positive against well-behaved peers.
     """
-    protocol = conn._protocol
+    # ``getattr`` (rather than direct attribute access) so the function
+    # tolerates partial mocks that omit ``_protocol`` entirely — the
+    # acquire-path pre-ping calls this on every dequeued conn, including
+    # ones built by tests that don't set up a protocol.
+    protocol = getattr(conn, "_protocol", None)
     if protocol is None:
         return True
     # Poisoned-decoder short-circuit: a wire desync means the next
@@ -745,12 +749,31 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
 
         # If connection is dead, discard and create a fresh one with leader discovery.
         # Also drain other idle connections — they likely point to the same dead server.
-        if not conn.is_connected:
+        #
+        # ``is_connected`` is the protocol-level handshake-complete flag;
+        # ``_socket_looks_dead`` is a non-blocking transport-level peek
+        # (transport.is_closing() / reader.at_eof() / poisoned-decoder
+        # short-circuit). An idle connection that has seen a clean peer
+        # FIN (leader flip with graceful close) passes ``is_connected``
+        # but trips ``_socket_looks_dead`` — without the second check,
+        # the pool hands out a zombie connection and the user's first
+        # query fails. The peek is one syscall; cheap to run on every
+        # acquire.
+        if not conn.is_connected or _socket_looks_dead(conn):
             logger.debug(
                 "pool.acquire: drain-idle triggered by stale conn=%r closing_idle=%d",
                 conn,
                 self._pool.qsize(),
             )
+            # Close the dead-but-transport-alive connection explicitly:
+            # ``_drain_idle`` only walks the idle queue, so a connection
+            # we just dequeued whose transport is half-closed (FIN seen)
+            # would otherwise leak its writer. ``close()`` is safe on a
+            # never-connected conn (``_protocol is None``) — it short-
+            # circuits to a noop. Shielded so an outer cancel does not
+            # leave the writer dangling.
+            with contextlib.suppress(*_POOL_CLEANUP_EXCEPTIONS):
+                await asyncio.shield(conn.close())
             # Wrap _drain_idle so a cancel mid-drain releases the dead
             # conn's reservation. Without this guard, a CancelledError
             # delivered to a checkpoint inside _drain_idle propagates

tests/test_pool_acquire_cleanup.py

@@ -53,6 +53,21 @@ def __init__(self, close_side_effect: Any = None) -> None:
         self._address = "stub:9001"
         self._close_side_effect = close_side_effect
         self.close_calls = 0
+        # Pretend healthy at acquire entry so the pre-ping does not
+        # short-circuit before the test's user-body has a chance to
+        # flip ``is_connected = False`` and exercise the broken-
+        # cleanup branch. Mirrors the slice of ``DqliteConnection``
+        # that ``_socket_looks_dead`` walks (protocol → writer/reader →
+        # transport.is_closing / reader.at_eof).
+        protocol = MagicMock()
+        protocol.is_wire_coherent = True
+        transport = MagicMock()
+        transport.is_closing.return_value = False
+        protocol._writer = MagicMock()
+        protocol._writer.transport = transport
+        protocol._reader = MagicMock()
+        protocol._reader.at_eof.return_value = False
+        self._protocol = protocol
 
     async def close(self) -> None:
         self.close_calls += 1

tests/test_pool_pre_ping_on_acquire.py

@@ -0,0 +1,132 @@
+"""Pin: ``ConnectionPool.acquire`` peeks the transport via
+``_socket_looks_dead`` *before* yielding an idle connection to the
+caller, so a connection that saw a clean peer FIN (leader flip with
+graceful close) does not get handed out as a zombie.
+
+The protocol-level ``is_connected`` is just ``self._protocol is not
+None`` — it does NOT detect a half-closed socket. Without the
+transport-level peek, the user's first query after a leader flip
+fails on a dequeue path that the pool could have self-healed.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from dqliteclient.pool import ConnectionPool
+
+
+def _alive_conn(*, dead: bool = False) -> MagicMock:
+    """Build a MagicMock that mimics the slice of ``DqliteConnection``
+    the pool's acquire flow touches.
+
+    With ``dead=True``, the transport peek tells ``_socket_looks_dead``
+    the connection is gone (transport.is_closing() returns True).
+    """
+    conn = MagicMock()
+    conn.is_connected = True
+    conn.close = AsyncMock()
+    conn._pool_released = False
+    proto = MagicMock()
+    proto.is_wire_coherent = True
+    transport = MagicMock()
+    transport.is_closing.return_value = dead
+    proto._writer = MagicMock()
+    proto._writer.transport = transport
+    proto._reader = MagicMock()
+    proto._reader.at_eof.return_value = False
+    conn._protocol = proto
+    return conn
+
+
+@pytest.mark.asyncio
+async def test_acquire_drops_dead_idle_connection_and_returns_fresh_one(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """An idle connection whose transport is closing must not be
+    yielded to the caller. The pool drains the idle queue and creates
+    a fresh connection."""
+    pool = ConnectionPool(["a:9001"], min_size=0, max_size=1)
+    # Skip real cluster bootstrap; we only exercise the acquire path.
+    monkeypatch.setattr("dqliteclient.pool.ConnectionPool.initialize", AsyncMock())
+    pool._initialized = True
+    # Seed one dead conn into the idle queue and reserve its slot.
+    dead = _alive_conn(dead=True)
+    pool._pool.put_nowait(dead)
+    pool._size = 1
+
+    fresh = _alive_conn(dead=False)
+    create_calls = 0
+
+    async def fake_create(self: object) -> MagicMock:
+        nonlocal create_calls
+        create_calls += 1
+        return fresh
+
+    monkeypatch.setattr("dqliteclient.pool.ConnectionPool._create_connection", fake_create)
+
+    async with pool.acquire() as conn:
+        assert conn is fresh, "acquire must drop the dead idle connection and yield a fresh one"
+        assert dead.close.called, "dead connection must have been closed"
+    assert create_calls == 1
+
+
+@pytest.mark.asyncio
+async def test_acquire_yields_alive_idle_connection_unchanged(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Regression pin: a healthy idle connection passes the peek and
+    is yielded unchanged — the pre-ping does NOT force unnecessary
+    reconnects."""
+    pool = ConnectionPool(["a:9001"], min_size=0, max_size=1)
+    monkeypatch.setattr("dqliteclient.pool.ConnectionPool.initialize", AsyncMock())
+    pool._initialized = True
+    alive = _alive_conn(dead=False)
+    pool._pool.put_nowait(alive)
+    pool._size = 1
+
+    create_calls = 0
+
+    async def fake_create(self: object) -> MagicMock:
+        nonlocal create_calls
+        create_calls += 1
+        return _alive_conn()
+
+    monkeypatch.setattr("dqliteclient.pool.ConnectionPool._create_connection", fake_create)
+
+    async with pool.acquire() as conn:
+        assert conn is alive, "alive idle connection must pass through unchanged"
+        assert not alive.close.called
+    assert create_calls == 0, "no fresh create when the idle connection passes the peek"
+
+
+@pytest.mark.asyncio
+async def test_acquire_peeks_via_eof(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """``_socket_looks_dead`` triggers on EOF too, not just
+    transport.is_closing()."""
+    pool = ConnectionPool(["a:9001"], min_size=0, max_size=1)
+    monkeypatch.setattr("dqliteclient.pool.ConnectionPool.initialize", AsyncMock())
+    pool._initialized = True
+    eof_conn = _alive_conn(dead=False)
+    eof_conn._protocol._reader.at_eof.return_value = True
+    pool._pool.put_nowait(eof_conn)
+    pool._size = 1
+
+    fresh = _alive_conn()
+
+    async def fake_create(self: object) -> MagicMock:
+        return fresh
+
+    monkeypatch.setattr("dqliteclient.pool.ConnectionPool._create_connection", fake_create)
+
+    async with pool.acquire() as conn:
+        assert conn is fresh
+        assert eof_conn.close.called
+
+    # Avoid leaking the asyncio queue close-task on test teardown
+    await asyncio.sleep(0)