connection: tolerate clean startup close

[dkropachev]

Summary

Fixes #942.

Version 3.29.11 started raising ConnectionShutdown when the server closed a socket cleanly during startup. Scylla maintenance mode deliberately has this socket shape for regular CQL: it accepts the connection and initial CQL frame, then closes because regular CQL is unavailable while the maintenance socket remains usable.

This change restores the agreed factory contract at both layers: Connection.factory() and Cluster.connection_factory() return the closed connection after a clean startup close. Actual connection owners validate that result before using it:

• control and pool creation reject it without adopting a dead socket
• host reconnection parks without retrying or falsely calling on_up
• a status UP event or a successful control reconnect to that endpoint resumes the parked host
• replacement recovery performs an authoritative host-down handoff, avoiding both a tight replacement loop and an empty pool stranded by custom conviction or another session's open-pool discount
• missing-shard attempts discard the closed result and always release their in-progress marker
• prepared-statement warming stops before issuing work on the closed connection

The existing pending-connection ownership changes remain: pool-owned connections stay tracked during startup and replacement keyspace setup so shutdown can close them promptly. The Twisted reactor tracks and cancels pending endpoint connects, wakes the startup waiter during shutdown, rejects a late connectionMade callback, and propagates genuine endpoint-connect failures immediately instead of waiting for the startup timeout.

Reproducer and coverage

The maintenance-mode integration test uses a real local TCP listener that accepts a CQL socket, reads the initial OPTIONS frame, and closes cleanly during startup. It verifies that both factory layers return the closed connection while Cluster.connect() still reports the host as unavailable.

Additional unit coverage verifies owner-side rejection, parked host reconnection, control-reconnect recovery, replacement handoff (including shutdown and open-pool-discount races), missing-shard cleanup for both endpoint paths, pending connection ownership, and Twisted cancellation/failure handling.

Driver surface

Touches connection startup behavior, cluster connection owners, host/pool recovery, shard-aware connection creation, and Twisted pending-connect cancellation. Protocol serialization and normal query execution are unchanged.

Compatibility / risk

The observable compatibility restoration is that callers of either factory layer can inspect a cleanly closed startup result, matching behavior before 3.29.11. Cluster-managed owners still require a serviceable connection and now enforce that requirement at the ownership boundary instead of changing the factory return contract.

Tests

• TZ=UTC uv run pytest -q tests/unit — 722 passed, 79 skipped
• CASSANDRA_VERSION=3.11.4 MAPPED_CASSANDRA_VERSION=3.11.4 uv run pytest -q tests/integration/standard/test_maintenance_mode_connection.py — 2 passed

[coderabbitai]

📝 Walkthrough

Walkthrough

This change hardens connection startup and shutdown, host-transition ordering, pool ownership, replacement recovery, session keyspace propagation, control-connection identity handling, endpoint relocation, and request fallback. It adds epoch and generation fencing, serialized callbacks, stale-attempt cleanup, compatibility handling for legacy hooks, and extensive unit and integration coverage for lifecycle races and maintenance-mode startup closures.

Sequence Diagram(s)

sequenceDiagram
  participant Session
  participant HostConnection
  participant Connection
  participant ControlConnection
  Session->>HostConnection: borrow or create connection
  HostConnection->>Connection: start, validate, and configure keyspace
  Connection-->>HostConnection: return usable or startup-closed result
  HostConnection-->>Session: publish or discard pool connection
  ControlConnection->>Session: refresh topology and host identity
  Session-->>ControlConnection: schedule recovery or endpoint relocation

Loading

Possibly related PRs

• scylladb/python-driver#637 — Both modify control-connection topology handling and related tests.
• scylladb/python-driver#936 — Both modify HostConnection and shard-aware endpoint handling.
• scylladb/python-driver#947 — Both modify prepared-statement keyspace metadata and fallback behavior.

Suggested labels: area/Driver_-_python-driver

Suggested reviewers: lorak-mmk, mykaul

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 13.60% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the primary change: tolerating a clean startup connection close.
Description check	✅ Passed	The description thoroughly summarizes the fix, rationale, testing, compatibility, and linked issue, despite not reproducing the checklist headings.
Linked Issues check	✅ Passed	The changes address issue #942 by restoring maintenance-mode startup behavior and preserving local CQL access, with targeted integration and unit coverage.
Out of Scope Changes check	✅ Passed	The implementation and tests remain focused on clean startup closes, connection ownership, recovery, pending connections, and maintenance-mode compatibility.

---

_{Comment @coderabbitai help to get the list of available commands.}

[Copilot]

Pull request overview

Restores pre-3.29.11 handling of clean server-side closes during connection startup.

Changes:

• Returns cleanly closed startup connections to their owner.
• Adds regression coverage for this behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
cassandra/connection.py	Removes immediate ConnectionShutdown for clean startup closes.
tests/unit/test_connection.py	Tests returning a cleanly closed connection.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

cassandra/connection.py:990

• Returning a closed connection is treated as a successful reconnect by existing owners. _ReconnectionHandler.run() invokes on_reconnection() and its callback for any returned value (cassandra/pool.py:281-305), and _HostReconnectionHandler then marks the host up (cassandra/pool.py:353-361); initial pool construction likewise installs the returned connection without checking is_closed (cassandra/pool.py:429-431). A clean startup close can therefore falsely mark a maintenance-mode host up and build a pool around a closed socket instead of preserving reconnection cadence. The owning paths need explicit closed-result handling (with coverage), or the factory needs a distinct failure/result contract.

        else:
            return conn

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (4)

cassandra/connection.py:975

• The opening sentence (“returns a connection once startup has completed”) conflicts with the documented behavior immediately below (direct callers may receive a closed connection when startup did not complete). Consider rewording the first sentence to reflect that the method may return a closed connection when the server cleanly closes during the handshake.

        A factory function which returns a connection once startup has
        completed, or raises an exception otherwise.

        Direct callers may receive a closed connection when the server accepts
        the socket and then cleanly closes it during the startup handshake.
        Callers that own pool startup or reconnection probes pass ``host_conn``
        or ``_raise_on_startup_close`` and receive :exc:`ConnectionShutdown`
        for that clean startup close instead.

cassandra/connection.py:998

• The exception is constructed using the input endpoint rather than the connection’s actual endpoint (conn.endpoint). Using conn.endpoint (both for message formatting and for the exception’s endpoint attribute) is typically more accurate if the connection normalizes/rewrites endpoints (e.g., via an endpoint factory or translation) and also avoids the slightly confusing % (endpoint,) tuple formatting.

        elif conn.is_closed and (host_conn is not None or raise_on_startup_close):
            raise ConnectionShutdown(
                "Connection to %s was closed during the startup handshake" % (endpoint,),
                endpoint)

tests/unit/test_connection.py:480

• server.first_frame[4] can raise IndexError if the server receives fewer than 5 bytes (e.g., if the client closes early or a timeout occurs). Adding an explicit assert len(server.first_frame) >= 5 (or parsing/validating the 9-byte header length) would make failures clearer and avoid masking the underlying cause.

            assert conn.is_closed
            assert server.received_frame.wait(2)
            assert server.error is None
            assert server.first_frame[4] == 0x05  # OPTIONS

tests/unit/test_connection.py:502

• This test asserts on the private/internal attribute _pending_connections, which can make the test brittle to refactors of internal connection tracking. If possible, assert the externally observable behavior you care about (e.g., that the factory raises and that the returned/created connection is closed) without coupling to _pending_connections’ internal representation.

            host_conn = Mock()
            host_conn.is_shutdown = False
            host_conn._pending_connections = []

            with pytest.raises(ConnectionShutdown) as exc_info:
                MaintenanceModeConnection.factory(
                    DefaultEndPoint('127.0.0.1', server.port),
                    timeout=2,
                    host_conn=host_conn)

            assert "closed during the startup handshake" in str(exc_info.value)
            assert server.received_frame.wait(2)
            assert server.error is None
            assert server.first_frame[4] == 0x05  # OPTIONS
            assert len(host_conn._pending_connections) == 1
            assert host_conn._pending_connections[0].is_closed

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (4)

cassandra/connection.py:975

• The first sentence says the factory “returns a connection once startup has completed,” but the docstring later states direct callers may receive a closed connection during the startup handshake. Please revise the opening description to reflect that the method can return a closed (non-serviceable) connection in the clean-startup-close case, so the contract is not self-contradictory.

        A factory function which returns a connection once startup has
        completed, or raises an exception otherwise.

        Direct callers may receive a closed connection when the server accepts
        the socket and then cleanly closes it during the startup handshake.
        Callers that own pool startup or reconnection probes pass ``host_conn``
        or ``_raise_on_startup_close`` and receive :exc:`ConnectionShutdown`
        for that clean startup close instead.

cassandra/connection.py:966

• _raise_on_startup_close is currently a “hidden” control flag extracted from **kwargs, which makes the factory behavior easier to call incorrectly (typos silently change behavior) and harder to discover for implementers of custom Connection subclasses. Consider making it an explicit (preferably keyword-only) parameter of factory() (defaulting to False) and documenting it alongside host_conn.

    def factory(cls, endpoint, timeout, host_conn = None, *args, **kwargs):

cassandra/connection.py:977

• _raise_on_startup_close is currently a “hidden” control flag extracted from **kwargs, which makes the factory behavior easier to call incorrectly (typos silently change behavior) and harder to discover for implementers of custom Connection subclasses. Consider making it an explicit (preferably keyword-only) parameter of factory() (defaulting to False) and documenting it alongside host_conn.

        raise_on_startup_close = kwargs.pop('_raise_on_startup_close', False)

tests/unit/test_connection.py:429

• close() uses a timed join(2) on a daemon thread and does not verify the thread actually stopped. This can leak background activity across tests and introduce intermittency under slow CI. Prefer a deterministic shutdown (e.g., signal + unblock accept/recv, then join() without a timeout, or assert not self.thread.is_alive() after joining) so failures are visible and cleanup is reliable.

            def close(self):
                self._sock.close()
                self.thread.join(2)

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

cassandra/connection.py:998

• This clean-close exception is still raised before reaching this branch for AsyncoreConnection: its close() assigns a generated ConnectionShutdown to last_error while startup is pending (cassandra/io/asyncorereactor.py:392-397), and factory() raises any last_error at line 987. Asyncore is the documented default when libev is unavailable (cassandra/cluster.py:953-956), so those users do not get the restored clean-close behavior described here. The synthetic connection in the new test leaves last_error unset and therefore misses this reactor-specific path; distinguish a clean-close-generated shutdown from a real startup error and cover Asyncore as well.

        elif conn.is_closed and (host_conn is not None or raise_on_startup_close):
            raise ConnectionShutdown(
                "Connection to %s was closed during the startup handshake" % (endpoint,),
                endpoint)

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (3)

cassandra/pool.py:614

• This call still runs while _replace() holds self._lock (line 599), but shutdown() needs that same lock before it can copy and close _pending_connections (lines 628-653). Consequently, registering the replacement as pending cannot let shutdown cancel a stalled startup; shutdown remains blocked until the connection attempt completes or times out. Move the blocking factory call outside the locked region, then reacquire the lock to install it only if the pool is still active.

                    connection = self._session.cluster.connection_factory(
                        self.host.endpoint,
                        host_conn=self,
                        on_orphaned_stream_released=self.on_orphaned_stream_released)

tests/unit/test_cluster.py:282

• This test name says a startup close is rejected, but the supplied connection has is_closed=False and the assertions verify the successful return path. Rename it to describe the open-connection case so failures and coverage are not misleading.

    def test_reconnection_factory_rejects_startup_close(self):

tests/unit/test_host_connection_pool.py:254

• No startup close occurs here: replacement_conn.is_closed is false, and the mocked cluster factory simply returns it. The test only verifies that _replace() forwards the pool as host_conn; rename it accordingly rather than claiming closed-startup rejection coverage.

    def test_replace_tracks_pending_connection_and_rejects_startup_close(self):

[Copilot]

Pull request overview

Copilot reviewed 11 out of 12 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 11 out of 12 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 11 out of 12 changed files in this pull request and generated no new comments.

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tests/unit/test_response_future.py (1)

1788-1788: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Use a Python 3.9-compatible no-log assertion.

The project supports Python 3.9, but assertNoLogs is only available from Python 3.10, causing AttributeError at both occurrences.

🤖 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/unit/test_response_future.py` at line 1788, Replace both assertNoLogs
usages in the affected tests with the project's Python 3.9-compatible no-log
assertion mechanism, preserving the cassandra.cluster logger and WARNING
threshold behavior.

🧹 Nitpick comments (3)

tests/unit/test_control_connection.py (1)

398-412: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

The modeled fallback path is never executed here.

_uses_default_failure_hooks returns True, so _signal_error takes the fenced _on_down_locked branch and self.cluster.signal_connection_failure (i.e. public_signal) is never called. The comment claims this test also proves the pre-fix fallback path would commit DOWN, but that branch is unreachable in this configuration. Consider looping over _uses_default_failure_hooks returning True/False so both stale-check paths are actually exercised.

🤖 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/unit/test_control_connection.py` around lines 398 - 412, The test setup
around public_signal only exercises the _uses_default_failure_hooks=True branch,
leaving the modeled fallback path untested. Update the test to run the scenario
with _uses_default_failure_hooks returning both True and False, ensuring the
fenced _on_down_locked path and the public signal_connection_failure path are
each executed while preserving the stale-check assertions.

tests/unit/test_response_future.py (1)

271-290: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Move callback assertions outside _set_result.

_set_result catches exceptions raised by pool.return_connection and stores them on the future, so assertions inside return_connection may not fail this test. Capture the observed state in the callback and assert it after _set_result returns.

🤖 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/unit/test_response_future.py` around lines 271 - 290, Update
test_successful_use_clears_quarantine_before_pool_return so return_connection
only captures the returned connection’s quarantine and keyspace state, without
assertions that can be swallowed by _set_result. After rf._set_result returns,
assert the captured state and returned connection, while preserving the existing
pool.return_connection call verification.

cassandra/cluster.py (1)

5277-5306: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Bind the per-iteration keyspace callback state explicitly.

callback closes over errors_returned/set_keyspace_event by name, so a late or duplicate callback from an earlier loop iteration would mutate the current iteration's state and could prematurely release set_keyspace_event.wait(...). Binding them as default arguments makes each iteration's callback self-contained (also silences Ruff B023).

♻️ Proposed fix

-                    def callback(pool, errors):
-                        errors_returned.extend(errors)
-                        set_keyspace_event.set()
+                    def callback(
+                            pool,
+                            errors,
+                            errors_returned=errors_returned,
+                            set_keyspace_event=set_keyspace_event):
+                        errors_returned.extend(errors)
+                        set_keyspace_event.set()

🤖 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 `@cassandra/cluster.py` around lines 5277 - 5306, Update the per-iteration
callback definition in the keyspace-setting flow to bind that iteration’s
errors_returned and set_keyspace_event through callback default arguments,
rather than resolving the enclosing variables at invocation time. Preserve the
existing extend and event-set behavior while ensuring late or duplicate
callbacks cannot affect another iteration’s state.

Source: Linters/SAST tools

🤖 Prompt for all review comments with 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.

Outside diff comments:
In `@tests/unit/test_response_future.py`:
- Line 1788: Replace both assertNoLogs usages in the affected tests with the
project's Python 3.9-compatible no-log assertion mechanism, preserving the
cassandra.cluster logger and WARNING threshold behavior.

---

Nitpick comments:
In `@cassandra/cluster.py`:
- Around line 5277-5306: Update the per-iteration callback definition in the
keyspace-setting flow to bind that iteration’s errors_returned and
set_keyspace_event through callback default arguments, rather than resolving the
enclosing variables at invocation time. Preserve the existing extend and
event-set behavior while ensuring late or duplicate callbacks cannot affect
another iteration’s state.

In `@tests/unit/test_control_connection.py`:
- Around line 398-412: The test setup around public_signal only exercises the
_uses_default_failure_hooks=True branch, leaving the modeled fallback path
untested. Update the test to run the scenario with _uses_default_failure_hooks
returning both True and False, ensuring the fenced _on_down_locked path and the
public signal_connection_failure path are each executed while preserving the
stale-check assertions.

In `@tests/unit/test_response_future.py`:
- Around line 271-290: Update
test_successful_use_clears_quarantine_before_pool_return so return_connection
only captures the returned connection’s quarantine and keyspace state, without
assertions that can be swallowed by _set_result. After rf._set_result returns,
assert the captured state and returned connection, while preserving the existing
pool.return_connection call verification.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: cfaf7120-c3e5-45cc-9b9f-70d14b366e77

📥 Commits

Reviewing files that changed from the base of the PR and between abfcaf8 and 0dcb688.

📒 Files selected for processing (12)

• cassandra/cluster.py
• cassandra/connection.py
• cassandra/io/twistedreactor.py
• cassandra/pool.py
• tests/integration/standard/test_maintenance_mode_connection.py
• tests/unit/io/test_twistedreactor.py
• tests/unit/test_cluster.py
• tests/unit/test_connection.py
• tests/unit/test_control_connection.py
• tests/unit/test_host_connection_pool.py
• tests/unit/test_response_future.py
• tests/unit/test_shard_aware.py

[dkropachev]

The assertNoLogs Python 3.9 issue predates this PR and is already present on master (introduced by b8b714ca3), so I’m leaving it unchanged here. It should be handled separately as a compatibility fix.

[Copilot]

Pull request overview

Copilot reviewed 11 out of 12 changed files in this pull request and generated no new comments.

[coderabbitai]

🧹 Nitpick comments (7)

cassandra/connection.py (2)

1958-1963: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Chain the converted validation error.

Without from, the traceback loses the originating ErrorMessage.

♻️ Proposed change

         except ProtocolRequestValidationException as validation_error:
-            raise validation_error.to_exception()
+            raise validation_error.to_exception() from validation_error

🤖 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 `@cassandra/connection.py` around lines 1958 - 1963, Update the exception
handling in the wait_for_response flow to chain the converted exception from
ProtocolRequestValidationException using the original validation_error as its
cause. Keep RequestValidationException propagation unchanged.

Source: Linters/SAST tools

---

1125-1145: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Extract the pending-connection fallback into helpers.

The "call _register_pending_connection else poke _pending_connections under _lock else poke it unlocked" ladder is duplicated for register and unregister inside factory. Two module-level helpers (_register_pending, _unregister_pending) would keep factory's control flow readable while preserving the third-party-pool compatibility behavior.

Also applies to: 1214-1234

🤖 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 `@cassandra/connection.py` around lines 1125 - 1145, Extract the
pending-connection compatibility ladder from factory into module-level
_register_pending and _unregister_pending helpers, covering both registration
and removal paths. Preserve the existing _register_pending_connection
preference, _lock-guarded fallback, and unlocked fallback for third-party pools,
then have factory call these helpers while retaining its shutdown and
connection-close behavior.

cassandra/pool.py (2)

778-782: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Redundant has_current_pool_fence guard in the lock list.

The not has_current_pool_fence case already returned above, so this condition is always true and only obscures the lock set.

♻️ Proposed change

-        locks = [
-            lock for lock in (cluster_lock, host_lock, session_lock)
-            if lock is not None and (
-                lock is not session_lock or has_current_pool_fence)]
+        locks = [
+            lock for lock in (cluster_lock, host_lock, session_lock)
+            if lock is not None]

🤖 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 `@cassandra/pool.py` around lines 778 - 782, Remove the redundant
has_current_pool_fence condition from the locks comprehension in the pool lock
acquisition flow, retaining only the non-None filter and existing session_lock
exclusion. Preserve the earlier return behavior for cases without a current pool
fence.

---

261-305: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

_transition_owner is write-only and lacks a class-level default.

Cluster._run_host_transition sets and clears host._transition_owner, but nothing reads it; and unlike the sibling transition fields it has no class attribute. Either drop it or add the default alongside the others.

♻️ Proposed change

     _transition_lock = None
     _transition_queue = None
     _transition_running = False
+    _transition_owner = None
     _transition_notification_queue = None

🤖 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 `@cassandra/pool.py` around lines 261 - 305, Remove the unused
_transition_owner assignment and initialization from
Cluster._run_host_transition and Host.__init__, or alternatively define a
class-level default alongside the other transition fields if the state is
required; keep the transition queue behavior unchanged.

tests/unit/test_connection.py (1)

330-342: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Bind responses explicitly in the closure.

The closure is only invoked within the same iteration, so behavior is correct today, but binding the loop variable as a default silences the linter and keeps it correct if the test is later refactored.

♻️ Proposed change

-                def send_response(message, request_id, callback):
-                    callback(responses.pop(0))
+                def send_response(
+                        message, request_id, callback,
+                        responses=responses):
+                    callback(responses.pop(0))

🤖 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/unit/test_connection.py` around lines 330 - 342, Update the
send_response closure in the test loop to bind the current responses value
explicitly through a default parameter, rather than relying on the loop-scoped
variable lookup. Preserve the existing callback behavior and response ordering
for each subTest.

Source: Linters/SAST tools

cassandra/cluster.py (1)

5563-5596: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Prefer an outer loop over tail recursion for the drain restart.

_drain_keyspace_completions re-enters itself from finally when a newly-ready head appears; a plain outer loop removes the (unlikely but unbounded) stack growth and reads more clearly.

♻️ Sketch

-    def _drain_keyspace_completions(self):
-        with self._keyspace_completion_lock:
-            if self._keyspace_completion_runner_active:
-                return
-            self._keyspace_completion_runner_active = True
-
-        rerun = False
-        try:
+    def _drain_keyspace_completions(self):
+        with self._keyspace_completion_lock:
+            if self._keyspace_completion_runner_active:
+                return
+            self._keyspace_completion_runner_active = True
+
+        while True:
+            try:
                 while True:
                     ...
-        finally:
-            with self._keyspace_completion_lock:
-                self._keyspace_completion_runner_active = False
-                rerun = ...
-            if rerun:
-                self._drain_keyspace_completions()
+            finally:
+                with self._keyspace_completion_lock:
+                    self._keyspace_completion_runner_active = False
+                    rerun = ...
+                    if rerun:
+                        self._keyspace_completion_runner_active = True
+            if not rerun:
+                return

🤖 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 `@cassandra/cluster.py` around lines 5563 - 5596, Replace the tail-recursive
restart in _drain_keyspace_completions with an outer loop that repeatedly drains
newly-ready keyspace completions. Keep the lock-protected runner-active state
and callback error handling intact, ensuring the method continues processing
when the queue head becomes dispatch-complete and ready without growing the call
stack.

tests/integration/standard/test_maintenance_mode_connection.py (1)

26-81: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Duplicated maintenance-mode fake CQL server across the two suites. Both files implement the same listener (bind ephemeral port, read a 9-byte frame, close without replying) to reproduce a maintenance-mode startup close; the shared root cause is one test double written twice.

• tests/integration/standard/test_maintenance_mode_connection.py#L26-L81: keep this as the single implementation (ideally moved to a shared test helper module so it can be imported).
• tests/unit/test_connection.py#L539-L668: drop the nested MaintenanceModeCqlServer/MaintenanceModeConnection real-socket server and reduce this test to the pending-connection ownership assertions using an in-process fake, relying on the integration test for the real-socket path.

🤖 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_maintenance_mode_connection.py` around lines
26 - 81, Deduplicate the maintenance-mode socket test double: retain the
implementation in tests/integration/standard/test_maintenance_mode_connection.py
lines 26-81, optionally relocating MaintenanceModeCqlServer to a shared test
helper for import. In tests/unit/test_connection.py lines 539-668, remove the
nested MaintenanceModeCqlServer and MaintenanceModeConnection real-socket
server, and reduce the test to pending-connection ownership assertions using an
in-process fake; the integration test remains responsible for real-socket
coverage.

🤖 Prompt for all review comments with 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.

Nitpick comments:
In `@cassandra/cluster.py`:
- Around line 5563-5596: Replace the tail-recursive restart in
_drain_keyspace_completions with an outer loop that repeatedly drains
newly-ready keyspace completions. Keep the lock-protected runner-active state
and callback error handling intact, ensuring the method continues processing
when the queue head becomes dispatch-complete and ready without growing the call
stack.

In `@cassandra/connection.py`:
- Around line 1958-1963: Update the exception handling in the wait_for_response
flow to chain the converted exception from ProtocolRequestValidationException
using the original validation_error as its cause. Keep
RequestValidationException propagation unchanged.
- Around line 1125-1145: Extract the pending-connection compatibility ladder
from factory into module-level _register_pending and _unregister_pending
helpers, covering both registration and removal paths. Preserve the existing
_register_pending_connection preference, _lock-guarded fallback, and unlocked
fallback for third-party pools, then have factory call these helpers while
retaining its shutdown and connection-close behavior.

In `@cassandra/pool.py`:
- Around line 778-782: Remove the redundant has_current_pool_fence condition
from the locks comprehension in the pool lock acquisition flow, retaining only
the non-None filter and existing session_lock exclusion. Preserve the earlier
return behavior for cases without a current pool fence.
- Around line 261-305: Remove the unused _transition_owner assignment and
initialization from Cluster._run_host_transition and Host.__init__, or
alternatively define a class-level default alongside the other transition fields
if the state is required; keep the transition queue behavior unchanged.

In `@tests/integration/standard/test_maintenance_mode_connection.py`:
- Around line 26-81: Deduplicate the maintenance-mode socket test double: retain
the implementation in
tests/integration/standard/test_maintenance_mode_connection.py lines 26-81,
optionally relocating MaintenanceModeCqlServer to a shared test helper for
import. In tests/unit/test_connection.py lines 539-668, remove the nested
MaintenanceModeCqlServer and MaintenanceModeConnection real-socket server, and
reduce the test to pending-connection ownership assertions using an in-process
fake; the integration test remains responsible for real-socket coverage.

In `@tests/unit/test_connection.py`:
- Around line 330-342: Update the send_response closure in the test loop to bind
the current responses value explicitly through a default parameter, rather than
relying on the loop-scoped variable lookup. Preserve the existing callback
behavior and response ordering for each subTest.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: a2755b6a-1f5d-4630-8574-af5e37fb539c

📥 Commits

Reviewing files that changed from the base of the PR and between 0dcb688 and 08e8726.

📒 Files selected for processing (12)

• cassandra/cluster.py
• cassandra/connection.py
• cassandra/io/twistedreactor.py
• cassandra/pool.py
• tests/integration/standard/test_maintenance_mode_connection.py
• tests/unit/io/test_twistedreactor.py
• tests/unit/test_cluster.py
• tests/unit/test_connection.py
• tests/unit/test_control_connection.py
• tests/unit/test_host_connection_pool.py
• tests/unit/test_response_future.py
• tests/unit/test_shard_aware.py