Unblock a concurrent _run_sync caller on force_close_transport

force_close_transport() stopped the loop without resolving an outstanding
run_coroutine_threadsafe future, so a concurrent _run_sync caller (under
check_same_thread=False) waited out the full sync_timeout (4×timeout, ~40s by
default) before a clean OperationalError. Set a dedicated _force_close_requested
flag before stopping the loop — distinct from _closed_flag, which close() sets
and then runs its own _run_sync(_close_async()) that must not bail — and wait on
the future in short slices so the caller observes the flag and raises
InterfaceError promptly. When the flag is unset the slice loop is behaviourally
identical: it returns on completion, propagates a coroutine error, and re-raises
the genuine TimeoutError at the deadline into the existing timeout arm.

The _run_sync timeout stubs keyed on a call counter that assumed a single
result() call; rework them to key on the actual control flow (done() checked for
the recovery arm, cancel() called for the bounded wait) so they hold under the
sliced wait.

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

Author: antoineleclair

src/dqlitedbapi/connection.py

@@ -7,6 +7,7 @@
 import math
 import os
 import threading
+import time
 import warnings
 import weakref
 from collections.abc import Coroutine, Iterable, Iterator, Sequence
@@ -130,6 +131,11 @@ def _join_budget_for_current_thread(
 # N=4 covers the worst case: handshake + open_database + send +
 # read+drain. Steady-state bottoms out at N=2.
 _SYNC_PHASES_MULTIPLIER: Final[int] = 4
+# Slice length for _run_sync's result wait: short enough that a concurrent
+# force_close_transport unblocks the caller promptly (it sets
+# _force_close_requested before stopping the loop), long enough that a normal
+# in-flight op adds no measurable wakeup overhead.
+_FORCE_CLOSE_POLL_INTERVAL: Final[float] = 0.05
 
 # Fallback join budget for ``_cleanup_loop_thread`` when the captured
 # ``close_timeout`` is missing or invalid.
@@ -1014,6 +1020,12 @@ def __init__(
         # Mutable flag the finalizer reads; a list avoids the finalizer
         # closing over ``self`` and preventing GC.
         self._closed_flag: list[bool] = [False]
+        # Set by force_close_transport before it abruptly stops the loop, so a
+        # concurrent _run_sync (check_same_thread=False) bails promptly instead of
+        # riding out sync_timeout on a future the stopped loop will never resolve.
+        # Deliberately NOT _closed_flag: close() sets that and then runs its own
+        # _run_sync(self._close_async()), which must not bail.
+        self._force_close_requested: bool = False
         # Box for late-publishing the inner handle into the finalizer's
         # captured args (mutated by _get_async_connection to a
         # weakref.ref). Not cleared on explicit close (finalizer is
@@ -1208,8 +1220,36 @@ def _run_sync[T](self, coro: Coroutine[Any, Any, T]) -> T:
                         f"event loop closed before coroutine could be scheduled: {e}"
                     ) from e
                 # Future.result() is a happens-before barrier; loop-thread
-                # writes are visible here.
-                return future.result(timeout=sync_timeout)
+                # writes are visible here. Wait in slices (rather than one
+                # result(timeout=sync_timeout)) so a concurrent
+                # force_close_transport — which sets _force_close_requested
+                # and then stops the loop, orphaning this future — unblocks us
+                # promptly instead of riding out the full sync_timeout. When the
+                # flag is not set this is behaviourally identical: it re-raises
+                # the genuine TimeoutError at the deadline into the arm below.
+                deadline = time.monotonic() + sync_timeout
+                while True:
+                    remaining = deadline - time.monotonic()
+                    try:
+                        return future.result(
+                            timeout=max(0.0, min(_FORCE_CLOSE_POLL_INTERVAL, remaining))
+                        )
+                    except TimeoutError:
+                        if self._force_close_requested:
+                            # Honour a result that already landed (avoid a
+                            # spurious failure / non-idempotent retry), else bail
+                            # with the closed-connection contract rather than the
+                            # timeout/invalidate path below.
+                            if future.done() and not future.cancelled():
+                                return future.result(timeout=0)
+                            future.cancel()
+                            raise InterfaceError(
+                                f"Connection force-closed during operation (id={id(self)})"
+                            ) from None
+                        if remaining <= _FORCE_CLOSE_POLL_INTERVAL:
+                            # Genuine sync_timeout: re-raise to the arm below
+                            # (recovered-result check + invalidate).
+                            raise
             except TimeoutError as e:
                 # Only future.result can raise this, so future is bound;
                 # assert is for mypy.
@@ -1570,6 +1610,9 @@ def force_close_transport(self) -> None:
         if self._closed:
             return
         self._closed = True
+        # Signal a concurrent in-flight _run_sync to bail BEFORE we stop the loop
+        # (its future would otherwise never resolve). Set before any teardown.
+        self._force_close_requested = True
         self._closed_flag[0] = True
         # Fork-after-init: same shape as close()'s pid guard.
         if get_current_pid() != self._creator_pid:

tests/integration/test_force_close_unblocks_run_sync.py

@@ -0,0 +1,55 @@
+"""force_close_transport() must unblock a concurrent in-flight _run_sync caller promptly,
+not leave it riding out the full sync_timeout (4×timeout) on a future the stopped loop will
+never resolve. (check_same_thread=False; force_close is the documented last-resort path.)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import threading
+import time
+
+import pytest
+
+import dqlitedbapi
+from dqlitedbapi import InterfaceError, OperationalError
+
+
+@pytest.mark.integration
+def test_force_close_unblocks_concurrent_run_sync(cluster_address: str) -> None:
+    # sync_timeout = 4 * 5.0 = 20s; pre-fix the concurrent caller wedged the full 20s.
+    conn = dqlitedbapi.connect(cluster_address, timeout=5.0, check_same_thread=False)
+    try:
+        conn.cursor().execute("SELECT 1").fetchone()  # warm up the loop + wire
+
+        result: dict[str, object] = {}
+        ready = threading.Event()
+
+        def slow_op() -> None:
+            ready.set()
+            t0 = time.monotonic()
+            try:
+                conn._run_sync(asyncio.sleep(30))  # stand-in for a slow in-flight wire op
+            except BaseException as e:  # noqa: BLE001
+                result["exc"] = e
+                result["elapsed"] = time.monotonic() - t0
+
+        t = threading.Thread(target=slow_op)
+        t.start()
+        ready.wait(timeout=2.0)
+        time.sleep(0.3)  # let _run_sync acquire the op_lock and enter the result wait
+
+        conn.force_close_transport()
+
+        t.join(timeout=10.0)
+        assert not t.is_alive(), "concurrent _run_sync caller did not unblock"
+        assert "exc" in result, "the in-flight op did not raise on force-close"
+        elapsed = result["elapsed"]
+        assert isinstance(elapsed, float)
+        # The whole point: unblocked promptly, NOT after the 20s sync_timeout.
+        assert elapsed < 3.0, f"caller wedged {elapsed:.1f}s (sync_timeout was 20s)"
+        assert isinstance(result["exc"], (InterfaceError, OperationalError))
+    finally:
+        with contextlib.suppress(Exception):
+            conn.close()

tests/test_run_sync_timeout.py

@@ -33,9 +33,13 @@ def test_run_sync_logs_unexpected_error_during_cancel_wait(
 
         class _StubFuture:
             def __init__(self) -> None:
-                self._calls = 0
+                self._cancel_called = False
 
             def cancel(self) -> bool:
+                # The except-TimeoutError arm calls cancel() before the bounded wait;
+                # flip so only that bounded wait raises the surprise error, not the
+                # slice-poll's repeated result() calls.
+                self._cancel_called = True
                 return True
 
             def done(self) -> bool:
@@ -46,10 +50,9 @@ def cancelled(self) -> bool:
                 return False
 
             def result(self, timeout: float | None = None) -> None:
-                self._calls += 1
-                if self._calls == 1:
-                    raise cf.TimeoutError()
-                raise RuntimeError("surprise bug during cancel")
+                if self._cancel_called:
+                    raise RuntimeError("surprise bug during cancel")
+                raise cf.TimeoutError()
 
         stub = _StubFuture()
 
@@ -128,22 +131,25 @@ class _LateSuccessFuture:
             """result() raises TimeoutError first, then returns the late success."""
 
             def __init__(self) -> None:
-                self._calls = 0
+                self._done_checked = False
 
             def cancel(self) -> bool:
                 return False  # cancel lost the race
 
             def done(self) -> bool:
+                # Only the except-TimeoutError recovery arm calls done(); flip so
+                # result() yields the late success only there, not during the
+                # slice-poll's repeated result() calls.
+                self._done_checked = True
                 return True
 
             def cancelled(self) -> bool:
                 return False
 
             def result(self, timeout: float | None = None) -> int:
-                self._calls += 1
-                if self._calls == 1:
-                    raise cf.TimeoutError()
-                return 1234
+                if self._done_checked:
+                    return 1234
+                raise cf.TimeoutError()
 
         stub = _LateSuccessFuture()
 
@@ -184,24 +190,27 @@ class _LateIntegrityErrorFuture:
             """result() raises TimeoutError first, then the recovered IntegrityError."""
 
             def __init__(self) -> None:
-                self._calls = 0
+                self._done_checked = False
 
             def cancel(self) -> bool:
                 return False
 
             def done(self) -> bool:
+                # Only the recovery arm calls done(); see _LateSuccessFuture.
+                self._done_checked = True
                 return True
 
             def cancelled(self) -> bool:
                 return False
 
             def result(self, timeout: float | None = None) -> Any:
-                self._calls += 1
-                if self._calls == 1:
-                    raise cf.TimeoutError()
-                raise IntegrityError(
-                    "UNIQUE constraint failed", code=2067, raw_message="UNIQUE constraint failed"
-                )
+                if self._done_checked:
+                    raise IntegrityError(
+                        "UNIQUE constraint failed",
+                        code=2067,
+                        raw_message="UNIQUE constraint failed",
+                    )
+                raise cf.TimeoutError()
 
         stub = _LateIntegrityErrorFuture()
 
@@ -235,22 +244,25 @@ def test_recovered_exception_carries_timeout_in_context(
 
         class _Stub:
             def __init__(self) -> None:
-                self._calls = 0
+                self._done_checked = False
 
             def cancel(self) -> bool:
                 return False
 
             def done(self) -> bool:
+                # Recovery happens via the except-TimeoutError arm (which calls done()),
+                # so the recovered IntegrityError's __context__ is the re-raised
+                # TimeoutError; raising it during a slice would lose that context.
+                self._done_checked = True
                 return True
 
             def cancelled(self) -> bool:
                 return False
 
             def result(self, timeout: float | None = None) -> Any:
-                self._calls += 1
-                if self._calls == 1:
-                    raise cf.TimeoutError()
-                raise IntegrityError("constraint failed", code=2067, raw_message="x")
+                if self._done_checked:
+                    raise IntegrityError("constraint failed", code=2067, raw_message="x")
+                raise cf.TimeoutError()
 
         stub = _Stub()
         monkeypatch.setattr(
@@ -295,9 +307,12 @@ def _invalidate(self, exc: BaseException | None = None) -> None:
 
         class _StubFuture:
             def __init__(self) -> None:
-                self._calls = 0
+                self._cancel_called = False
 
             def cancel(self) -> bool:
+                # The except-TimeoutError arm calls cancel() before the bounded wait;
+                # see the logs-unexpected-error stub above.
+                self._cancel_called = True
                 return True
 
             def done(self) -> bool:
@@ -307,10 +322,9 @@ def cancelled(self) -> bool:
                 return False
 
             def result(self, timeout: float | None = None) -> None:
-                self._calls += 1
-                if self._calls == 1:
-                    raise cf.TimeoutError()
-                raise cf.CancelledError()  # bounded-wait: coroutine unwound cleanly
+                if self._cancel_called:
+                    raise cf.CancelledError()  # bounded-wait: coroutine unwound cleanly
+                raise cf.TimeoutError()
 
         stub = _StubFuture()