Harden transaction + pool cancellation discipline

Coordinated fix for six cross-cutting cancellation/lifecycle bugs in
DqliteConnection.transaction() and ConnectionPool. Seven related
issues land together because they share the same code blocks and
invariants; splitting them would produce repeated rewrites of the
same methods across multiple PRs.

transaction() rewrites (connection.py):

- When body raises AND ROLLBACK also raises, invalidate the
  connection (set _protocol=None) so the pool discards it instead
  of returning to the queue with a live server-side transaction.
  The previous `finally` block cleared _in_transaction
  unconditionally, masking the dangling transaction from the pool's
  _reset_connection check.

- The body-raised ROLLBACK no longer uses suppress(BaseException).
  CancelledError / KeyboardInterrupt / SystemExit propagate
  (structured-concurrency contract — TaskGroup / asyncio.timeout()
  rely on it). Non-cancellation exceptions still invalidate the
  connection so the original body exception propagates cleanly.

- Cancellation matrix — cancel before BEGIN, in body, during COMMIT,
  during ROLLBACK — all four phases now produce consistent
  post-conditions: CancelledError propagates, connection state is
  terminal (either usable or invalidated).

ConnectionPool rewrites (pool.py):

- initialize() now uses asyncio.gather(return_exceptions=True) and
  closes survivors before re-raising the first failure. The
  previous default-exceptions path cancelled sibling tasks but did
  NOT close already-succeeded connections, leaking transports.

- The check-`_closed`-then-clear-`closed_event` sequence in
  acquire() now runs under self._lock, eliminating the window where
  a concurrent close() could set() the event between check and
  clear (which would lose the close signal and stall parked waiters
  until the per-poll timeout).

- acquire()'s except branch now wraps cleanup in a try/finally with
  asyncio.shield(self._release_reservation()). The finally runs
  even when a second cancellation lands mid-cleanup; shield keeps
  the _size decrement uninterruptible so the pool never loses a
  reservation under rapid-fire cancellation.

- CRITICAL transport leak fix (flagged by concurrency review on the
  initial diff): every `_pool_released = True` now comes AFTER
  `await conn.close()`. DqliteConnection.close() has an early-return
  guard on _pool_released; setting the flag first made every
  pool-side close() a silent no-op, leaking the transport. Applies
  to all 8 cleanup branches across acquire() and _release().

connect() failure path:

- New _abort_protocol() helper uses
  asyncio.wait_for(protocol.wait_closed(), timeout=0.5) so the
  transport drains under normal conditions but never hangs on an
  unresponsive peer. Different call site than the leader-query
  finally block (where unbounded wait_closed was correctly rejected
  for hang reasons); connect() is a retry-loop path that benefits
  from bounded draining.

Tests (all RED before this commit, GREEN after):

- tests/test_transaction_cancellation.py (7 tests):
  - TestRollbackFailureInvalidatesConnection
  - TestRollbackHappyPath: regression for body-raise + clean-rollback
  - TestRollbackCancellationPropagates
  - TestTransactionCancellationPhases: 4-phase matrix

- tests/test_pool_cancellation.py (4 tests):
  - TestInitializePartialFailureClosesSurvivors
  - TestAcquireCancellationRestoresSize
  - TestCloseWakesAllWaiters: message-text assertion (CI-latency
    robust) rather than timing

The _FakeConn test double replicates the real DqliteConnection's
_pool_released close-guard via close_effective. Without this, the
transport-leak bug (the critical finding on the initial diff) would
have been invisible to tests.

Two concurrency reviews and one test-quality review ran on the
implementation. The first review caught the _pool_released ordering
bug; the fix applies to every cleanup branch. A follow-up review
verified all 8 branches close-before-flag.

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

Author: antoineleclair

src/dqliteclient/connection.py

@@ -176,8 +176,7 @@ async def connect(self) -> None:
                 await self._protocol.handshake()
                 self._db_id = await self._protocol.open_database(self._database)
             except OperationalError as e:
-                self._protocol.close()
-                self._protocol = None
+                await self._abort_protocol()
                 if e.code in _LEADER_ERROR_CODES:
                     # Leader-change errors during OPEN are transport-level
                     # problems — the caller needs to reconnect elsewhere, not
@@ -187,8 +186,7 @@ async def connect(self) -> None:
                     ) from e
                 raise
             except BaseException:
-                self._protocol.close()
-                self._protocol = None
+                await self._abort_protocol()
                 raise
         finally:
             self._in_use = False
@@ -210,6 +208,25 @@ async def close(self) -> None:
         protocol.close()
         await protocol.wait_closed()
 
+    async def _abort_protocol(self) -> None:
+        """Tear down a half-open protocol during a connect failure path.
+
+        Close the writer, then give ``wait_closed`` a bounded budget so
+        the transport drains under normal conditions but never hangs on
+        an unresponsive peer (ISSUE-72). Cycle-2 ISSUE-38 rejected an
+        unbounded ``wait_closed`` in the leader-query finally block
+        because leader-query is a discovery path that runs against
+        arbitrary nodes; connect is a retry-loop path that benefits
+        from draining, so the two sites take different decisions.
+        """
+        protocol = self._protocol
+        if protocol is None:
+            return
+        self._protocol = None
+        protocol.close()
+        with contextlib.suppress(BaseException):
+            await asyncio.wait_for(protocol.wait_closed(), timeout=0.5)
+
     async def __aenter__(self) -> "DqliteConnection":
         await self.connect()
         return self
@@ -383,7 +400,26 @@ async def fetchval(self, sql: str, params: Sequence[Any] | None = None) -> Any:
 
     @asynccontextmanager
     async def transaction(self) -> AsyncIterator[None]:
-        """Context manager for transactions."""
+        """Context manager for transactions.
+
+        Cancellation contract (ISSUE-75 / ISSUE-79):
+        - Cancellation during BEGIN: state cleared, CancelledError
+          propagates.
+        - Cancellation during the body: ROLLBACK is attempted. If
+          ROLLBACK itself is cancelled, the connection is invalidated
+          and CancelledError propagates (structured-concurrency
+          contract — TaskGroup / asyncio.timeout() require this).
+        - Cancellation during COMMIT: connection invalidated
+          (server-side state ambiguous), CancelledError propagates.
+        - Cancellation during ROLLBACK (body already raised): connection
+          invalidated, CancelledError propagates and supersedes the
+          body exception (Python chains it via ``__context__``).
+
+        Non-cancellation ROLLBACK failure (ISSUE-73): connection is
+        invalidated so the pool discards it instead of reusing a
+        Python-side "_in_transaction=False" connection with live
+        server-side transaction state.
+        """
         if self._in_transaction:
             raise InterfaceError("Nested transactions are not supported; use SAVEPOINT directly")
 
@@ -409,10 +445,35 @@ async def transaction(self) -> AsyncIterator[None]:
                 # pool discards it instead of recycling an unknown state.
                 self._invalidate()
             else:
-                # Body raised before COMMIT; try to roll back. Swallow
-                # rollback errors; the original exception is more important.
-                with contextlib.suppress(BaseException):
+                # Body raised before COMMIT; try to roll back.
+                #
+                # Narrow suppression to Exception (NOT BaseException):
+                # CancelledError / KeyboardInterrupt / SystemExit must
+                # propagate. Previously ``suppress(BaseException)``
+                # swallowed cancellation, breaking structured-concurrency
+                # contracts (ISSUE-75).
+                #
+                # If ROLLBACK fails for any reason (including the narrow
+                # cancellation catch below), the connection's transaction
+                # state is unknowable from our side and the connection
+                # must be invalidated so the pool discards it on return
+                # (ISSUE-73). The original body exception is still the
+                # one that propagates, except for cancellation which
+                # takes precedence.
+                try:
                     await self.execute("ROLLBACK")
+                except (asyncio.CancelledError, KeyboardInterrupt, SystemExit):
+                    # Rollback interrupted mid-flight. Server-side tx is
+                    # in an unknown state; invalidate and propagate the
+                    # higher-priority signal.
+                    self._invalidate()
+                    raise
+                except Exception:
+                    # Rollback failed for a non-cancellation reason.
+                    # Invalidate so the pool discards on return, then
+                    # re-raise the ORIGINAL body exception (below)
+                    # — rollback failure is a secondary concern.
+                    self._invalidate()
             raise
         finally:
             self._tx_owner = None

src/dqliteclient/pool.py

@@ -154,6 +154,13 @@ async def initialize(self) -> None:
 
         Idempotent: concurrent callers share the same initialization —
         only one performs the TCP work, the others await its result.
+
+        Partial-failure behavior (ISSUE-77): ``asyncio.gather`` with the
+        default ``return_exceptions=False`` cancels sibling tasks on
+        first failure but does NOT close connections that already
+        succeeded — they leak as orphaned transports. Use
+        ``return_exceptions=True`` so every task resolves, then close
+        survivors explicitly before re-raising the first failure.
         """
         # Hold the lock across the gather so a second concurrent
         # initialize() call observes _initialized=True after the first
@@ -163,17 +170,31 @@ async def initialize(self) -> None:
                 return
             if self._min_size > 0:
                 self._size += self._min_size
-                try:
-                    # Create min_size connections concurrently so startup
-                    # latency doesn't scale with min_size × per-connect RTT.
-                    conns = await asyncio.gather(
-                        *(self._create_connection() for _ in range(self._min_size))
-                    )
-                except BaseException:
+                # Create min_size connections concurrently so startup
+                # latency doesn't scale with min_size × per-connect RTT.
+                results = await asyncio.gather(
+                    *(self._create_connection() for _ in range(self._min_size)),
+                    return_exceptions=True,
+                )
+                successes: list[DqliteConnection] = []
+                failures: list[BaseException] = []
+                for r in results:
+                    if isinstance(r, BaseException):
+                        failures.append(r)
+                    else:
+                        successes.append(r)
+                if failures:
+                    # Close the connections that did succeed — they are
+                    # unowned now that initialize is aborting.
+                    for conn in successes:
+                        with contextlib.suppress(BaseException):
+                            await conn.close()
                     self._size -= self._min_size
                     self._signal_state_change()
-                    raise
-                for conn in conns:
+                    # Re-raise the first observed failure as the root
+                    # cause; additional failures chain as context.
+                    raise failures[0]
+                for conn in successes:
                     await self._pool.put(conn)
             self._initialized = True
 
@@ -292,10 +313,13 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
                 )
             # Race the queue against the state-change event so any pool
             # state change (close, size decrement, drain) wakes waiters
-            # promptly. Clear right before parking so we only wake on
-            # subsequent signals; the loop top re-checks _closed.
-            closed_event = self._get_closed_event()
-            if not self._closed:
+            # promptly. The check-_closed-then-clear pair runs under
+            # the lock so a concurrent close() can't set() the event
+            # between our read and our clear (ISSUE-74).
+            async with self._lock:
+                closed_event = self._get_closed_event()
+                if self._closed:
+                    raise DqliteConnectionError("Pool is closed")
                 closed_event.clear()
             get_task: asyncio.Task[DqliteConnection] = asyncio.create_task(self._pool.get())
             closed_task = asyncio.create_task(closed_event.wait())
@@ -337,29 +361,51 @@ async def acquire(self) -> AsyncIterator[DqliteConnection]:
         try:
             yield conn
         except BaseException:
-            if conn.is_connected and not self._closed:
-                # Connection is healthy — user code raised a non-connection error.
-                # Roll back any open transaction, then return to pool.
-                if not await self._reset_connection(conn):
-                    conn._pool_released = True
-                    with contextlib.suppress(Exception):
+            # Cleanup must complete even if a second cancellation lands
+            # mid-await (ISSUE-76). ``returned_to_queue`` tracks whether
+            # the reservation was transferred to an in-queue connection;
+            # if not, the ``finally`` below releases it. ``asyncio.shield``
+            # around ``_release_reservation`` makes the decrement itself
+            # uninterruptible so ``_size`` stays consistent under rapid-
+            # fire cancellation.
+            #
+            # Order note: ``conn.close()`` has a guard that returns early
+            # when ``_pool_released`` is True (so user code can't
+            # accidentally close a connection they returned to the
+            # pool). The pool itself MUST therefore close BEFORE setting
+            # the flag, or the transport leaks.
+            returned_to_queue = False
+            try:
+                if conn.is_connected and not self._closed:
+                    # Connection is healthy — user code raised a non-connection
+                    # error. Roll back any open transaction, then return to pool.
+                    if await self._reset_connection(conn):
+                        try:
+                            self._pool.put_nowait(conn)
+                        except asyncio.QueueFull:
+                            with contextlib.suppress(Exception):
+                                await conn.close()
+                            conn._pool_released = True
+                        else:
+                            conn._pool_released = True
+                            returned_to_queue = True
+                    else:
+                        with contextlib.suppress(Exception):
+                            await conn.close()
+                        conn._pool_released = True
+                else:
+                    # Connection is broken (invalidated by execute/fetch error
+                    # handlers). Drain other idle connections — they likely
+                    # point to the same dead server.
+                    with contextlib.suppress(BaseException):
+                        await self._drain_idle()
+                    with contextlib.suppress(BaseException):
                         await conn.close()
-                    await self._release_reservation()
-                    raise
-                conn._pool_released = True
-                try:
-                    self._pool.put_nowait(conn)
-                except asyncio.QueueFull:
-                    await conn.close()
-                    await self._release_reservation()
-            else:
-                # Connection is broken (invalidated by execute/fetch error handlers).
-                # Drain other idle connections — they likely point to the same dead server.
-                conn._pool_released = True
-                await self._drain_idle()
-                with contextlib.suppress(BaseException):
-                    await conn.close()
-                await self._release_reservation()
+                    conn._pool_released = True
+            finally:
+                if not returned_to_queue:
+                    with contextlib.suppress(BaseException):
+                        await asyncio.shield(self._release_reservation())
             raise
         else:
             await self._release(conn)
@@ -402,26 +448,38 @@ async def _reset_connection(self, conn: DqliteConnection) -> bool:
         return True
 
     async def _release(self, conn: DqliteConnection) -> None:
-        """Return a connection to the pool or close it."""
+        """Return a connection to the pool or close it.
+
+        ``conn.close()`` has an early-return guard against
+        ``_pool_released=True``, so close MUST run before the flag is
+        set — otherwise the transport leaks (ISSUE-76 review found
+        this bug across every branch that closes a pool-owned
+        connection).
+        """
         if self._closed:
-            conn._pool_released = True
             await conn.close()
+            conn._pool_released = True
             await self._release_reservation()
             return
 
         if not await self._reset_connection(conn):
-            conn._pool_released = True
             with contextlib.suppress(Exception):
                 await conn.close()
+            conn._pool_released = True
             await self._release_reservation()
             return
 
-        conn._pool_released = True
+        # Healthy connection returning to queue: no close; just flip
+        # the flag and enqueue. If the queue is full we must close
+        # before setting the flag.
         try:
             self._pool.put_nowait(conn)
         except asyncio.QueueFull:
             await conn.close()
+            conn._pool_released = True
             await self._release_reservation()
+        else:
+            conn._pool_released = True
 
     async def execute(self, sql: str, params: Sequence[Any] | None = None) -> tuple[int, int]:
         """Execute a SQL statement using a pooled connection."""