Skip to content

(fix) test: Fix KeyError in test_idle_heartbeat with connection repla… #762

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mykaul wants to merge 1 commit into
base: master
Choose a base branch
from
Open

(fix) test: Fix KeyError in test_idle_heartbeat with connection repla… #762

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 44 additions & 4 deletions tests/integration/standard/test_cluster.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Comment on lines +803 to +810

@coderabbitai coderabbitai Bot Jul 31, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

python - <<'PY'
import gc

class Connection:
    pass

snapshot = Connection()
snapshot_id = id(snapshot)
del snapshot
gc.collect()

for _ in range(10000):
    replacement = Connection()
    if id(replacement) == snapshot_id:
        print("id reuse observed")
        break
    del replacement
else:
    raise SystemExit("id reuse was not observed in this run")
PY

Repository: scylladb/python-driver

Length of output: 198

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- symbol references ---'
rg -n -C 8 'connection_request_ids|known_connections|missing_from_current|unknown_in_current|test_idle_heartbeat' tests/integration/standard/test_cluster.py

printf '%s\n' '--- relevant structure ---'
ast-grep outline tests/integration/standard/test_cluster.py --match 'def test_idle_heartbeat' --view expanded
sed -n '770,845p' tests/integration/standard/test_cluster.py

Repository: scylladb/python-driver

Length of output: 11204

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- connection holder and replacement paths ---'
rg -n -C 6 '_open_connection_to_missing_shard|return_connection|_replace|defunct|remove|_connections|get_connections' cassandra tests/integration/standard/test_cluster.py

printf '%s\n' '--- connection class definitions and references ---'
rg -n -C 4 'class (HostConnection|Connection)|Connection\(' cassandra | head -n 240

printf '%s\n' '--- deterministic ID-reuse probes ---'
python3 - <<'PY'
import gc

class Connection:
    pass

def probe(factory, name):
    first = factory()
    first_id = id(first)
    del first
    gc.collect()
    for i in range(100000):
        replacement = factory()
        if id(replacement) == first_id:
            print(name, "id reuse observed at iteration", i)
            return
        del replacement
    print(name, "id reuse not observed")

probe(object, "object")
probe(Connection, "Connection")
PY

Repository: scylladb/python-driver

Length of output: 50379

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- targeted pool symbols ---'
rg -n 'class HostConnection|def return_connection|def _replace|def _open_connection_to_missing_shard|_connections\.pop|_connections\[' cassandra/pool.py cassandra/cluster.py

printf '%s\n' '--- probe output ---'
tail -n 40 /tmp/coderabbit-shell-logs/shell-output-ul1sAj

Repository: scylladb/python-driver

Length of output: 3321

🏁 Script executed:

#!/bin/bash
set -euo pipefail

sed -n '373,450p' cassandra/pool.py
sed -n '530,625p' cassandra/pool.py
sed -n '690,715p' cassandra/pool.py
sed -n '748,785p' cassandra/pool.py

Repository: scylladb/python-driver

Length of output: 11129

Preserve snapshot connection objects during reconciliation.

Because connection_request_ids stores only integer IDs, Python can reuse an ID after a removed connection is collected. Store strong references in snapshot_connections and compare connections with is when computing missing_from_current, unknown_in_current, and known_connections.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/integration/standard/test_cluster.py` around lines 803 - 810, The
reconciliation logic around snapshot_connections must retain strong references
to the original connection objects instead of relying on integer IDs, preventing
ID reuse from misidentifying replacements. Update missing_from_current,
unknown_in_current, and known_connections to compare connections by identity
using is while preserving the existing snapshot and replacement behavior.

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"
)
Comment on lines +820 to +823

@coderabbitai coderabbitai Bot Jul 31, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Make the failure message match the failure condition.

known_connections is empty when no snapshotted connection remains in the current list. The asynchronous removal-to-insertion gap can cause this without proving that every connection was replaced. Report the observed condition directly.

-            "All connections were replaced during the test; "
+            "No snapshotted connections remained in the current pools; "
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
assert len(known_connections) > 0, (
"All connections were replaced during the test; "
"no heartbeats could be validated"
)
assert len(known_connections) > 0, (
"No snapshotted connections remained in the current pools; "
"no heartbeats could be validated"
)
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/integration/standard/test_cluster.py` around lines 820 - 823, Update
the assertion message for known_connections so it accurately states that no
snapshotted connection remains in the current connection list, rather than
claiming all connections were replaced. Keep the existing len(known_connections)
> 0 failure condition unchanged.


# 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,
Expand Down
Loading