fix: match interrupted connections by the connection, not a reusable id

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

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,

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]

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."""

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