Skip to content

Commit e11836c

Browse files
mykaulclaude
andcommitted
test: fix TcpProxy close/race in test_client_routes.py NLB test helper
TcpProxy.stop()/drop_connections() used to close a connection's sockets directly, from the manager thread, while that connection's own forwarder thread could still be blocked in select()/recv() on those exact file descriptors -- a classic close-under-concurrent-user race. Since fds are process-global, closing them could let the OS silently recycle the fd number into a brand new connection before the stale forwarder thread's blocked call unwound, causing it to read/write/close a socket that no longer belonged to it (observed here as unhandled "ValueError: file descriptor cannot be a negative integer (-1)" crashes in _forward_loop once a socket was closed out from under it). stop() also only ever joined the accept-loop thread, never the per-connection forwarder threads it had just closed sockets out from under. Fix, scoped entirely to this test helper (no driver code touched): - TcpProxy._connections now maps (client_sock, target_sock) -> the forwarder thread serving that pair. - stop()/drop_connections() now shut down (SHUT_RDWR) both sockets -- safe to do concurrently with a blocked select()/recv(), unlike close() -- and then join() every forwarder thread before returning. Only the forwarder thread itself ever closes its own sockets now, and only after it has fully stopped using them. - _handle_new_connection starts the forwarder thread before publishing it into _connections, so a concurrent stop()/drop_connections() can never observe (and try to join) a thread that hasn't started yet. - NLBEmulator.add_node()/remove_node() now serialize against each other via an RLock, so a new proxy/connection can never be created while another thread's remove_node() is still tearing one down. - NLBEmulator._live_addresses() (read by rr_handler(), the discovery port's round-robin accept handler, from the discovery TcpProxy's own accept-loop thread) now also snapshots self._node_proxies under the same _lock that add_node()/remove_node() mutate it under. Previously it iterated the dict unlocked, so a concurrent add_node()/remove_node() could raise "RuntimeError: dictionary changed size during iteration" on that thread. A follow-up pass over the same synchronization path (Copilot automated review on the PR, flagged low-confidence so not posted as formal review threads, but both genuine) found one more real gap and a missing test: - _shutdown_and_join_connections() joined every forwarder thread with a 5s timeout, then unconditionally popped *every* connection out of self._connections regardless of whether its thread had actually exited. A thread that didn't finish in time was dropped from tracking anyway, so active_connections under-reported live connections, and a later stop()/drop_connections() could never retry shutting it down -- permanently leaking that thread and its fds. Fixed by only popping entries whose thread is confirmed dead (`not thread.is_alive()`) after the join; still-alive entries stay tracked until _forward_loop's own self-removal (already lock-protected and idempotent) reaps them, so a subsequent shutdown call can retry. - Added tests/unit/test_tcp_proxy.py, a checked-in deterministic regression test (TcpProxy has no CCM/cluster dependency, only sockets, so it runs as a plain fast unit test against a local dummy TCP echo backend). It covers: (a) the exact regression above -- neutering _shutdown_pair and shrinking one thread's join wait to deterministically force the "still alive after the timeout" path, and asserting the connection stays tracked until a retried drop_connections() actually reaps it -- and (b) a concurrent stress test that hammers drop_connections() from multiple threads while other threads continuously open/close real connections, asserting no unhandled exceptions and no forwarder threads left alive once stop() returns. Validation: - Standalone stress harness (no CCM needed) driving concurrent clients through TcpProxy while repeatedly calling drop_connections()/stop()+restart from another thread: pre-fix, 40 iterations produced 162 unhandled ValueError crashes; post-fix, 40 iterations (same parameters) and a follow-up 150-iteration run produced zero corruptions/exceptions/leaked threads. - Standalone stress harness driving concurrent readers directly exercising _live_addresses() against concurrent add_node()/ remove_node() churn: pre-fix, 313 "dictionary changed size during iteration" RuntimeErrors over 447k calls in 5s; post-fix, zero errors over 1.9M+ calls across three separate runs. - Full end-to-end runs of TestFullNodeReplacementThroughNlb::test_should_survive_full_node_replacement_through_nlb against a real CCM cluster, both before and after the fix, to check for behavioral regressions and reproduce the reported flakiness. - New tests/unit/test_tcp_proxy.py: 30/30 clean runs with the active_connections fix in place; with the fix reverted, the targeted regression test failed deterministically 15/15 runs (active_connections incorrectly reported 0 instead of 1 for a still-alive forwarder thread), confirming the test actually catches the bug it targets. - Full tests/unit/ suite: 722 passed, 88 skipped, 0 failed. Fixes #948. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1 parent 9b5b037 commit e11836c

2 files changed

Lines changed: 351 additions & 21 deletions

File tree

tests/integration/standard/test_client_routes.py

Lines changed: 92 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ def __init__(self, listen_host, listen_port, target_host, target_port):
7373
self._running = False
7474
self._thread = None
7575
self._lock = threading.Lock()
76-
self._connections = set()
76+
self._connections = {} # (client_sock, target_sock) -> forwarder thread
7777
self.total_connections = 0
7878

7979
def start(self):
@@ -92,16 +92,12 @@ def start(self):
9292
self.target_host, self.target_port)
9393

9494
def stop(self):
95-
self._running = False
9695
if self._server_sock:
9796
try:
9897
self._server_sock.close()
9998
except Exception:
10099
pass
101-
with self._lock:
102-
for csock, tsock in list(self._connections):
103-
self._close_pair(csock, tsock)
104-
self._connections.clear()
100+
self._shutdown_and_join_connections(stopping=True)
105101
if self._thread:
106102
self._thread.join(timeout=5)
107103
log.info("TcpProxy stopped %s:%d", self.listen_host, self.listen_port)
@@ -120,12 +116,45 @@ def retarget(self, new_host, new_port):
120116

121117
def drop_connections(self):
122118
"""Forcibly close all active connections."""
123-
with self._lock:
124-
for csock, tsock in list(self._connections):
125-
self._close_pair(csock, tsock)
126-
self._connections.clear()
119+
self._shutdown_and_join_connections()
127120
log.info("TcpProxy %s:%d dropped all connections", self.listen_host, self.listen_port)
128121

122+
def _shutdown_and_join_connections(self, stopping=False):
123+
"""
124+
Shut down (not close) each connection's sockets to unblock its
125+
forwarder thread, then join it. Only the forwarder thread itself
126+
closes its sockets, avoiding a close-vs-still-in-use fd-reuse race.
127+
128+
stopping=True (stop() only) flips _running to False under the same
129+
lock as the connections snapshot, so no connection registered by
130+
_handle_new_connection can be missed.
131+
"""
132+
with self._lock:
133+
if stopping:
134+
self._running = False
135+
connections = list(self._connections.items())
136+
for (csock, tsock), _thread in connections:
137+
self._shutdown_pair(csock, tsock)
138+
finished_keys = []
139+
for (csock, tsock), thread in connections:
140+
thread.join(timeout=5)
141+
if thread.is_alive():
142+
# Do NOT drop this entry from self._connections: it is
143+
# still a live thread owning open fds. Leaving it tracked
144+
# lets active_connections reflect reality and lets a
145+
# subsequent stop()/drop_connections() retry the shutdown
146+
# and join. _forward_loop() removes its own entry (under
147+
# _lock) once it actually exits, so there's no leak here.
148+
log.warning(
149+
"TcpProxy %s:%d: forwarder thread %s did not exit "
150+
"within timeout; leaked fds are possible",
151+
self.listen_host, self.listen_port, thread.name)
152+
else:
153+
finished_keys.append((csock, tsock))
154+
with self._lock:
155+
for key in finished_keys:
156+
self._connections.pop(key, None)
157+
129158
def _run(self):
130159
while self._running:
131160
try:
@@ -153,14 +182,32 @@ def _handle_new_connection(self, client_sock, target_host=None, target_port=None
153182
client_sock.close()
154183
return
155184

156-
with self._lock:
157-
self._connections.add((client_sock, target_sock))
158-
self.total_connections += 1
159-
160185
t = threading.Thread(target=self._forward_loop,
161186
args=(client_sock, target_sock),
162187
daemon=True)
163-
t.start()
188+
# Register then start() atomically under _lock, in that order:
189+
# otherwise a short-lived thread could finish (and clean up)
190+
# before being registered, leaking the entry, or run unseen by
191+
# a concurrent stop()/drop_connections(). Also re-check
192+
# _running, to reject connections after shutdown has begun.
193+
with self._lock:
194+
if not self._running:
195+
target_sock.close()
196+
client_sock.close()
197+
return
198+
self._connections[(client_sock, target_sock)] = t
199+
self.total_connections += 1
200+
try:
201+
t.start()
202+
except Exception as e:
203+
# Undo registration: join()-ing an unstarted thread later
204+
# would raise RuntimeError.
205+
self._connections.pop((client_sock, target_sock), None)
206+
self.total_connections -= 1
207+
log.warning("TcpProxy %s:%d failed to start forwarder thread: %s",
208+
self.listen_host, self.listen_port, e)
209+
client_sock.close()
210+
target_sock.close()
164211

165212
def _forward_loop(self, client_sock, target_sock):
166213
try:
@@ -178,7 +225,7 @@ def _forward_loop(self, client_sock, target_sock):
178225
pass
179226
finally:
180227
with self._lock:
181-
self._connections.discard((client_sock, target_sock))
228+
self._connections.pop((client_sock, target_sock), None)
182229
self._close_pair(client_sock, target_sock)
183230

184231
@staticmethod
@@ -189,6 +236,15 @@ def _close_pair(csock, tsock):
189236
except Exception:
190237
pass
191238

239+
@staticmethod
240+
def _shutdown_pair(csock, tsock):
241+
"""Best-effort shutdown (not close) to interrupt a thread blocked in select()/recv()."""
242+
for s in (csock, tsock):
243+
try:
244+
s.shutdown(socket.SHUT_RDWR)
245+
except OSError:
246+
pass
247+
192248

193249
class NLBEmulator:
194250
"""
@@ -227,7 +283,8 @@ def __init__(self, discovery_port=0,
227283
self._node_proxies = {}
228284
self._discovery_proxy = None
229285
self._rr_index = 0
230-
self._lock = threading.Lock()
286+
# RLock: add_node() holds it while _add_node_proxy() re-acquires it.
287+
self._lock = threading.RLock()
231288
self._running = False
232289

233290
def start(self, node_addresses):
@@ -288,13 +345,17 @@ def stop(self):
288345
log.info("NLB stopped")
289346

290347
def add_node(self, node_id, addr):
291-
self._add_node_proxy(node_id, addr)
348+
# Serialize against remove_node(): TcpProxy.stop() blocks until
349+
# joined, so by the time we get the lock any freed fds are reaped.
350+
with self._lock:
351+
self._add_node_proxy(node_id, addr)
292352

293353
def remove_node(self, node_id):
294354
with self._lock:
295355
proxy = self._node_proxies.pop(node_id, None)
356+
if proxy:
357+
proxy.stop()
296358
if proxy:
297-
proxy.stop()
298359
log.info("NLB removed node %d", node_id)
299360

300361
def node_port(self, node_id):
@@ -328,8 +389,18 @@ def _add_node_proxy(self, node_id, addr):
328389
node_id, self.LISTEN_HOST, port, addr, self.native_port)
329390

330391
def _live_addresses(self):
331-
"""IPs of nodes with active proxies."""
332-
return [p.target_host for p in self._node_proxies.values()]
392+
"""
393+
IPs of nodes with active proxies.
394+
395+
Snapshots under _lock: rr_handler() (the discovery port's accept
396+
handler) calls this from the discovery TcpProxy's own accept-loop
397+
thread, concurrently with add_node()/remove_node() mutating
398+
_node_proxies from other threads. Without the lock, a node being
399+
added/removed mid-iteration can raise "RuntimeError: dictionary
400+
changed size during iteration".
401+
"""
402+
with self._lock:
403+
return [p.target_host for p in self._node_proxies.values()]
333404

334405
def post_client_routes(contact_point, routes):
335406
"""

0 commit comments

Comments
 (0)