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/_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/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}" + ) 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()