test: track TcpProxy forwarders through socket cleanup

[dkropachev]

Summary

Fixes: #962.

• keep each TcpProxy connection registered until its forwarder finishes closing both sockets
• preserve accurate active_connections reporting throughout socket cleanup
• add a deterministic regression test that blocks _close_pair() and verifies the still-running forwarder remains tracked

Scope and compatibility

This changes only the TcpProxy integration-test helper and its unit coverage. Driver runtime behavior, wire protocols, and Scylla/Cassandra compatibility are unaffected.

Validation

• uv run pytest -q tests/unit/test_tcp_proxy.py (4 passed)
• git diff --check

[coderabbitai]

📝 Walkthrough

Walkthrough

TcpProxy now uses a lifecycle lock to serialize manager-side shutdown with forwarder-owned socket closure. Forwarders close socket pairs before removing connections from the registry. Regression tests cover connection retention during blocked cleanup and prevent socket closing from overlapping connection shutdown.

Sequence Diagram(s)

sequenceDiagram
participant TcpProxy
participant Shutdown
participant Forwarder
participant SocketPair
TcpProxy->>Shutdown: acquire lifecycle lock
Shutdown->>SocketPair: shut down sockets
Forwarder->>Forwarder: wait for lifecycle lock
Forwarder->>SocketPair: close sockets
Forwarder->>TcpProxy: remove connection

Loading

Possibly related issues

• test: track TcpProxy connections during backend connect #960 — Addresses related TcpProxy bookkeeping and shutdown races for connections in progress or cleanup.

Possibly related PRs

• scylladb/python-driver#944 — Modifies related connection and socket lifecycle shutdown coordination.
• scylladb/python-driver#949 — Provides related TcpProxy lifecycle and shutdown synchronization changes.
• scylladb/python-driver#966 — Modifies and tests the shared tests.tcp_proxy.TcpProxy helper.

Suggested reviewers:mykaul, sylwiaszunejko

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The changes keep connections registered through socket cleanup and add the required deterministic regression test for issue #962.
Out of Scope Changes check	✅ Passed	The changes are limited to the TcpProxy test helper and its unit coverage, matching the linked issue scope.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Title check	✅ Passed	The title clearly describes tracking TcpProxy forwarders during socket cleanup, which is the main change.
Description check	✅ Passed	The description explains the change, scope, linked issue, regression test, and validation results; the missing checklist items are non-critical.

---

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

[Copilot]

Pull request overview

Keeps TcpProxy connections tracked until socket cleanup completes, fixing inaccurate connection reporting and shutdown synchronization.

Changes:

• Reorders forwarder cleanup and connection deregistration.
• Adds deterministic regression coverage for blocked socket cleanup.

Reviewed changes

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

File	Description
tests/integration/standard/test_client_routes.py	Deregisters connections after closing sockets.
tests/unit/test_tcp_proxy.py	Tests tracking throughout socket cleanup.

---

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

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unit/test_tcp_proxy.py (1)

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

Avoid materializing _connections for one entry.

Ruff reports RUF015 on this line. Use next(iter(...)) to avoid creating a list.

Proposed change

- connection, thread = list(self.proxy._connections.items())[0]+ connection, thread = next(iter(self.proxy._connections.items()))

🤖 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_tcp_proxy.py` at line 187, Update the connection extraction
in the test around self.proxy._connections to use next(iter(...)) and retrieve
the single key-value pair without materializing the dictionary items into a
list.

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.
Inline comments:
In `@tests/unit/test_tcp_proxy.py`:
- Around line 192-195: Update blocking_close_pair to wait indefinitely on
allow_cleanup before calling real_close_pair, removing the timeout so cleanup
cannot proceed until the test assertions and finally block release the barrier.
---
Nitpick comments:
In `@tests/unit/test_tcp_proxy.py`:
- Line 187: Update the connection extraction in the test around
self.proxy._connections to use next(iter(...)) and retrieve the single key-value
pair without materializing the dictionary items into a list.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: c69fef7b-02cf-4866-b549-6152e33c0890

📥 Commits

Reviewing files that changed from the base of the PR and between d99dc46 and 4d6df3a.

📒 Files selected for processing (2)

• tests/integration/standard/test_client_routes.py
• tests/unit/test_tcp_proxy.py

[Copilot]

Pull request overview

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

[nikagra]

Verified and LGTM. The one-line reorder is correct and minimal, and the new test genuinely fails on the pre-fix ordering rather than passing on both.

The invariant this establishes

Entry absent from _connections => both sockets are already closed. The old pop-then-close window left a live forwarder holding open fds untracked, so a concurrent stop()/drop_connections() could snapshot straight past it while active_connections under-reported.

Why the new order is safe

Worth recording, since it's the non-obvious half. Close-before-pop means a concurrent stop() can now snapshot an entry whose sockets are already closed and call _shutdown_pair() on them. That's a harmless no-op rather than an fd-reuse hazard: CPython sets sock_fd = INVALID_SOCKETbefore the OS close(2), so shutdown() on a closed socket raises OSError EBADF (which _shutdown_pair catches) and never hands a recycled fd to the kernel. Confirmed locally - errno 9, fileno() == -1. So the _shutdown_and_join_connections docstring guarantee, "only the forwarder thread itself closes its sockets", still holds.

Two other properties the swap preserves:

• _close_pair() is still called outside_lock, so active_connections stays reachable during cleanup. No deadlock - and this is what makes the new test observable at all.
• _close_pair() swallows per-socket exceptions, so the pop() is still unconditionally reached; a failing close() can't leak the entry.

Verification

Check	Result
New test logic on PR head	40/40 pass
Same logic on pre-fix ordering	fails: active_connections == 0, expected 1
Existing stress test against new ordering	2178 connections, 0 escaping exceptions, active_connections == 0, 0 live forwarders

Method: extracted TcpProxy and exec'd it standalone against a loopback echo server, so the pre/post comparison isolates the two-line swap and nothing else.

Also confirming patch.object(..., new=staticmethod(blocking_close_pair)) is the right construction here - mock reads target.__dict__[name], so it restores the staticmethod descriptor rather than degrading it to a bound method. Written as a plain function it would silently break self._close_pair(csock, tsock) for the sibling tests, so the wrapper is doing real work.

Non-blocking notes

Take or leave any of these:

• The class docstring (lines 105-115) still describes only #948, though the module docstring was updated to "#948 and #962". The new test's one-line docstring also departs from the file's convention of spelling out exactly what regresses - both sibling tests do.
• Line 187 reads _connections without _lock; test_concurrent_stop_and_drop_leaves_no_live_forwarders takes the lock for the same read. Safe as written (single connection, idle accept loop) and consistent with the #948 test at line 144.
• The test asserts the mechanism - the entry stays tracked. The harm named in #962 is the consequence: a concurrent stop()/drop_connections() misses the forwarder. A thread calling drop_connections() while _close_pair() is blocked would cover that directly.
• Pre-existing, out of scope: _handle_new_connection's t.start() failure path pops before closing - the mirror image of what this fixes. It's inside _lock, so no observer can see the intermediate state; correct as-is, just asymmetric now.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[Copilot]

Pull request overview

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

[Copilot]

Pull request overview

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

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Note

Quiet mode is enabled, so only the most important comments were posted inline. Other review comments are grouped below.

🟡 Other comments (2)

tests/unit/test_tcp_proxy.py-17-17 (1)

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

Add Fixes: #962`` to the PR description.

The supplied PR description has no required issue-closing annotation.

As per coding guidelines, add appropriate Fixes: annotations to the pull request description.

🤖 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_tcp_proxy.py` at line 17, Update the pull request description
to include the issue-closing annotation “Fixes: `#962`” for the shutdown/join
synchronization change. Do not modify the implementation or tests.

Source: Coding guidelines

tests/unit/test_tcp_proxy.py-238-242 (1)

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

Suppress the intentional broad exception.

Line 241 triggers Ruff BLE001. Add an explicit # noqa: BLE001 if this helper must collect every drop_connections() failure.

As per coding guidelines, ensure all commits compile, pass static checks, and pass tests.

🤖 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_tcp_proxy.py` around lines 238 - 242, Update the
drop_connections helper to explicitly suppress Ruff BLE001 on the intentional
broad exception handler, while continuing to collect every failure in
dropper_errors. Keep the existing self.proxy.drop_connections() call and
error-collection behavior unchanged.

Sources: Coding guidelines, 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.
Other comments:
In `@tests/unit/test_tcp_proxy.py`:
- Line 17: Update the pull request description to include the issue-closing
annotation “Fixes: `#962`” for the shutdown/join synchronization change. Do not
modify the implementation or tests.
- Around line 238-242: Update the drop_connections helper to explicitly suppress
Ruff BLE001 on the intentional broad exception handler, while continuing to
collect every failure in dropper_errors. Keep the existing
self.proxy.drop_connections() call and error-collection behavior unchanged.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: QUIET

Plan: Pro Plus

Run ID: 570fd2b2-0dbf-4e0b-b8cf-10329c15b8f6

📥 Commits

Reviewing files that changed from the base of the PR and between 1d7e601 and 32505c6.

📒 Files selected for processing (2)

• tests/tcp_proxy.py
• tests/unit/test_tcp_proxy.py

[dkropachev]

Regarding the CodeRabbit BLE001 note: the broad catch is intentional because this worker records any escaping drop_connections() failure. The repository Ruff configuration does not enable BLE001, and uv run ruff check tests/unit/test_tcp_proxy.py passes, so no # noqa is needed.