Do not clear _in_use from _invalidate

The flag's contract is "owned by the task that claimed it between
claim-site and clear-site". Clearing it from an out-of-band path
(call_soon_threadsafe from the dbapi sync-timeout arm, heartbeat
invalidations, transaction()'s rollback-failure arm) lets a
concurrent task enter a _run_protocol-protected critical section
while the original claimant is still mid-await — and the
claimant's own finally then clobbers the new claim. The in-flight
task observes the nulled _protocol on its next read/write and
clears the flag via its own finally.

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

Author: antoineleclair

src/dqliteclient/connection.py

@@ -2323,14 +2323,17 @@ def _invalidate(self, cause: BaseException | None = None) -> None:
         If ``cause`` is provided, it is remembered so a later caller that
         hits "Not connected" can chain it as ``__cause__`` for diagnostics.
 
-        Also clears the ``_in_use`` slot: invalidation can be invoked
-        out-of-band (e.g. scheduled from the dbapi sync-timeout path
-        via ``call_soon_threadsafe``), bypassing ``_run_protocol``'s
-        finally that normally resets the flag. An invalidated connection
-        is no longer holding a meaningful in-flight operation, so the
-        flag and the liveness state must stay consistent — otherwise
-        the next call deterministically raises "another operation is
-        in progress" on a connection that is in fact dead.
+        Does NOT touch the ``_in_use`` slot: the flag's contract is
+        "owned by the task that claimed it between claim-site and
+        clear-site". Clearing it here — when ``_invalidate`` was
+        invoked out-of-band (e.g. ``call_soon_threadsafe`` from the
+        dbapi sync-timeout path) — would let a second task into a
+        ``_run_protocol``-protected critical section while the original
+        in-flight task is still mid-``await``. The in-flight task
+        unblocks promptly once we null ``self._protocol`` below (its
+        next read raises a transport error), and its own
+        ``finally`` then clears ``_in_use`` — restoring the
+        "transient" wait the next caller would see anyway.
 
         Synchronous writer-close + async bounded drain: ``protocol.close()``
         is synchronous (writer.close()), but ``wait_closed()`` is a
@@ -2432,7 +2435,13 @@ async def _bounded_drain() -> None:
                     self._pending_drain = None
         self._protocol = None
         self._db_id = None
-        self._in_use = False
+        # NOTE: ``_in_use`` is intentionally NOT cleared here. See the
+        # method docstring — the flag's task-ownership contract requires
+        # that only the task that claimed it clears it (typically via
+        # ``_run_protocol``'s finally at line ~2728 or ``connect()``'s
+        # finally arms). Clearing it from this out-of-band path can let
+        # a concurrent task enter ``_run_protocol`` while the original
+        # claimant is still mid-await.
         # Atomic clear of transaction-state flags: an external
         # invalidation (heartbeat path, KeyboardInterrupt mid-yield,
         # ``call_soon_threadsafe``) lands here without going through

tests/test_connection.py

@@ -422,19 +422,24 @@ async def test_connect_emits_debug_logs_on_handshake_and_open(
         assert any("handshake ok" in m and "localhost:9001" in m for m in messages)
         assert any("db opened" in m and "localhost:9001" in m for m in messages)
 
-    async def test_invalidate_clears_in_use_flag(self, connected_connection) -> None:
-        """_invalidate must reset _in_use alongside dropping the protocol.
-
-        If an out-of-band invalidation (e.g. scheduled from a sync-side
-        timeout path) fires while _in_use=True, leaving the flag set
-        deterministically blocks the next call on this connection with
-        "another operation is in progress" — a misleading error for a
-        connection that is in fact dead and needs reconnecting.
+    async def test_invalidate_preserves_in_use_flag(self, connected_connection) -> None:
+        """_invalidate must NOT clear _in_use.
+
+        The flag's contract is "owned by the task that claimed it
+        between claim-site and clear-site". Clearing it from an
+        out-of-band path (``call_soon_threadsafe``, heartbeat,
+        ``transaction()`` rollback-failure arm) would let a second task
+        enter a ``_run_protocol``-protected critical section while the
+        original claimant is still mid-``await``. The in-flight task
+        observes the nulled ``_protocol`` and runs its own ``finally``
+        to clear the flag; the next caller transiently sees "another
+        operation is in progress" until that finally runs — which is
+        the correct behaviour for an in-flight claim.
         """
         conn, _, _ = connected_connection
         conn._in_use = True
         conn._invalidate(OperationalError("synthetic", 0))
-        assert conn._in_use is False
+        assert conn._in_use is True
         assert conn._protocol is None
 
     async def test_invalidate_closes_transport(self, connected_connection) -> None:

tests/test_invalidate_does_not_clobber_in_use.py

@@ -0,0 +1,104 @@
+"""Pin: ``_invalidate`` does NOT clear ``_in_use``.
+
+The flag's contract is "owned by the task that claimed it between
+claim-site and clear-site". Clearing it from an out-of-band path
+(``call_soon_threadsafe`` from the dbapi sync-timeout arm,
+heartbeat invalidations, ``transaction()`` rollback-failure arm)
+would let a second task enter a ``_run_protocol``-protected
+critical section while the original claimant is still mid-``await``.
+
+The in-flight task observes the nulled ``_protocol`` at its next
+read / write and runs its own ``finally`` to clear ``_in_use`` —
+the post-fix shape never lets a concurrent caller see a
+spuriously-cleared flag mid-await.
+"""
+
+from __future__ import annotations
+
+import weakref
+
+import pytest
+
+from dqliteclient.connection import DqliteConnection
+from dqliteclient.exceptions import InterfaceError
+
+
+@pytest.mark.asyncio
+async def test_invalidate_preserves_in_use_so_concurrent_claimant_is_rejected() -> None:
+    """Task A holds ``_in_use=True`` (the in-flight claimant). A
+    concurrent ``_invalidate`` lands. The flag must NOT be cleared
+    here — a second task hitting ``_check_in_use`` must still see
+    the in-progress claim and raise.
+    """
+    import asyncio
+
+    loop = asyncio.get_running_loop()
+
+    conn = DqliteConnection.__new__(DqliteConnection)
+    conn._in_use = True  # simulate Task A's mid-_run_protocol claim
+    conn._protocol = object()  # type: ignore[assignment]
+    conn._db_id = 5
+    conn._in_transaction = False
+    conn._tx_owner = None
+    conn._savepoint_stack = []
+    conn._savepoint_implicit_begin = False
+    conn._has_untracked_savepoint = False
+    conn._invalidation_cause = None
+    conn._bound_loop_ref = weakref.ref(loop)
+    conn._close_timeout = 5.0
+    conn._pending_drain = None
+    conn._address = "127.0.0.1:1"
+    import os
+
+    conn._creator_pid = os.getpid()
+    conn._pool_released = False
+
+    # Out-of-band _invalidate (e.g. from call_soon_threadsafe).
+    conn._invalidate(RuntimeError("simulated heartbeat invalidation"))
+
+    # Load-bearing: the flag was NOT cleared. A second task hitting
+    # ``_check_in_use`` must observe the still-claimed state.
+    assert conn._in_use is True, (
+        "_invalidate must not clear _in_use; otherwise a concurrent "
+        "task can enter _run_protocol's critical section while the "
+        "original claimant is still mid-await"
+    )
+
+    # ``_check_in_use`` rejects with InterfaceError.
+    with pytest.raises(InterfaceError, match="another operation is in progress"):
+        conn._check_in_use()
+
+
+@pytest.mark.asyncio
+async def test_invalidate_with_no_active_claimant_leaves_in_use_false() -> None:
+    """When no task holds the flag (``_in_use=False`` going in),
+    ``_invalidate`` must not flip it to True. The fix is a delete, not
+    a swap.
+    """
+    import asyncio
+
+    loop = asyncio.get_running_loop()
+
+    conn = DqliteConnection.__new__(DqliteConnection)
+    conn._in_use = False
+    conn._protocol = object()  # type: ignore[assignment]
+    conn._db_id = 5
+    conn._in_transaction = False
+    conn._tx_owner = None
+    conn._savepoint_stack = []
+    conn._savepoint_implicit_begin = False
+    conn._has_untracked_savepoint = False
+    conn._invalidation_cause = None
+    conn._bound_loop_ref = weakref.ref(loop)
+    conn._close_timeout = 5.0
+    conn._pending_drain = None
+    conn._address = "127.0.0.1:1"
+    import os
+
+    conn._creator_pid = os.getpid()
+    conn._pool_released = False
+
+    conn._invalidate(RuntimeError("simulated"))
+
+    # The flag was untouched — and is still safely false.
+    assert conn._in_use is False