mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
fix(mcp): close the shutdown drain race + close the owned loop
Review feedback: (1) gating the drain on a main-thread truthiness check of _background_tasks could skip cancellation when a spawn queued via call_soon_threadsafe had not reached the set yet — submit whenever the loop is RUNNING and snapshot on the loop, where FIFO callback order guarantees earlier-queued spawns have landed; (2) shutdown stopped the loop thread but never closed the loop or cleared _loop/_thread, leaking selector resources for embedders that cycle managers — close + clear when we own the thread and it actually stopped (loud warning when it does not); unowned loops (tests wiring _loop directly) stay untouched; (3) the bare await-in-suppress drain loops become asyncio.gather(return_exceptions=True) in both the shutdown drain and the test fixture.
This commit is contained in:
@@ -163,9 +163,7 @@ def running_loop_mgr():
|
||||
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
|
||||
for t in tasks:
|
||||
t.cancel()
|
||||
for t in tasks:
|
||||
with suppress(BaseException):
|
||||
await t
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
with suppress(Exception):
|
||||
asyncio.run_coroutine_threadsafe(_cancel_pending(), loop).result(timeout=5)
|
||||
@@ -2042,6 +2040,37 @@ class TestShutdownCleanup:
|
||||
assert mgr._resource_map == {}
|
||||
assert mgr._prompt_map == {}
|
||||
|
||||
def test_shutdown_closes_owned_loop_and_clears_refs(self):
|
||||
"""When the manager owns the loop thread, shutdown must close the loop
|
||||
(selector resources leak otherwise) and drop both refs; a second
|
||||
shutdown is then a clean no-op."""
|
||||
import threading as _threading
|
||||
|
||||
mgr = MCPClientManager({})
|
||||
loop = asyncio.new_event_loop()
|
||||
thread = _threading.Thread(target=loop.run_forever, daemon=True)
|
||||
thread.start()
|
||||
mgr._loop = loop
|
||||
mgr._thread = thread
|
||||
|
||||
mgr.shutdown()
|
||||
assert loop.is_closed()
|
||||
assert mgr._loop is None
|
||||
assert mgr._thread is None
|
||||
mgr.shutdown() # idempotent
|
||||
|
||||
def test_shutdown_leaves_unowned_loop_open(self):
|
||||
"""Tests (and any embedder) that wire ``_loop`` directly without a
|
||||
thread own the loop's lifecycle — shutdown must not close it."""
|
||||
mgr = MCPClientManager({})
|
||||
loop = asyncio.new_event_loop()
|
||||
mgr._loop = loop
|
||||
try:
|
||||
mgr.shutdown()
|
||||
assert not loop.is_closed()
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TCP probe and unreachable server handling
|
||||
|
||||
@@ -2721,16 +2721,21 @@ class MCPClientManager:
|
||||
"""Close all MCP sessions and stop the background loop."""
|
||||
# Cancel tracked background tasks (catalog refreshes etc.) FIRST —
|
||||
# they are pure auxiliaries, and draining them up front means the
|
||||
# stack teardown below can't race an in-flight refresh.
|
||||
if self._loop and self._background_tasks:
|
||||
# stack teardown below can't race an in-flight refresh. Submitted
|
||||
# whenever a loop exists — NOT gated on a main-thread truthiness
|
||||
# check of ``_background_tasks``: a spawn queued via
|
||||
# call_soon_threadsafe may not have reached the set yet, but ready
|
||||
# callbacks run in FIFO order, so by the time the drain coroutine
|
||||
# snapshots the set ON the loop, every earlier-queued spawn has
|
||||
# landed. ``is_running()`` guard: on a stopped loop nothing can
|
||||
# execute the drain — submitting would just stall on the future.
|
||||
if self._loop is not None and self._loop.is_running():
|
||||
|
||||
async def _cancel_background() -> None:
|
||||
tasks = list(self._background_tasks)
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
for task in tasks:
|
||||
with contextlib.suppress(BaseException):
|
||||
await task
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
future = asyncio.run_coroutine_threadsafe(_cancel_background(), self._loop)
|
||||
try:
|
||||
@@ -2793,6 +2798,18 @@ class MCPClientManager:
|
||||
self._loop.call_soon_threadsafe(self._loop.stop)
|
||||
if self._thread:
|
||||
self._thread.join(timeout=5)
|
||||
if self._thread.is_alive():
|
||||
# Closing a still-running loop raises; the daemon thread dies
|
||||
# with the process, so leaving the loop open is the lesser
|
||||
# evil. Loud because a stuck loop thread is itself a bug.
|
||||
log.warning("MCP loop thread did not stop within 5s; loop left open")
|
||||
else:
|
||||
if self._loop is not None:
|
||||
self._loop.close()
|
||||
self._loop = None
|
||||
self._thread = None
|
||||
# When no thread was started (tests wire ``_loop`` directly) the loop
|
||||
# is not ours to close — the stop above is all the owner needs.
|
||||
|
||||
# Clear all state
|
||||
self._background_tasks.clear()
|
||||
|
||||
Reference in New Issue
Block a user