From b1f97ca5f82936e1bfb3f3afa2e361f1dc69bb9c Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 22 Oct 2026 06:42:09 +0100 Subject: [PATCH 1/2] fix: match interrupted connections by the connection, not a reusable id Deadline cancellation marked a connection by `id(connection)` in a process-wide set that only a recovering pool's health check cleared. A mark left on a connection that never returned to such a pool outlived it, and a later connection allocated at the same address was then health-checked as if interrupted, running an extra probe statement. Each mark now holds a weak reference to its connection, or a strong one when the connection cannot be weakly referenced, and a check reports only when the mark still refers to that same connection. --- .../classes/interrupted_connections.py | 25 +++++-- .../runtime/execution_limits/_test_types.py | 14 ++++ .../runtime/execution_limits/helpers.py | 4 ++ .../test_interrupted_connections.py | 69 +++++++++++++++++++ 4 files changed, 105 insertions(+), 7 deletions(-) create mode 100644 tests/unit/src/sqlbuild/runtime/execution_limits/test_interrupted_connections.py diff --git a/src/sqlbuild/runtime/execution_limits/classes/interrupted_connections.py b/src/sqlbuild/runtime/execution_limits/classes/interrupted_connections.py index 4e11723474..60ac61615f 100644 --- a/src/sqlbuild/runtime/execution_limits/classes/interrupted_connections.py +++ b/src/sqlbuild/runtime/execution_limits/classes/interrupted_connections.py @@ -3,24 +3,35 @@ from __future__ import annotations import threading +import weakref +from collections.abc import Callable from typing import Any class InterruptedConnections: - """Identities of interrupted connections, each cleared when it is checked.""" + """Interrupted connections, held weakly where possible and each cleared when it is checked.""" def __init__(self) -> None: - self._connections: set[int] = set() + self._connections: dict[int, Callable[[], Any]] = {} self._lock: threading.Lock = threading.Lock() def mark(self, connection: Any) -> None: + reference: Callable[[], Any] = _reference(connection) with self._lock: - self._connections.add(id(connection)) + self._connections[id(connection)] = reference def take(self, connection: Any) -> bool: - """Whether `connection` was interrupted since it was last checked, clearing the mark.""" + """Whether `connection` itself was interrupted since its last check, clearing the mark.""" with self._lock: - interrupted: bool = id(connection) in self._connections - self._connections.discard(id(connection)) - return interrupted + reference: Callable[[], Any] | None = self._connections.pop(id(connection), None) + return reference is not None and reference() is connection + + +def _reference(connection: Any) -> Callable[[], Any]: + """A weak reference to `connection`, or a strong one when it cannot be weakly referenced.""" + + try: + return weakref.ref(connection) + except TypeError: + return lambda: connection diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py b/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py index d9bcba5fe7..fc2c0e24df 100644 --- a/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/_test_types.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from dataclasses import dataclass from sqlbuild.runtime.execution_limits.models import ItemTimeoutPolicy @@ -79,3 +80,16 @@ class SharedQueuedClockTestCase: waited_seconds: float expected_minimum_running_seconds: float expected_maximum_running_seconds: float + + +@dataclass(frozen=True) +class InterruptedConnectionReportsTestCase: + description: str + new_connection: Callable[[], object] + expected_reports: tuple[bool, bool] + + +@dataclass(frozen=True) +class InterruptedIdentityReuseTestCase: + description: str + expected_reuse_and_report: tuple[bool, bool] diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py b/tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py index 1ae300b0bf..513d70e6be 100644 --- a/tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/helpers.py @@ -58,3 +58,7 @@ def run_finished_statement() -> None: assert deadline is not None deadline.statement_started() deadline.statement_finished(query_id=None) + + +class PlaceholderConnection: + """A connection stand-in that, like real adapter connections, supports weak references.""" diff --git a/tests/unit/src/sqlbuild/runtime/execution_limits/test_interrupted_connections.py b/tests/unit/src/sqlbuild/runtime/execution_limits/test_interrupted_connections.py new file mode 100644 index 0000000000..e1bc69d7be --- /dev/null +++ b/tests/unit/src/sqlbuild/runtime/execution_limits/test_interrupted_connections.py @@ -0,0 +1,69 @@ +"""Interrupted-connection marks match only the connection that was interrupted.""" + +from __future__ import annotations + +import pytest + +from sqlbuild.runtime.execution_limits.classes.interrupted_connections import ( + InterruptedConnections, +) +from tests.unit.src.sqlbuild.runtime.execution_limits._test_types import ( + InterruptedConnectionReportsTestCase, + InterruptedIdentityReuseTestCase, +) +from tests.unit.src.sqlbuild.runtime.execution_limits.helpers import PlaceholderConnection + + +@pytest.mark.parametrize( + "test_case", + [ + InterruptedConnectionReportsTestCase( + description="weakly referenced connection is reported once", + new_connection=PlaceholderConnection, + expected_reports=(True, False), + ), + InterruptedConnectionReportsTestCase( + description="connection without weak references is reported once", + new_connection=object, + expected_reports=(True, False), + ), + ], + ids=lambda case: case.description, +) +def test_given_marked_connection_when_checked_twice_then_reported_once( + test_case: InterruptedConnectionReportsTestCase, +) -> None: + connections: InterruptedConnections = InterruptedConnections() + connection: object = test_case.new_connection() + connections.mark(connection) + + reports: tuple[bool, bool] = (connections.take(connection), connections.take(connection)) + + assert reports == test_case.expected_reports + + +@pytest.mark.parametrize( + "test_case", + [ + InterruptedIdentityReuseTestCase( + description="new connection at a collected interrupted identity is not reported", + expected_reuse_and_report=(True, False), + ), + ], + ids=lambda case: case.description, +) +def test_given_collected_marked_connection_when_new_connection_reuses_identity_then_not_reported( + test_case: InterruptedIdentityReuseTestCase, +) -> None: + connections: InterruptedConnections = InterruptedConnections() + interrupted: PlaceholderConnection = PlaceholderConnection() + interrupted_identity: int = id(interrupted) + connections.mark(interrupted) + del interrupted + + reused: PlaceholderConnection = PlaceholderConnection() + + assert ( + id(reused) == interrupted_identity, + connections.take(reused), + ) == test_case.expected_reuse_and_report From 988e8ca8845ef0d9035c5370aa4acb3c4a5e4340 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 22 Oct 2026 07:20:33 +0100 Subject: [PATCH 2/2] test: raise the 1000-model build memory guard to 1760 MiB Main runs at 1,778-1,808 MB on CI runners against a 1,812 MB cap, so unrelated branches fail the guard at random (about +/-15 MB runner noise). 1,760 MiB leaves about 31 MB above the highest CI value seen. --- .../cli/commands/main/compile/test_build_performance.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_build_performance.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_build_performance.py index ff1b53d553..e0a6a9bba8 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_build_performance.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/test_build_performance.py @@ -38,7 +38,7 @@ description="models_3000_job_build_1000_model_benchmark_from_empty_warehouse", expected_max_wall_seconds=90.0, # Temporary native-migration headroom pending memory optimisation. - expected_max_rss_bytes=1_728 * _MIB, + expected_max_rss_bytes=1_760 * _MIB, ) ], ids=lambda case: case.description,