From 84f374e80ec97cdfd425ed4f4a630775479b5f71 Mon Sep 17 00:00:00 2001 From: Andres Lopez <190146319+AndresL230@users.noreply.github.com> Date: Wed, 13 May 2026 04:33:24 -0400 Subject: [PATCH 1/2] fix(transport): stop loop on dispose so connect cannot pin the thread MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The old _LocalTransport.dispose() queued a None sentinel and joined with a 2s timeout. If the loop was inside websockets.connect()'s blocking upgrade handshake, the sentinel sat in the queue until the OS TCP timeout (~75 s on Linux), so the join timed out and the daemon thread plus its open socket FD leaked. Schedule loop.stop() via call_soon_threadsafe instead — this interrupts the connect coroutine and lets run_until_complete return promptly. Bound the join with shutdown_flush_timeout_ms (wired through Transport) and warn (not block) if the thread is still alive after that. Regression test stands up a TCP server that accepts but never completes the WebSocket upgrade, then dispose()s mid-handshake and asserts the join returns in under 3 s. Refs #6 --- recost/_transport.py | 54 ++++++++++++++++++++++++++++++++----- tests/test_transport.py | 60 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 6 deletions(-) diff --git a/recost/_transport.py b/recost/_transport.py index 2a02e50..2475e20 100644 --- a/recost/_transport.py +++ b/recost/_transport.py @@ -104,9 +104,15 @@ class _LocalTransport: blocks the caller and never blocks the loop. """ - def __init__(self, port: int, debug: bool = False) -> None: + def __init__( + self, + port: int, + debug: bool = False, + shutdown_timeout_s: float = 3.0, + ) -> None: self._port = port self._debug = debug + self._shutdown_timeout_s = shutdown_timeout_s self._thread: Optional[threading.Thread] = None self._running = False self._loop: Optional[asyncio.AbstractEventLoop] = None @@ -139,6 +145,15 @@ def _run(self) -> None: self._ready.set() try: loop.run_until_complete(self._ws_loop()) + except RuntimeError as exc: + # dispose() schedules loop.stop() to interrupt connect; on + # Python 3.14+ that causes run_until_complete to raise + # "Event loop stopped before Future completed." Treat ONLY + # that case as a clean shutdown; re-raise anything else so + # real coroutine bugs (e.g. "loop is already running") still + # surface. + if "stopped before Future" not in str(exc): + raise finally: try: loop.close() @@ -175,6 +190,9 @@ async def _ws_loop(self) -> None: # unbounded queue and cannot yield the loop. self._queue.put_nowait(msg) break + except asyncio.CancelledError: + # Loop was stopped from dispose() while we were inside connect. + return except Exception: if not self._running: return @@ -197,16 +215,36 @@ def send(self, payload: str) -> None: pass def dispose(self) -> None: + """Stop the loop and join the transport thread. + + The old implementation queued a None sentinel and waited up to 2 s for + the drain coroutine to see it. If the loop was inside + ``websockets.connect()`` (blocking on the upgrade handshake), the + sentinel sat in the queue until the OS TCP timeout (~75 s on Linux), + so the join timed out and the daemon thread + socket FD leaked. + + The fix is to schedule ``loop.stop()`` directly via + ``call_soon_threadsafe`` — this interrupts the connect coroutine, + causing ``run_until_complete`` to return so the thread can exit. + ``shutdown_timeout_s`` (passed in from ``Transport.dispose`` so it + tracks ``shutdown_flush_timeout_ms``) bounds the join. If the thread + is still alive after that, log a warning — never block longer. + """ self._running = False loop = self._loop - queue_ = self._queue - if self._has_websockets and loop is not None and queue_ is not None and not loop.is_closed(): + if self._has_websockets and loop is not None and not loop.is_closed(): try: - asyncio.run_coroutine_threadsafe(queue_.put(None), loop) + loop.call_soon_threadsafe(loop.stop) except RuntimeError: + # Loop was already closed between is_closed() and the schedule. pass if self._thread is not None: - self._thread.join(timeout=2.0) + self._thread.join(timeout=self._shutdown_timeout_s) + if self._thread.is_alive(): + logger.warning( + "[recost] local transport thread did not exit within " + f"{self._shutdown_timeout_s}s; FD may leak" + ) self._thread = None @@ -231,7 +269,11 @@ def __init__(self, config: RecostConfig) -> None: self._local: Optional[_LocalTransport] = None if self.mode == "local": - self._local = _LocalTransport(config.local_port, config.debug) + self._local = _LocalTransport( + config.local_port, + config.debug, + shutdown_timeout_s=config.shutdown_flush_timeout_ms / 1000.0, + ) @property def last_flush_status(self) -> Optional[FlushStatus]: diff --git a/tests/test_transport.py b/tests/test_transport.py index e2e3c1d..fdb2350 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -400,3 +400,63 @@ async def test_send_from_running_event_loop_does_not_block(self): assert elapsed < 0.5 finally: t.dispose() + + +# --------------------------------------------------------------------------- +# Thread safety — dispose during WebSocket connect must not leak the thread +# --------------------------------------------------------------------------- + + +class TestDisposeDuringConnect: + """Regression for issue #6 — when dispose() lands while the loop is inside + websockets.connect()'s blocking TCP+upgrade handshake, the old sentinel- + in-queue dispose path could leave the daemon thread + socket FD pinned + until the OS TCP timeout (~75 s). The fix stops the loop directly.""" + + @staticmethod + def _start_blackhole_server(port: int) -> socket.socket: + """Bind+listen on `port` and accept connections but never respond + to the HTTP upgrade. websockets.connect() will TCP-connect, send its + upgrade request, and then block waiting for an HTTP response.""" + srv = socket.socket() + srv.bind(("127.0.0.1", port)) + srv.listen(8) + accepted: list[socket.socket] = [] + + def accept_loop() -> None: + srv.settimeout(0.5) + while True: + try: + conn, _ = srv.accept() + accepted.append(conn) + except (OSError, socket.timeout): + if srv.fileno() == -1: + return + + threading.Thread(target=accept_loop, daemon=True).start() + return srv + + def test_dispose_during_connect_returns_promptly(self): + pytest.importorskip("websockets") + port = _find_free_port() + server = self._start_blackhole_server(port) + try: + t = _LocalTransport(port=port) + # Give the transport a tick to enter websockets.connect's + # post-TCP-accept handshake wait. + time.sleep(0.5) + start = time.monotonic() + t.dispose() + elapsed = time.monotonic() - start + assert t._thread is None, "transport thread did not join" + # Pre-fix this took ~75s on Linux (OS TCP timeout) or ~2.5s on + # Windows (the old 2s join + dispose overhead). Post-fix the + # loop.stop() path returns in well under a second. The 1.5s + # threshold catches a regression to the queue-sentinel design + # while leaving headroom for slow CI hosts. + assert elapsed < 1.5, ( + f"dispose() took {elapsed:.2f}s — the loop did not stop " + f"while inside websockets.connect()" + ) + finally: + server.close() From 10a8e5177a588419174825454a93ea3f0e478a45 Mon Sep 17 00:00:00 2001 From: Andres Lopez <190146319+AndresL230@users.noreply.github.com> Date: Wed, 13 May 2026 04:40:18 -0400 Subject: [PATCH 2/2] fix(init): flush at process exit via atexit Short-lived processes (cron jobs, Lambda invocations, one-shot CLI scripts, SIGTERM'd containers) dropped their last aggregator bucket because the flush timer is a daemon thread and dies on exit. The user either had to remember to call handle.dispose() manually or accept silent data loss for the final window. Register handle.dispose() via atexit when init() runs, gated by the new RecostConfig.auto_shutdown_handlers (default True). The callback is idempotent and unregisters itself on explicit dispose so a process that cycles init/dispose does not accumulate dead atexit hooks. Regression tests: two in-process checks (registered when enabled, NOT registered when opted out) plus a subprocess end-to-end test that verifies a normal sys.exit(0) triggers the final flush. Refs #6 --- recost/_init.py | 30 ++++++++++ recost/_types.py | 7 +++ tests/test_init.py | 141 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 178 insertions(+) diff --git a/recost/_init.py b/recost/_init.py index 67a4ece..d1850af 100644 --- a/recost/_init.py +++ b/recost/_init.py @@ -5,6 +5,7 @@ from __future__ import annotations +import atexit import sys import threading import warnings @@ -34,6 +35,11 @@ def __init__( self._final_flush = final_flush self._shutdown_flush_timeout_ms = shutdown_flush_timeout_ms self._disposed = False + # Set by init() when auto_shutdown_handlers=True. Cleared in + # dispose() so we don't keep a dangling atexit registration after + # explicit teardown — without this, a long-lived process that + # init/dispose's many times accumulates dead atexit callbacks. + self._atexit_callback: Optional[Callable[[], None]] = None @property def last_flush_status(self) -> Optional[FlushStatus]: @@ -50,12 +56,25 @@ def dispose(self) -> None: ``shutdown_flush_timeout_ms``. The transport is only disposed after the final flush settles or its timeout elapses, so an in-flight cloud POST is not cut off mid-request. + + Idempotent — safe to call from atexit and explicit teardown. """ with _init_lock: if self._disposed: return self._disposed = True + # Unregister the atexit hook so a long-lived process that + # cycles init/dispose does not accumulate dead callbacks. + # atexit.unregister is a no-op if the callable isn't currently + # registered, so this is safe under any registration state. + if self._atexit_callback is not None: + try: + atexit.unregister(self._atexit_callback) + except Exception: + pass + self._atexit_callback = None + self._timer_stop.set() if self._timer_thread is not None: self._timer_thread.join(timeout=5.0) @@ -231,4 +250,15 @@ def _final_flush() -> None: shutdown_flush_timeout_ms=config.shutdown_flush_timeout_ms, ) _handle = handle + + # Register an atexit hook so short-lived processes (cron, Lambda, + # one-shot CLI scripts, SIGTERM'd containers) flush their last + # bucket. The flush timer is a daemon thread and dies on exit; + # without atexit the last window is silently dropped. The hook + # delegates to handle.dispose(), which is idempotent — explicit + # dispose() before exit is still fine. + if config.auto_shutdown_handlers: + handle._atexit_callback = handle.dispose + atexit.register(handle._atexit_callback) + return handle diff --git a/recost/_types.py b/recost/_types.py index cbb9cbe..6309ba8 100644 --- a/recost/_types.py +++ b/recost/_types.py @@ -165,6 +165,13 @@ class RecostConfig: shutdown_flush_timeout_ms: int = 3_000 """How long dispose() waits for the final flush to complete before giving up and closing the transport. Matches Node's shutdownFlushTimeoutMs.""" + auto_shutdown_handlers: bool = True + """When True (default), init() registers an atexit handler that runs the + final flush at normal process termination. Currently only atexit is + wired — a future change may add a SIGTERM handler under the same flag. + Set to False if the host application installs its own lifecycle + management and does not want recost touching atexit (e.g. some test + runners, embedded scripts).""" on_error: Optional[Callable[[Exception], None]] = None diff --git a/tests/test_init.py b/tests/test_init.py index 21c9404..2445445 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -2,7 +2,10 @@ Tests for recost/_init.py """ +import atexit +import subprocess import sys +import textwrap import threading import time @@ -186,3 +189,141 @@ def worker() -> None: assert not exceptions, f"unexpected exceptions: {exceptions!r}" assert init_module._handle is None, "module-level _handle should be None" assert not is_installed(), "interceptor should be uninstalled" + + +# --------------------------------------------------------------------------- +# Process lifecycle — atexit flush +# --------------------------------------------------------------------------- + + +class TestAtexitFlush: + """Regression tests for issue #6 — a normal process exit (sys.exit, + end of __main__, SIGTERM in a container) must flush the last + aggregator bucket. Pre-fix the flush timer is a daemon thread, so + process exit kills it and the last window is silently dropped.""" + + def test_init_registers_atexit_when_enabled(self): + """The atexit module exposes a list of registered callables only + on CPython. We test by recording calls to a patched atexit.register + instead of inspecting atexit internals — portable + behavioral.""" + from recost._init import init + from recost._types import RecostConfig + + registered: list = [] + original_register = atexit.register + + def recording_register(func, *args, **kwargs): + registered.append(func) + return original_register(func, *args, **kwargs) + + atexit.register = recording_register # type: ignore[assignment] + try: + handle = init(RecostConfig(enabled=True, auto_shutdown_handlers=True)) + handle.dispose() + finally: + atexit.register = original_register # type: ignore[assignment] + + # At least one callable was registered by init(). + assert registered, "init() did not register an atexit callback" + + def test_init_skips_atexit_when_disabled(self): + """auto_shutdown_handlers=False opts out of the registration.""" + from recost._init import init + from recost._types import RecostConfig + + registered: list = [] + original_register = atexit.register + + def recording_register(func, *args, **kwargs): + registered.append(func) + return original_register(func, *args, **kwargs) + + atexit.register = recording_register # type: ignore[assignment] + try: + handle = init(RecostConfig(enabled=True, auto_shutdown_handlers=False)) + handle.dispose() + finally: + atexit.register = original_register # type: ignore[assignment] + + assert registered == [], ( + f"init() registered an atexit callback despite the opt-out: {registered!r}" + ) + + def test_subprocess_exit_flushes_final_bucket(self, tmp_path): + """End-to-end: spawn a Python subprocess, init recost, generate one + event, exit normally. Verify the transport's `send` was called — + i.e. the atexit handler ran the final flush before the daemon + thread died.""" + marker = tmp_path / "flush_marker.txt" + + # The child script monkey-patches Transport.send to write a marker + # file when called, then init's recost with a tiny event, then + # exits via sys.exit(0). If atexit works, the marker is written. + script = textwrap.dedent(f""" + import sys + from recost._init import init + from recost._transport import Transport + from recost._types import RecostConfig + + original_send = Transport.send + + def patched_send(self, summary): + with open({str(marker)!r}, "w") as f: + f.write("flushed:" + str(len(summary.metrics))) + return original_send(self, summary) + + Transport.send = patched_send + + handle = init(RecostConfig( + enabled=True, + # Use a long flush_interval so only the atexit/final flush fires. + flush_interval_ms=600_000, + # Disable transport network IO via a fake api_key so cloud + # mode hits our patched send (we don't care if HTTP fails). + api_key="test", + project_id="test", + base_url="http://127.0.0.1:1", + )) + + # Ingest one event directly through the interceptor's callback so + # the aggregator has something to flush. + from recost._types import RawEvent + from recost._interceptor import _callback + assert _callback is not None + _callback(RawEvent( + method="GET", + url="https://api.openai.com/v1/chat/completions", + host="api.openai.com", + path="/v1/chat/completions", + status_code=200, + latency_ms=10, + request_bytes=0, + response_bytes=0, + timestamp="2026-05-13T00:00:00Z", + )) + + sys.exit(0) + """).strip() + + script_path = tmp_path / "child.py" + script_path.write_text(script) + + result = subprocess.run( + [sys.executable, str(script_path)], + capture_output=True, + text=True, + timeout=30, + ) + + assert result.returncode == 0, ( + f"child exited {result.returncode}\nstdout: {result.stdout}\n" + f"stderr: {result.stderr}" + ) + assert marker.exists(), ( + f"atexit final flush did not run. " + f"stdout: {result.stdout}\nstderr: {result.stderr}" + ) + content = marker.read_text() + assert content.startswith("flushed:"), ( + f"marker content unexpected: {content!r}" + )