From def3f2b72b37d52c6fdbbd39a9060be868bebef5 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Sun, 4 Oct 2026 11:59:22 +0300 Subject: [PATCH] feat: shuffle the bulk connect order per connection, not per factory (#35) --- CONTEXT.md | 10 ++-- db_retry/connections.py | 46 +++++++------------ tests/test_connection_factory.py | 79 ++++++++++++++++++++------------ 3 files changed, 71 insertions(+), 64 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index b622acf..4e8806b 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -28,8 +28,8 @@ _Avoid_: retry, as a count. The knob is spelled `retries` for callers and cannot breaking them, but every number in this package counts attempts. **Primary host**: -The host a `ConnectionPlan`'s first connect attempt is aimed at — which, for a multi-host DSN, is -*every* host in shuffled order, handed to asyncpg to walk itself. "Primary" is a position in the -two-stage connect (one bulk attempt, then host-by-host through `failover`), not a PostgreSQL -replication role; replication role is `target_session_attrs`, where `read-write` selects a writable -node and `prefer-standby` a replica. +The host a connection's first connect attempt is aimed at — which, for a multi-host DSN, is +*every* host, in an order shuffled afresh for each connection, handed to asyncpg to walk itself. +"Primary" is a position in the two-stage connect (one bulk attempt, then host-by-host through +`failover`), not a PostgreSQL replication role; replication role is `target_session_attrs`, where +`read-write` selects a writable node and `prefer-standby` a replica. diff --git a/db_retry/connections.py b/db_retry/connections.py index de1c9ea..8111dc2 100644 --- a/db_retry/connections.py +++ b/db_retry/connections.py @@ -22,8 +22,6 @@ class ConnectionPlan: connect_args: typing.Mapping[str, typing.Any] target_session_attrs: SessionAttribute | None - primary_host: str | list[str] - primary_port: int | list[int] | None failover: tuple[tuple[str, int], ...] @@ -42,45 +40,29 @@ def build_connection_plan(url: sqlalchemy.URL) -> ConnectionPlan: target_session_attrs: SessionAttribute | None = ( SessionAttribute(raw_target_session_attrs) if raw_target_session_attrs else None ) - raw_hosts: str | list[str] = connect_args.pop("host") - raw_ports: int | list[int] | None = connect_args.pop("port", None) - primary_host: str | list[str] - primary_port: int | list[int] | None - failover: tuple[tuple[str, int], ...] if hosts_and_ports: - random.shuffle(hosts_and_ports) - primary_host = list(map(itemgetter(0), hosts_and_ports)) - primary_port = list(map(itemgetter(1), hosts_and_ports)) - failover = tuple(hosts_and_ports) - else: - primary_host = raw_hosts - primary_port = raw_ports - failover = () + del connect_args["host"], connect_args["port"] return ConnectionPlan( connect_args=types.MappingProxyType(connect_args), target_session_attrs=target_session_attrs, - primary_host=primary_host, - primary_port=primary_port, - failover=failover, + failover=tuple(hosts_and_ports), ) async def _connect( plan: ConnectionPlan, - host: str | list[str], - port: int | list[int] | None, timeout: float, # noqa: ASYNC109 + **address: str | int | list[str] | list[int], ) -> "ConnectionType": return await asyncpg.connect( **plan.connect_args, - host=host, - port=port, + **address, timeout=timeout, target_session_attrs=plan.target_session_attrs, ) -def _reshuffled(failover: tuple[tuple[str, int], ...]) -> list[tuple[str, int]]: +def _shuffled(failover: tuple[tuple[str, int], ...]) -> list[tuple[str, int]]: return random.sample(failover, len(failover)) @@ -91,17 +73,23 @@ def build_connection_factory( plan: typing.Final = build_connection_plan(url) async def _connection_factory() -> "ConnectionType": + if not plan.failover: + return await _connect(plan, timeout) + + bulk_order: typing.Final = _shuffled(plan.failover) try: - return await _connect(plan, plan.primary_host, plan.primary_port, timeout) + return await _connect( + plan, + timeout, + host=list(map(itemgetter(0), bulk_order)), + port=list(map(itemgetter(1), bulk_order)), + ) except TimeoutError: - if not plan.failover: - raise - logger.warning("Failed to fetch asyncpg connection. Trying host by host.") - for host, port in _reshuffled(plan.failover): + for host, port in _shuffled(plan.failover): try: - return await _connect(plan, host, port, timeout) + return await _connect(plan, timeout, host=host, port=port) except (TimeoutError, OSError, asyncpg.TargetServerAttributeNotMatched) as exc: logger.warning("Failed to fetch asyncpg connection from %s, %s", host, exc) msg: typing.Final = f"None of the hosts match the target attribute requirement {plan.target_session_attrs}" diff --git a/tests/test_connection_factory.py b/tests/test_connection_factory.py index ab4bbbe..600b8ce 100644 --- a/tests/test_connection_factory.py +++ b/tests/test_connection_factory.py @@ -1,4 +1,5 @@ import os +import random import typing from unittest import mock @@ -70,44 +71,64 @@ async def test_connection_factory_failure_and_success(monkeypatch: pytest.Monkey assert result is mock_connection +def _six_hosts() -> tuple[list[tuple[str, int]], sqlalchemy.URL]: + hosts_and_ports: typing.Final = [(f"host{n}", 5431 + n) for n in range(1, 7)] + query: typing.Final = "&".join(f"host={host}:{port}" for host, port in hosts_and_ports) + return hosts_and_ports, sqlalchemy.make_url(f"postgresql+asyncpg://user:password@/database?{query}") + + def test_build_connection_plan_multihost() -> None: url: typing.Final = sqlalchemy.make_url( - "postgresql+asyncpg://user:password@/database?host=host1:5432&host=host2:5432&target_session_attrs=read-write" + "postgresql+asyncpg://user:password@/database?host=host1:5432&host=host2:5433&target_session_attrs=read-write" ) plan: typing.Final[ConnectionPlan] = build_connection_plan(url) - assert set(plan.failover) == {("host1", 5432), ("host2", 5432)} - assert isinstance(plan.primary_host, list) - assert isinstance(plan.primary_port, list) + assert plan.failover == (("host1", 5432), ("host2", 5433)) assert plan.target_session_attrs == SessionAttribute("read-write") assert "host" not in plan.connect_args assert "port" not in plan.connect_args assert "target_session_attrs" not in plan.connect_args -def test_host_and_port_stay_paired_through_the_shuffle() -> None: - """INVARIANT: no plan ever pairs one host's name with another host's port. +def test_build_connection_plan_is_deterministic() -> None: + hosts_and_ports, url = _six_hosts() + assert build_connection_plan(url) == build_connection_plan(url) + assert build_connection_plan(url).failover == tuple(hosts_and_ports) - Broken by shuffling ``primary_host`` and ``primary_port`` as two independent lists, or by - deriving ``failover`` from a second shuffle of its own rather than from the one that produced - the primary order. Both read as tidier code and both silently mis-pair. Nothing downstream can - catch it: ``_connect`` hands whatever pair it is given straight to asyncpg, so a swap surfaces - as a refused connection that is indistinguishable from a host being down -- on the failover - path, which by definition only runs when hosts are already failing. This is also why the DSN - here gives each host a distinct port; with matching ports a swap is unobservable. Six hosts - rather than two for the same reason: the assertions are order-independent, so a mis-pairing - build can still be let through by a shuffle that happens to come out in step, and six hosts put - that at one run in 720 instead of one in two. + +async def test_each_connection_gets_a_fresh_bulk_order(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("db_retry.connections.random", random.Random(0)) # noqa: S311 + connect: typing.Final = mock.AsyncMock(return_value=mock.AsyncMock(spec=asyncpg.Connection)) + monkeypatch.setattr("asyncpg.connect", connect) + _, url = _six_hosts() + factory: typing.Final = build_connection_factory(url=url, timeout=1.0) + await factory() + await factory() + first, second = (call.kwargs["host"] for call in connect.await_args_list) + assert sorted(first) == sorted(second) + assert first != second + + +async def test_host_and_port_stay_paired_through_the_shuffle(monkeypatch: pytest.MonkeyPatch) -> None: + """INVARIANT: no connect attempt ever pairs one host's name with another host's port. + + Broken by shuffling the bulk attempt's hosts and ports as two independent lists, or by drawing + them from different shuffles. Both read as tidier code and both silently mis-pair. Nothing + downstream can catch it: ``_connect`` hands whatever it is given straight to asyncpg, so a swap + surfaces as a refused connection that is indistinguishable from a host being down -- on the + failover path, which by definition only runs when hosts are already failing. This is also why + each host here has a distinct port; with matching ports a swap is unobservable. Six hosts + rather than two for the same reason: a mis-pairing build can still be let through by shuffles + that happen to come out in step, and six hosts put that at one run in 720 instead of one in two. """ - expected: typing.Final = {(f"host{n}", 5431 + n) for n in range(1, 7)} - hosts: typing.Final = "&".join(f"host={host}:{port}" for host, port in sorted(expected)) - plan: typing.Final[ConnectionPlan] = build_connection_plan( - sqlalchemy.make_url(f"postgresql+asyncpg://user:password@/database?{hosts}") - ) - assert isinstance(plan.primary_host, list) - assert isinstance(plan.primary_port, list) - assert set(plan.failover) == expected - assert set(zip(plan.primary_host, plan.primary_port, strict=True)) == expected - assert list(zip(plan.primary_host, plan.primary_port, strict=True)) == list(plan.failover) + connect: typing.Final = mock.AsyncMock(side_effect=TimeoutError) + monkeypatch.setattr("asyncpg.connect", connect) + hosts_and_ports, url = _six_hosts() + factory: typing.Final = build_connection_factory(url=url, timeout=1.0) + with pytest.raises(asyncpg.TargetServerAttributeNotMatched): + await factory() + bulk, *failover = connect.await_args_list + assert sorted(zip(bulk.kwargs["host"], bulk.kwargs["port"], strict=True)) == hosts_and_ports + assert sorted((call.kwargs["host"], call.kwargs["port"]) for call in failover) == hosts_and_ports def test_build_connection_plan_connect_args_is_read_only() -> None: @@ -122,9 +143,7 @@ def test_build_connection_plan_single_host() -> None: url: typing.Final = sqlalchemy.make_url(f"postgresql+asyncpg://user:password@host1:{port}/database") plan: typing.Final[ConnectionPlan] = build_connection_plan(url) assert plan.failover == () - assert plan.primary_host == "host1" - assert plan.primary_port == port + assert plan.connect_args["host"] == "host1" + assert plan.connect_args["port"] == port assert plan.target_session_attrs is None - assert "host" not in plan.connect_args - assert "port" not in plan.connect_args assert "target_session_attrs" not in plan.connect_args