diff --git a/hermes_state.py b/hermes_state.py index 36274ddceb..fe284fab1f 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -29,6 +29,7 @@ import sqlite3 import sys import threading import time +import weakref from collections import deque from contextlib import contextmanager from pathlib import Path @@ -3173,6 +3174,10 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) _IMPORT_MAX_TOTAL_MESSAGES = 50_000 _IMPORT_MAX_SESSION_BYTES = 5 * 1024 * 1024 _IMPORT_MAX_TOTAL_BYTES = 25 * 1024 * 1024 + # Demand-started accounting workers retire after an idle window so their + # bound targets do not keep abandoned SessionDB instances (and SQLite + # descriptors) alive forever. A later enqueue starts a fresh worker. + _TOKEN_WRITER_IDLE_SECONDS = 30.0 @staticmethod def _store_system_prompt(conn, system_prompt: Optional[str]) -> Optional[str]: @@ -3311,6 +3316,7 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) self._token_writer_thread: Optional[threading.Thread] = None self._token_writer_stop = False self._token_writer_busy = False + self._token_atexit_hook: Optional[Callable[[], None]] = None initialization_complete = False try: if read_only: @@ -4388,12 +4394,9 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) #45383). Read-only connections never request a checkpoint. """ self._stop_token_writer() - # The atexit hook holds a strong reference to this instance (bound - # method); without unregistering, every closed SessionDB stays - # reachable until interpreter exit. Bound methods compare equal by - # (instance, function), so this removes exactly our registration; - # no-op when the writer never started. - atexit.unregister(self._drain_token_queue_at_exit) + hook, self._token_atexit_hook = self._token_atexit_hook, None + if hook is not None: + atexit.unregister(hook) # Drain the read-only connection pool. Setting the closed flag # under the lock first means a reader still in flight closes its own # connection on release instead of re-populating a pool that has @@ -4426,20 +4429,19 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) conn, self._conn = self._conn, None self._close_connection_quietly(conn) + def __enter__(self) -> "SessionDB": + return self + + def __exit__(self, _exc_type, _exc, _tb) -> None: + self.close() + def __del__(self) -> None: """Safety net: close the connection if the caller forgot. - ``atexit.register`` in ``__init__`` pins this instance alive until - interpreter exit, which prevents GC from collecting orphaned - ``SessionDB`` instances on exception paths. When callers forget - ``.close()``, the sqlite FDs leak until the process exits (EMFILE). - - A ``__del__`` finalizer is the last-resort guard: it fires when the - GC collects the object, which *can* happen once ``atexit`` is - unregistered (via ``close()``) **or** when the atexit-held - reference is the only remaining root and the interpreter is - shutting down. During normal interpreter teardown the order of - module cleanup is undefined, so we guard every attribute access. + The async accounting worker retires when idle and its atexit hook + holds only a weak reference, so neither can pin an otherwise orphaned + instance. During interpreter teardown the order of module cleanup is + undefined, so every attribute access remains guarded. Delegates to ``close()`` so the read pool, token writer, and atexit hook are all cleaned up — not just the writer connection. @@ -7257,7 +7259,16 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) ) self._token_writer_thread = thread thread.start() - atexit.register(self._drain_token_queue_at_exit) + if self._token_atexit_hook is None: + self_ref = weakref.ref(self) + + def _drain_at_exit() -> None: + db = self_ref() + if db is not None: + db._drain_token_queue_at_exit() + + self._token_atexit_hook = _drain_at_exit + atexit.register(_drain_at_exit) self._token_queue_cond.notify_all() if writer_stopped: # Writer permanently stopped (close() ran; a stop-flagged but @@ -7323,9 +7334,18 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) def _token_writer_loop(self) -> None: while True: with self._token_queue_cond: + idle_deadline = time.monotonic() + self._TOKEN_WRITER_IDLE_SECONDS while not self._token_queue and not self._token_writer_stop: - self._token_queue_cond.wait() + remaining = idle_deadline - time.monotonic() + if remaining <= 0: + # Publish retirement under the same lock used by + # queue_token_counts() to decide whether to spawn. An + # enqueue cannot strand a delta behind an exiting worker. + self._token_writer_thread = None + return + self._token_queue_cond.wait(remaining) if not self._token_queue: + self._token_writer_thread = None return # stop requested and fully drained # busy is set BEFORE the queue is cleared: the lock-free # fast path in flush_token_counts() reads queue-then-busy, diff --git a/tests/agent/test_async_token_accounting.py b/tests/agent/test_async_token_accounting.py index 5ad0a6fbed..e19ff3c291 100644 --- a/tests/agent/test_async_token_accounting.py +++ b/tests/agent/test_async_token_accounting.py @@ -338,12 +338,54 @@ class TestRouteSwitchBarrier: class TestDurability: + def test_idle_writer_restarts_for_later_delta(self, tmp_path, monkeypatch): + db = SessionDB(db_path=tmp_path / "writer-restart.db") + monkeypatch.setattr(SessionDB, "_TOKEN_WRITER_IDLE_SECONDS", 0.01) + try: + db.create_session("s-restart", "test") + db.queue_token_counts("s-restart", input_tokens=1, api_call_count=1) + assert db.flush_token_counts() + + deadline = time.monotonic() + 2.0 + while db._token_writer_thread is not None and time.monotonic() < deadline: + time.sleep(0.01) + assert db._token_writer_thread is None + + db.queue_token_counts("s-restart", input_tokens=2, api_call_count=1) + assert db.flush_token_counts() + totals = _totals(db, "s-restart") + assert totals["input_tokens"] == 3 + assert totals["api_call_count"] == 2 + finally: + db.close() + + def test_abandoned_db_releases_writer_and_connection(self, tmp_path, monkeypatch): + """An async-accounting writer must not pin an abandoned SessionDB.""" + import gc + import time + import weakref + + from hermes_cli.sqlite_safe_read import has_live_connection + + db_path = tmp_path / "abandoned.db" + monkeypatch.setattr(SessionDB, "_TOKEN_WRITER_IDLE_SECONDS", 0.01, raising=False) + db = SessionDB(db_path=db_path) + db.create_session("s-abandoned", "test") + db.queue_token_counts("s-abandoned", input_tokens=1, api_call_count=1) + assert db.flush_token_counts() + + ref = weakref.ref(db) + del db + deadline = time.monotonic() + 2.0 + while ref() is not None and time.monotonic() < deadline: + gc.collect() + time.sleep(0.01) + + assert ref() is None + assert has_live_connection(db_path) is False def test_close_unregisters_atexit_hook(self, tmp_path): - """close() must unregister the atexit drain hook: it holds a strong - reference (bound method) that would otherwise pin every closed - SessionDB — and its sqlite connection object — until interpreter - exit in multi-open/close processes.""" + """close() unregisters the now-weak atexit drain hook immediately.""" import gc import weakref diff --git a/tests/test_hermes_state.py b/tests/test_hermes_state.py index 6f03a945b6..89abd59d18 100644 --- a/tests/test_hermes_state.py +++ b/tests/test_hermes_state.py @@ -104,6 +104,17 @@ def _no_fts_rebuild_throttle(monkeypatch): class TestConnectionLifecycle: + def test_context_manager_closes_connection(self, tmp_path): + from hermes_cli.sqlite_safe_read import has_live_connection + + db_path = tmp_path / "context-managed.db" + with SessionDB(db_path=db_path) as managed: + managed.create_session("s-context", source="test") + assert has_live_connection(db_path) is True + + assert managed._conn is None + assert has_live_connection(db_path) is False + def test_failed_writable_open_does_not_leak_tracked_connection( self, tmp_path, monkeypatch ):