Cache ClusterClient per (seed, governors); thread governors through _resolve_leader

Two compounding problems in the dbapi leader-discovery probe path
(connection._resolve_leader, called by _build_and_connect):

1. Configuration silently dropped at the first hop. _resolve_leader
   constructed ClusterClient(store, timeout=timeout) with no other
   governors, so an operator who set trust_server_heartbeat=True saw
   the FIRST round-trip — leader discovery — running with the default
   opt-out. Likewise max_total_rows / max_continuation_frames were
   dropped for admin paths reachable through the resolved client.

2. Single-flight collapse / leader-tracker fast-path dead at this
   layer. Each _resolve_leader call constructed a brand-new
   ClusterClient and discarded it on return; the _find_leader_tasks
   slot map and _last_known_leader cache went with it. Under N
   concurrent SA pool checkouts after a leader flip, the cluster
   sees N independent leader-discovery sweeps where one would
   suffice.

Thread the full governor set through _resolve_leader, and add a
process-wide ClusterClient cache keyed by (address, timeout,
max_total_rows, max_continuation_frames, trust_server_heartbeat).
The cache is wholesale-invalidated on fork via the same _current_pid
token DqliteConnection uses for fork-safety; bounded at 32 entries
(LRU-ish drop) so adversarial governor-fragmentation cannot leak
unbounded.

The existing cluster-level fixes (ISSUE-1402 single-flight, ISSUE-1403
last-known-leader) presupposed ClusterClient reuse; without this fix
those benefits were largely defeated at the dbapi layer.

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

Author: antoineleclair

src/dqlitedbapi/connection.py

@@ -147,7 +147,97 @@ def _validate_close_timeout(close_timeout: float) -> None:
         )
 
 
-async def _resolve_leader(address: str, *, timeout: float) -> str:
+# Process-wide ``ClusterClient`` cache for the leader-discovery probe.
+# Keyed by the full ``(address, governor)`` tuple so two configurations
+# never share state. Without the cache, every dbapi ``connect()`` /
+# every SA pool slot warm-up constructs a fresh ``ClusterClient`` —
+# discarding both the single-flight ``_find_leader_tasks`` slot map
+# AND the ``_last_known_leader`` fast-path cache. Under N concurrent
+# SA pool checkouts after a leader flip, the cluster sees N
+# independent leader-discovery sweeps where one would suffice.
+#
+# Fork-safety: the cache is wholesale-invalidated on fork via the same
+# ``_current_pid`` token that ``DqliteConnection`` uses (see
+# ``dqliteclient.connection`` lines 45-78). The first ``_resolve_leader``
+# call in a child process observes the pid mismatch and clears the
+# inherited cache; the parent's ``ClusterClient`` instances would
+# otherwise carry parent-allocated ``asyncio.Lock`` / ``asyncio.Task``
+# references that the child's event loop cannot make progress on.
+#
+# Strong-reference (``dict``, not ``WeakValueDictionary``): the
+# ClusterClient must outlive a single ``find_leader`` call so the
+# fast-path cache survives across calls; nothing else holds a
+# reference. The cap (``_RESOLVE_LEADER_CACHE_MAX``) bounds the worst
+# case to a single distinct configuration per dbapi ``connect()`` call;
+# typical SA deployments use one config per Engine, so the cap is
+# only reached by adversarial / highly-fragmented usage.
+_RESOLVE_LEADER_CACHE: dict[tuple[object, ...], ClusterClient] = {}
+_RESOLVE_LEADER_CACHE_PID: int = os.getpid()
+_RESOLVE_LEADER_CACHE_MAX: Final[int] = 32
+
+
+def _get_resolve_leader_cluster(
+    *,
+    address: str,
+    timeout: float,
+    max_total_rows: int | None,
+    max_continuation_frames: int | None,
+    trust_server_heartbeat: bool,
+) -> ClusterClient:
+    """Return a process-shared :class:`ClusterClient` for the
+    leader-discovery probe, keyed by the (address, governor) tuple.
+
+    The single-flight collapse and ``_last_known_leader`` fast-path
+    inside ``ClusterClient`` only amortise across callers of the
+    *same* instance. Constructing a fresh client per ``connect()``
+    defeats both. A process-wide cache restores the invariant.
+
+    Cleared wholesale on fork: ``ClusterClient`` instances inherit
+    parent ``asyncio.Lock`` / pending ``asyncio.Task`` references
+    that are bound to the parent's event loop and cannot make
+    progress in the child. The pid check is cheap (Python int
+    equality) and runs only on the cache-lookup path.
+    """
+    global _RESOLVE_LEADER_CACHE_PID
+    pid = _client_conn_mod._current_pid
+    if pid != _RESOLVE_LEADER_CACHE_PID:
+        _RESOLVE_LEADER_CACHE.clear()
+        _RESOLVE_LEADER_CACHE_PID = pid
+
+    key: tuple[object, ...] = (
+        address,
+        timeout,
+        max_total_rows,
+        max_continuation_frames,
+        trust_server_heartbeat,
+    )
+    cluster = _RESOLVE_LEADER_CACHE.get(key)
+    if cluster is None:
+        if len(_RESOLVE_LEADER_CACHE) >= _RESOLVE_LEADER_CACHE_MAX:
+            # Evict an arbitrary oldest entry. Worst case: the
+            # evicted client's fast-path cache is lost; the next
+            # ``find_leader`` against that key rediscovers in one
+            # sweep. Acceptable bound on the cache size.
+            _RESOLVE_LEADER_CACHE.pop(next(iter(_RESOLVE_LEADER_CACHE)))
+        cluster = ClusterClient(
+            MemoryNodeStore([address]),
+            timeout=timeout,
+            max_total_rows=max_total_rows,
+            max_continuation_frames=max_continuation_frames,
+            trust_server_heartbeat=trust_server_heartbeat,
+        )
+        _RESOLVE_LEADER_CACHE[key] = cluster
+    return cluster
+
+
+async def _resolve_leader(
+    address: str,
+    *,
+    timeout: float,
+    max_total_rows: int | None = _DEFAULT_MAX_TOTAL_ROWS,
+    max_continuation_frames: int | None = _DEFAULT_MAX_CONTINUATION_FRAMES,
+    trust_server_heartbeat: bool = False,
+) -> str:
     """Resolve the cluster's current leader address from a seed.
 
     Bootstraps from the user-supplied ``address`` (the URL host:port
@@ -160,14 +250,29 @@ async def _resolve_leader(address: str, *, timeout: float) -> str:
     cluster has a healthy leader at a different address; the SA
     pool's reconnect-after-pre-ping path cannot recover.
 
+    Threads the governor set used by the subsequent
+    :class:`DqliteConnection` so the leader-discovery probe runs
+    with the same configuration as the eventual data session. Without
+    forwarding, an operator who set ``trust_server_heartbeat=True``
+    finds the *first* round-trip — leader discovery — running with
+    the default opt-out, defeating the very setting they enabled.
+    Likewise ``max_total_rows`` / ``max_continuation_frames`` matter
+    for admin paths (``cluster_info`` / ``dump``) reachable through
+    the resolved client.
+
     Wraps the seed in a single-node :class:`MemoryNodeStore` and
     delegates to :meth:`ClusterClient.find_leader`. Returns the
     leader's address on success; raises the underlying
     ``ClusterError`` / ``ClusterPolicyError`` for the surrounding
     error-translation arms in :func:`_build_and_connect` to handle.
     """
-    store = MemoryNodeStore([address])
-    cluster = ClusterClient(store, timeout=timeout)
+    cluster = _get_resolve_leader_cluster(
+        address=address,
+        timeout=timeout,
+        max_total_rows=max_total_rows,
+        max_continuation_frames=max_continuation_frames,
+        trust_server_heartbeat=trust_server_heartbeat,
+    )
     return await cluster.find_leader()
 
 
@@ -202,7 +307,13 @@ async def _build_and_connect(
     connections; the dbapi handles the redirect transparently.
     """
     try:
-        leader_address = await _resolve_leader(address, timeout=timeout)
+        leader_address = await _resolve_leader(
+            address,
+            timeout=timeout,
+            max_total_rows=max_total_rows,
+            max_continuation_frames=max_continuation_frames,
+            trust_server_heartbeat=trust_server_heartbeat,
+        )
     except _client_exc.ClusterPolicyError as e:
         # Operator allowlist rejected a redirect target. Surface as
         # InterfaceError with the canonical prefix — symmetric with

tests/conftest.py

@@ -1,8 +1,27 @@
 """Pytest configuration for dqlite-dbapi tests."""
 
 import sys
+from collections.abc import Iterator
 from pathlib import Path
 
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def _clear_resolve_leader_cache() -> Iterator[None]:
+    """Clear the process-wide ``_resolve_leader`` ClusterClient cache
+    between tests. The cache is keyed by (address, governors) so two
+    tests that mock ``ClusterClient`` against the same seed address
+    would otherwise share a stale cached instance from the first
+    test's patch context.
+    """
+    from dqlitedbapi import connection as _conn_mod
+
+    _conn_mod._RESOLVE_LEADER_CACHE.clear()
+    yield
+    _conn_mod._RESOLVE_LEADER_CACHE.clear()
+
+
 # Add python-dqlite-dev's testlib to sys.path so tests (in particular
 # the leader-redirect integration suite) can import shared utilities
 # from ``dqlitetestlib``. ``python-dqlite-dev`` is expected as a

tests/test_async_close_race.py

@@ -34,7 +34,7 @@ async def fake_query_raw_typed(_sql: str, _params: object) -> tuple[list, list,
     with (
         patch(
             "dqlitedbapi.connection._resolve_leader",
-            new=AsyncMock(side_effect=lambda address, *, timeout: address),
+            new=AsyncMock(side_effect=lambda address, *, timeout, **_kw: address),
         ),
         patch("dqlitedbapi.connection.DqliteConnection") as MockDqliteConn,
     ):

tests/test_close_timeout_plumbing.py

@@ -83,7 +83,7 @@ def test_propagates_to_underlying_dqlite_connection(self) -> None:
         with (
             patch(
                 "dqlitedbapi.connection._resolve_leader",
-                new=AsyncMock(side_effect=lambda address, *, timeout: address),
+                new=AsyncMock(side_effect=lambda address, *, timeout, **_kw: address),
             ),
             patch("dqlitedbapi.connection.DqliteConnection") as MockConn,
         ):

tests/test_connection_closed_checks.py

@@ -195,7 +195,7 @@ async def _raise_policy(*args: object, **kwargs: object) -> None:
         # ``DqliteConnection.connect`` arm. Then patch
         # ``DqliteConnection.connect`` to raise the policy error this
         # test pins.
-        async def _identity_resolve(address: str, *, timeout: float) -> str:
+        async def _identity_resolve(address: str, *, timeout: float, **_kw: object) -> str:
             return address
 
         monkeypatch.setattr(

tests/test_leader_redirect_on_connect.py

@@ -222,6 +222,40 @@ async def test_build_and_connect_mid_flip_leader_change_propagates() -> None:
             )
 
 
+@pytest.mark.asyncio
+async def test_resolve_leader_threads_governors_to_cluster_client() -> None:
+    """``_resolve_leader`` must forward ``trust_server_heartbeat`` /
+    ``max_total_rows`` / ``max_continuation_frames`` to the
+    ``ClusterClient`` it constructs. Without forwarding, the FIRST
+    round-trip — leader discovery — runs with default governors
+    even when the operator opted into different ones, defeating
+    the per-connection setting at exactly the path the security
+    audit flagged.
+    """
+    captured: dict[str, object] = {}
+
+    def fake_cluster_client(store: object, **kwargs: object) -> object:
+        captured.update(kwargs)
+        client = MagicMock()
+        client.find_leader = AsyncMock(return_value="leader:9999")
+        return client
+
+    with patch("dqlitedbapi.connection.ClusterClient", fake_cluster_client):
+        result = await _resolve_leader(
+            "seed:9001",
+            timeout=5.0,
+            max_total_rows=None,
+            max_continuation_frames=42,
+            trust_server_heartbeat=True,
+        )
+
+    assert result == "leader:9999"
+    assert captured["timeout"] == 5.0
+    assert captured["max_total_rows"] is None
+    assert captured["max_continuation_frames"] == 42
+    assert captured["trust_server_heartbeat"] is True
+
+
 # --- end-to-end via Connection (sync surface) ---
 
 
@@ -231,7 +265,7 @@ def test_connection_connect_uses_leader_address(monkeypatch: pytest.MonkeyPatch)
     inner ``DqliteConnection``."""
     captured_addresses: list[str] = []
 
-    async def fake_resolve(address: str, *, timeout: float) -> str:
+    async def fake_resolve(address: str, *, timeout: float, **_kw: object) -> str:
         return "leader:9999"
 
     def capture_dqlite_connection(address: str, *args: object, **kwargs: object) -> AsyncMock:

tests/test_max_total_rows_plumbing.py

@@ -47,7 +47,9 @@ def test_propagates_to_underlying_dqlite_connection(self) -> None:
         with (
             patch(
                 "dqlitedbapi.connection._resolve_leader",
-                new=AsyncMock(side_effect=lambda address, *, timeout: address),
+                new=AsyncMock(
+                    side_effect=lambda address, *, timeout, **_kw: address,
+                ),
             ),
             patch("dqlitedbapi.connection.DqliteConnection") as MockConn,
         ):

tests/test_resolve_leader_cluster_cache.py

@@ -0,0 +1,131 @@
+"""``_resolve_leader`` shares a process-wide ``ClusterClient`` per
+``(address, governors)`` tuple.
+
+Without sharing, every dbapi ``connect()`` / SA pool slot warm-up
+constructs a fresh ``ClusterClient`` and discards it on return —
+defeating the single-flight ``_find_leader_tasks`` collapse and the
+``_last_known_leader`` fast-path that ``ClusterClient`` already
+implements. Under N concurrent SA pool checkouts after a leader flip,
+the cluster sees N independent leader-discovery sweeps where one
+would suffice.
+
+This test pins the per-key reuse and the per-governor isolation, and
+validates wholesale invalidation on the fork-pid token used by the
+underlying ``DqliteConnection`` fork-safety machinery.
+"""
+
+import os
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from dqliteclient import connection as _client_conn_mod
+from dqlitedbapi import connection as _conn_mod
+from dqlitedbapi.connection import _resolve_leader
+
+
+@pytest.mark.asyncio
+async def test_resolve_leader_reuses_cluster_client_for_same_key() -> None:
+    """Two ``_resolve_leader`` calls with the same address+governors
+    must share a single ``ClusterClient`` instance — that is what
+    keeps the leader-tracker fast-path effective at the dbapi layer."""
+    construct_count = 0
+
+    def fake_cluster_client(_store: object, **_kwargs: object) -> MagicMock:
+        nonlocal construct_count
+        construct_count += 1
+        client = MagicMock()
+        client.find_leader = AsyncMock(return_value="leader:9999")
+        return client
+
+    with patch("dqlitedbapi.connection.ClusterClient", fake_cluster_client):
+        await _resolve_leader("seed:9001", timeout=5.0)
+        await _resolve_leader("seed:9001", timeout=5.0)
+        await _resolve_leader("seed:9001", timeout=5.0)
+
+    # Three calls but only ONE ClusterClient construction.
+    assert construct_count == 1
+
+
+@pytest.mark.asyncio
+async def test_resolve_leader_isolates_distinct_governors() -> None:
+    """Different governor tuples must produce distinct
+    ClusterClient instances — sharing would cross-contaminate
+    the trust_server_heartbeat opt-in."""
+    constructed_kwargs: list[dict[str, object]] = []
+
+    def fake_cluster_client(_store: object, **kwargs: object) -> MagicMock:
+        constructed_kwargs.append(kwargs)
+        client = MagicMock()
+        client.find_leader = AsyncMock(return_value="leader:9999")
+        return client
+
+    with patch("dqlitedbapi.connection.ClusterClient", fake_cluster_client):
+        await _resolve_leader("seed:9001", timeout=5.0, trust_server_heartbeat=False)
+        await _resolve_leader("seed:9001", timeout=5.0, trust_server_heartbeat=True)
+        # Repeat each — cache hits.
+        await _resolve_leader("seed:9001", timeout=5.0, trust_server_heartbeat=False)
+        await _resolve_leader("seed:9001", timeout=5.0, trust_server_heartbeat=True)
+
+    # Two distinct keys → two constructions; the repeats are cached.
+    assert len(constructed_kwargs) == 2
+    heartbeat_settings = {kw["trust_server_heartbeat"] for kw in constructed_kwargs}
+    assert heartbeat_settings == {False, True}
+
+
+@pytest.mark.asyncio
+async def test_resolve_leader_cache_invalidates_on_fork_pid_change() -> None:
+    """Fork in a child process must wholesale-clear the cache:
+    a ClusterClient inherited from the parent carries
+    parent-allocated asyncio.Lock / Task references that the
+    child's event loop cannot make progress on. The fork-pid
+    token that DqliteConnection uses for fork-safety is the
+    sentinel."""
+    construct_count = 0
+
+    def fake_cluster_client(_store: object, **_kwargs: object) -> MagicMock:
+        nonlocal construct_count
+        construct_count += 1
+        client = MagicMock()
+        client.find_leader = AsyncMock(return_value="leader:9999")
+        return client
+
+    with patch("dqlitedbapi.connection.ClusterClient", fake_cluster_client):
+        await _resolve_leader("seed:9001", timeout=5.0)
+        # Simulate fork: the client-side _current_pid is what
+        # _refresh_pid_cache writes after_in_child. Bump it to a
+        # value that cannot collide with the cache's recorded pid.
+        with patch.object(_client_conn_mod, "_current_pid", os.getpid() + 1):
+            await _resolve_leader("seed:9001", timeout=5.0)
+
+    # Two constructions: one pre-fork, one post-fork.
+    assert construct_count == 2
+
+
+@pytest.mark.asyncio
+async def test_resolve_leader_cache_evicts_at_max_size() -> None:
+    """LRU-ish eviction: a cache that grew past the cap must drop
+    its oldest entry rather than leak unbounded under adversarial
+    governor-fragmentation."""
+    construct_count = 0
+
+    def fake_cluster_client(_store: object, **_kwargs: object) -> MagicMock:
+        nonlocal construct_count
+        construct_count += 1
+        client = MagicMock()
+        client.find_leader = AsyncMock(return_value="leader:9999")
+        return client
+
+    cap = _conn_mod._RESOLVE_LEADER_CACHE_MAX
+    with patch("dqlitedbapi.connection.ClusterClient", fake_cluster_client):
+        # Fill to cap with distinct timeouts.
+        for i in range(cap):
+            await _resolve_leader("seed:9001", timeout=float(i + 1))
+        assert len(_conn_mod._RESOLVE_LEADER_CACHE) == cap
+
+        # One more entry must NOT push the cache past the cap.
+        await _resolve_leader("seed:9001", timeout=float(cap + 100))
+        assert len(_conn_mod._RESOLVE_LEADER_CACHE) == cap
+
+    # Exactly cap+1 constructions occurred.
+    assert construct_count == cap + 1