fix: keep pool slots, test transactions and shared clocks correct after timeouts

A connection that can neither be recovered nor replaced after a timeout
is now returned to its pool marked broken, so later work fails with the
driver's error instead of waiting forever for a lost slot; each later
return retries the replacement, and microbatch subworker slots are
always released.

A PostgreSQL unit test claims its transaction only after BEGIN
succeeds, so a refused BEGIN can no longer make a later release roll
back a caller's transaction.

A queued BigQuery job pauses an item's clock only while every statement
under that deadline is queued, so one queued microbatch worker no
longer pauses a model whose other workers are running.

Author: kvlonge

src/sqlbuild/adapters/bigquery/classes/bigquery_connection.py

@@ -99,13 +99,13 @@ def _wait_while_pending(*, job: Any, deadline: ItemDeadline | None) -> None:
 
     if deadline is None or getattr(job, "state", None) != PENDING_JOB_STATE:
         return
-    deadline.hold()
+    deadline.statement_queued()
     try:
         while getattr(job, "state", None) == PENDING_JOB_STATE:
             time.sleep(PENDING_JOB_POLL_SECONDS)
             job.reload()
     finally:
-        deadline.release()
+        deadline.statement_dequeued()
 
 
 def _job_id(*, job: Any) -> str | None:

src/sqlbuild/adapters/postgres/classes/postgres_adapter.py

@@ -148,8 +148,9 @@ def get_columns_for_relations(
     def create_sql_test_output(self, *, connection: Any, sql: str, name: str) -> str | None:
         """Store a unit test's output in a temporary table of one transaction, for poolers."""
 
-        if connection.claim_sql_test_transaction():
+        if connection.needs_sql_test_transaction():
             _ = self.execute(connection=connection, sql="BEGIN")
+            connection.claim_sql_test_transaction()
         try:
             _ = self.execute(
                 connection=connection,

src/sqlbuild/adapters/postgres/classes/postgres_connection.py

@@ -56,13 +56,15 @@ def _apply_statement_timeout(self) -> None:
         self._local_timeout_set = True
         self._local_timeout_deadline = id(deadline)
 
-    def claim_sql_test_transaction(self) -> bool:
-        """Whether a unit test must open its own transaction, which it then owns."""
+    def needs_sql_test_transaction(self) -> bool:
+        """Whether a unit test must open its own transaction before storing outputs."""
+
+        return not self._sql_test_transaction and not self.in_transaction()
+
+    def claim_sql_test_transaction(self) -> None:
+        """Record that the transaction a unit test just opened is its own to end."""
 
-        if self._sql_test_transaction or self.in_transaction():
-            return False
         self._sql_test_transaction = True
-        return True
 
     def release_sql_test_transaction(self) -> bool:
         """Whether a unit test owned the open transaction, which it must now end."""

src/sqlbuild/executor/build/classes/build_scheduler.py

@@ -680,8 +680,10 @@ def _run_microbatch_batch_worker(
         try:
             return execute(batch, connection)
         finally:
-            self._connection_pool.put(connection)
-            self._release_microbatch_subworker()
+            try:
+                self._connection_pool.put(connection)
+            finally:
+                self._release_microbatch_subworker()
 
     def _reserve_microbatch_subworker(self) -> bool:
         with self._microbatch_coordinator_lock:

src/sqlbuild/executor/scheduling/classes/recovering_connection_pool.py

@@ -16,4 +16,8 @@ def __init__(self, *, recovery: ConnectionRecovery) -> None:
         self._recovery: ConnectionRecovery = recovery
 
     def put(self, item: Any, block: bool = True, timeout: float | None = None) -> None:
-        super().put(self._recovery.usable(item), block=block, timeout=timeout)
+        returned: Any = item
+        try:
+            returned = self._recovery.usable(item)
+        finally:
+            super().put(returned, block=block, timeout=timeout)

src/sqlbuild/runtime/execution_limits/classes/connection_recovery.py

@@ -23,24 +23,41 @@ def __init__(self, *, adapter: Any, connect: Callable[[], Any]) -> None:
         self._adapter: Any = adapter
         self._connect: Callable[[], Any] = connect
         self._replacements: list[Any] = []
+        self._broken: set[int] = set()
         self._lock: threading.Lock = threading.Lock()
 
     def usable(self, connection: Any) -> Any:
-        """`connection` when it is healthy, otherwise a newly opened replacement."""
+        """`connection` if healthy, else a replacement, else the original kept marked broken."""
 
-        if not interrupted_connections.take(connection):
+        interrupted: bool = interrupted_connections.take(connection)
+        with self._lock:
+            broken: bool = id(connection) in self._broken
+        if not interrupted and not broken:
             return connection
         with execution_deadlines_suspended():
             if self._responds(connection):
+                self._forget_broken(connection)
                 return connection
             try:
                 self._adapter.rollback(connection)
             except Exception as error:
                 _LOGGER.debug("rollback after an interrupted statement failed: %s", error)
             if self._responds(connection):
+                self._forget_broken(connection)
                 return connection
+        try:
+            replacement: Any = self._connect()
+        except Exception as error:
+            with self._lock:
+                self._broken.add(id(connection))
+            _LOGGER.warning(
+                "A connection left unusable by a timed-out statement could not be replaced (%s); "
+                "work that uses it will fail",
+                error,
+            )
+            return connection
         _LOGGER.warning("A connection left unusable by a timed-out statement was replaced")
-        replacement: Any = self._connect()
+        self._forget_broken(connection)
         with self._lock:
             self._replacements.append(replacement)
         return replacement
@@ -57,6 +74,10 @@ def close_replacements(self) -> None:
             except Exception as error:
                 _LOGGER.debug("closing a replacement connection failed: %s", error)
 
+    def _forget_broken(self, connection: Any) -> None:
+        with self._lock:
+            self._broken.discard(id(connection))
+
     def _responds(self, connection: Any) -> bool:
         try:
             _ = self._adapter.execute(connection=connection, sql=CONNECTION_PROBE_SQL)

src/sqlbuild/runtime/execution_limits/classes/item_deadline.py

@@ -46,6 +46,7 @@ def __init__(
         self._running_since: float | None = None
         self._clock_started: bool = False
         self._holds: int = 0
+        self._queued_statements: int = 0
         self._closed: bool = False
         self._active_statements: int = 0
         self._awaiting_confirmation: bool = False
@@ -102,6 +103,7 @@ def statement_finished(self, *, query_id: str | None) -> None:
                 deadline._active_statements = max(0, deadline._active_statements - 1)
                 if query_id:
                     deadline._query_ids.append(query_id)
+                deadline._update_clock()
 
     def statement_submitted(self, *, query_id: str | None) -> None:
         if not query_id:
@@ -111,18 +113,32 @@ def statement_submitted(self, *, query_id: str | None) -> None:
                 deadline._query_ids.append(query_id)
 
     def hold(self) -> None:
-        """Stop the clock, for warehouse queueing or cleanup that must not count."""
+        """Stop the clock for cleanup while no statement runs under this deadline."""
 
         for deadline in self.chain():
             with deadline._lock:
                 deadline._holds += 1
-                deadline._pause_clock()
+                deadline._update_clock()
 
     def release(self) -> None:
         for deadline in self.chain():
             with deadline._lock:
                 deadline._holds = max(0, deadline._holds - 1)
-                deadline._resume_clock()
+                deadline._update_clock()
+
+    def statement_queued(self) -> None:
+        """Mark a running statement as waiting in a warehouse queue; only queued work pauses."""
+
+        for deadline in self.chain():
+            with deadline._lock:
+                deadline._queued_statements += 1
+                deadline._update_clock()
+
+    def statement_dequeued(self) -> None:
+        for deadline in self.chain():
+            with deadline._lock:
+                deadline._queued_statements = max(0, deadline._queued_statements - 1)
+                deadline._update_clock()
 
     def close(self) -> None:
         """Stop the clock and the watchdog when the item's scope ends."""
@@ -198,6 +214,7 @@ def _start_statement(self) -> None:
         with self._lock:
             self._active_statements += 1
             if self._clock_started or self._closed:
+                self._update_clock()
                 return
             self._clock_started = True
             self._resume_clock()
@@ -212,9 +229,21 @@ def _pause_clock(self) -> None:
             self._running_since = None
 
     def _resume_clock(self) -> None:
-        if self._clock_started and not self._closed and not self._holds:
+        if self._clock_started and not self._closed and not self._paused():
             self._running_since = self._running_since or time.monotonic()
 
+    def _paused(self) -> bool:
+        """Paused while held or queued, unless another statement is running meanwhile."""
+
+        held: bool = bool(self._holds or self._queued_statements)
+        return held and self._active_statements <= self._queued_statements
+
+    def _update_clock(self) -> None:
+        if self._paused():
+            self._pause_clock()
+        else:
+            self._resume_clock()
+
     def _watch(self) -> None:
         while True:
             remaining: float = self._own_remaining(grace=self._grace)

tests/unit/src/sqlbuild/adapters/postgres/_test_types.py

@@ -152,3 +152,12 @@ class FailedSqlTestOutputTestCase:
     description: str
     expected_statements: tuple[str, ...]
     expected_in_transaction: bool
+
+
+@dataclass(frozen=True)
+class RefusedBeginTestCase:
+    """A unit test whose BEGIN the deadline refuses, followed by work in a caller transaction."""
+
+    description: str
+    expected_statements: tuple[str, ...]
+    expected_in_transaction: bool

tests/unit/src/sqlbuild/adapters/postgres/test_deadline_transactions.py

@@ -20,6 +20,7 @@
 from tests.unit.src.sqlbuild.adapters.postgres._test_types import (
     DeadlineTransactionTestCase,
     FailedSqlTestOutputTestCase,
+    RefusedBeginTestCase,
     SqlTestOutputTransactionTestCase,
 )
 
@@ -154,5 +155,39 @@ def test_given_failing_output_query_when_storing_then_transaction_is_rolled_back
     assert connection.in_transaction() is test_case.expected_in_transaction
 
 
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        RefusedBeginTestCase(
+            description="a refused BEGIN leaves no ownership that could end a caller transaction",
+            expected_statements=("BEGIN", f"DROP TABLE IF EXISTS {_OUTPUT}"),
+            expected_in_transaction=True,
+        )
+    ],
+    ids=lambda case: case.description,
+)
+def test_given_refused_begin_when_storing_output_then_later_release_keeps_caller_transaction(
+    test_case: RefusedBeginTestCase,
+) -> None:
+    adapter: PostgresAdapter = PostgresAdapter()
+    raw: CancellableRawConnection = CancellableRawConnection()
+    connection: _PostgresConnection = _PostgresConnection(raw)
+
+    with item_deadline_scope(
+        kind=ItemKind.TEST, name="orders", policy=_POLICY, adapter=adapter, connection=connection
+    ) as deadline:
+        assert deadline is not None
+        deadline.expire()
+        with pytest.raises(ItemTimeoutError):
+            _ = adapter.create_sql_test_output(
+                connection=connection, sql="SELECT 1 AS order_id", name=_OUTPUT
+            )
+    _ = adapter.execute(connection=connection, sql="BEGIN")
+    adapter.drop_sql_test_output(connection=connection, relation=_OUTPUT)
+
+    assert without_limits(raw.executed) == test_case.expected_statements
+    assert connection.in_transaction() is test_case.expected_in_transaction
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-vv"])

tests/unit/src/sqlbuild/executor/scheduling/classes/_test_types.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class UnreplaceableConnectionTestCase:
+    """A timed-out item whose connection can neither be recovered nor replaced."""
+
+    description: str
+    later_items: int
+    expected_completions: tuple[tuple[str, str], ...]
+    expected_reconnect_attempts: int
+    expected_maximum_seconds: float