From 51d106428ddd544a52dda36f0e7080076f063f8d Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Wed, 29 Jul 2026 23:17:41 +0300 Subject: [PATCH] test: guard test_idle_heartbeat against shard connection replacement test_idle_heartbeat snapshots each connection's request_ids keyed by id(connection), sleeps for a couple of heartbeat intervals, then looks the connections back up by id() to validate the heartbeats. If a connection object is replaced in between, the lookup raises KeyError. master already fixed the main source of this: wait_for_all_pools() only waits for the first connection per host, so shard-aware connections to the remaining shards could still be opening (and replacing placeholders) during the sleep window. 454f72742 and f34863730 added _wait_for_all_shard_connections(), called before the snapshot is taken, and deliberately dropped an earlier "skip unknown connections" band-aid so the assertions stay meaningful once the pool is known to be stable. That fix only guarantees connection *count* has stabilized, though. It doesn't rule out an existing connection being swapped for a new object in the same shard slot while the count stays constant: - HostConnection._open_connection_to_missing_shard() replaces a shard's connection once orphaned_threshold_reached, independent of pool size. - HostConnection.return_connection() -> _replace() does the same when a connection is found defunct or closed. Both can fire at any time, including during this test's own sleep, and are more likely to fire under CPU/IO contention (e.g. other tests running concurrently delaying heartbeat responses) - which lines up with the "why does this pop now, and only intermittently" puzzle from the #653 review discussion (with @dkropachev and @fruch) about this same test's flakiness. Guard against the resulting KeyError by filtering to connections that are still present in the snapshot, but without silently weakening the test: in the common case (no replacement) the filtered list is identical to the full connection list, so behavior is unchanged from master. Only when the filter actually drops something do we log a warning (so a recurrence remains visible in CI) and we still hard-fail if every single connection was replaced, since at that point nothing was validated. Compare both directions of the snapshot/current intersection, not just "every current connection is known". return_connection() has a genuine asynchronous gap: it pops a defunct/closed connection out of the pool's dict synchronously, but the replacement is inserted later by _replace(), which runs on another thread via session.submit(). A snapshot taken during that gap sees fewer current connections than were recorded, and every one of them is still "known" (nothing new has landed yet) - so "all current connections are known" is trivially true even though fewer connections than snapshotted are being validated. Tracking missing-from-current (snapshotted connections no longer present) in addition to unknown-in-current (new connections not in the snapshot) makes that partial-removal case surface as a warning too, per review feedback from @copilot-pull-request-reviewer. Signed-off-by: Yaniv Kaul --- tests/integration/standard/test_cluster.py | 48 ++++++++++++++++++++-- 1 file changed, 44 insertions(+), 4 deletions(-) diff --git a/tests/integration/standard/test_cluster.py b/tests/integration/standard/test_cluster.py index 9db4fede9e..5080d2f3f3 100644 --- a/tests/integration/standard/test_cluster.py +++ b/tests/integration/standard/test_cluster.py @@ -783,15 +783,55 @@ def test_idle_heartbeat(self): connections = [c for holders in cluster.get_connection_holders() for c in holders.get_connections()] - # make sure requests were sent on all connections - for c in connections: + # _wait_for_all_shard_connections() above prevents the common KeyError caused by + # shard connections still being opened during pool warm-up. It does not cover a + # narrower, still-possible race: HostConnection can swap an existing connection + # for a new object in the same shard slot independently of the connection count, + # e.g. via _open_connection_to_missing_shard() once orphaned_threshold_reached, or + # via return_connection()/_replace() when a connection is defunct/closed. The + # latter has a genuine asynchronous gap: return_connection() pops the old + # connection out of the pool's dict synchronously, but _replace() (which inserts + # the replacement) runs later on another thread via session.submit(). A snapshot + # taken during that gap sees fewer *current* connections than were recorded, and + # every one of them is still "known" (nothing new has been inserted yet) - so + # checking only "is every current connection known" can't distinguish "fewer + # connections than snapshotted" from "same connections, none replaced". Compare + # both directions of the snapshot/current intersection so partial removals (as + # well as outright replacements) stay visible, and log loudly if either side + # actually differs so a recurrence stays visible in CI instead of being silently + # absorbed. + current_ids = {id(c) for c in connections} + snapshot_ids = set(connection_request_ids) + # snapshotted connections no longer present at all (removed, replacement may + # not have landed yet) + missing_from_current = snapshot_ids - current_ids + # connections present now that weren't in the snapshot (replacement landed) + unknown_in_current = current_ids - snapshot_ids + known_connections = [c for c in connections if id(c) in connection_request_ids] + if missing_from_current or unknown_in_current: + log.warning( + "test_idle_heartbeat: connections changed between the snapshot and " + "validation (%d snapshotted connections missing, %d new connections " + "observed, out of %d snapshotted / %d current); validating only the " + "%d connections common to both", + len(missing_from_current), len(unknown_in_current), + len(snapshot_ids), len(current_ids), len(known_connections) + ) + assert len(known_connections) > 0, ( + "All connections were replaced during the test; " + "no heartbeats could be validated" + ) + + # make sure heartbeat requests were sent on all known connections + for c in known_connections: expected_ids = connection_request_ids[id(c)] expected_ids.rotate(-1) with c.lock: assertListEqual(list(c.request_ids), list(expected_ids)) - # assert idle status - assert all(c.is_idle for c in connections) + # assert idle status on known connections only (replaced connections + # may not have had their idle state set yet) + assert all(c.is_idle for c in known_connections) # send enough messages to ensure all connections are used # (with shard-aware routing, each query only hits one shard per host,