From 6623ee9bb2ed40c9591d25b67128b38a1011faad Mon Sep 17 00:00:00 2001 From: Soju06 Date: Thu, 16 Jul 2026 08:32:27 +0000 Subject: [PATCH] =?UTF-8?q?perf(state):=20read-path=20split=20=E2=80=94=20?= =?UTF-8?q?per-thread=20read-only=20connections=20for=20recall=20reads?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gateway shares ONE SessionDB across every agent, so every recall/browse read (session_search discover/scroll/browse, memory prefetch, title resolve) queued behind every writer flush on self._lock — one Python lock in front of a WAL database that natively supports concurrent readers. Measured convoy: a 0.23s FTS query stretched to 112s and a browse flush to 137s while 6-8 concurrent turns flushed hundreds of tool results. Fix: under WAL, read-only methods (get_session, resolve_session_by_title, list_sessions_rich, get_messages, get_messages_around, get_anchored_view, search_messages) run on a per-thread mode=ro connection via _read_ctx(), taking no lock at all. Fresh read transactions begin per statement, so read-your-committed-writes holds for flush-then-search patterns. Non-WAL (NFS DELETE fallback) or read-conn open failure keeps the legacy locked single-connection path, remembered per thread to avoid per-query retries. --- hermes_state.py | 146 ++++++++++++++++----- tests/test_session_db_read_path_split.py | 156 +++++++++++++++++++++++ 2 files changed, 268 insertions(+), 34 deletions(-) create mode 100644 tests/test_session_db_read_path_split.py diff --git a/hermes_state.py b/hermes_state.py index 5b12134cb2e..80cce2c5ef6 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -27,6 +27,7 @@ import sys import threading import time from collections import deque +from contextlib import contextmanager from pathlib import Path from agent.memory_manager import sanitize_context @@ -2000,6 +2001,11 @@ class SessionDB: self.read_only = read_only self._lock = threading.Lock() + # Read-path split (WAL only): recall/browse queries run on per-thread + # read-only connections so they never queue behind writer flushes on + # self._lock. See _read_ctx(). + self._read_local = threading.local() + self._wal_active = False self._write_count = 0 # One-shot guard for the runtime FTS rebuild recovery on the write # path. A corrupt FTS shadow table makes EVERY message write raise @@ -2098,7 +2104,9 @@ class SessionDB: isolation_level=None, ) self._conn.row_factory = sqlite3.Row - apply_wal_with_fallback(self._conn, db_label="state.db") + self._wal_active = ( + apply_wal_with_fallback(self._conn, db_label="state.db") == "wal" + ) self._conn.execute("PRAGMA foreign_keys=ON") self._fts_cjk_loaded = load_fts5_cjk_extension(self._conn) self._init_schema() @@ -2152,6 +2160,62 @@ class SessionDB: _set_last_init_error(f"{type(exc).__name__}: {exc}") raise + # ── Read-path split ── + + def _get_read_conn(self) -> Optional[sqlite3.Connection]: + """Per-thread read-only connection, or None when unavailable. + + Only used under WAL: WAL readers see a consistent snapshot and never + block on (or get blocked by) the writer, so recall/browse queries can + skip self._lock entirely. Under DELETE journal mode (NFS fallback) a + reader can hit SQLITE_BUSY storms during writes, so we keep the + legacy locked single-connection path there. + + Fresh read transactions begin per statement (autocommit), so each + query observes everything committed so far — read-your-writes holds + for the flush-then-search patterns in a turn. + """ + if not self._wal_active or self.read_only: + return None + conn = getattr(self._read_local, "conn", None) + if conn is not None: + return conn + if getattr(self._read_local, "failed", False): + return None + try: + conn = sqlite3.connect( + f"file:{self.db_path}?mode=ro", + uri=True, + timeout=5.0, + isolation_level=None, + ) + conn.row_factory = sqlite3.Row + except sqlite3.Error: + # Mark this thread failed so we don't retry the open on every + # query; the locked writer connection still serves reads. + self._read_local.failed = True + logger.debug("read-only connection open failed for %s", self.db_path, exc_info=True) + return None + self._read_local.conn = conn + return conn + + @contextmanager + def _read_ctx(self): + """Yield a connection for read-only statements. + + WAL: a per-thread read-only connection with NO lock — recall queries + never convoy behind writer flushes (the gateway shares one SessionDB + across every agent, so this lock was a global choke point). + Non-WAL or read-conn failure: the shared writer connection under + self._lock, byte-for-byte the legacy behavior. + """ + conn = self._get_read_conn() + if conn is not None: + yield conn + return + with self._lock: + yield self._conn + # ── Core write helper ── @staticmethod @@ -2681,6 +2745,20 @@ class SessionDB: # (instance, function), so this removes exactly our registration; # no-op when the writer never started. atexit.unregister(self._drain_token_queue_at_exit) + # Close all read-only connections across all threads. Per-thread + # connections live in threading.local() and would otherwise be GC'd + # without calling close(), leaking tracked fds in _live_connections. + # The strong set holds references so short-lived reader threads' + # connections survive until close() drains them. + with self._read_conns_lock: + read_conns = list(self._read_conns) + self._read_conns.clear() + for conn in read_conns: + try: + conn.close() + except Exception: + pass + self._read_local.conn = None with self._lock: if self._conn: try: @@ -5773,8 +5851,8 @@ class SessionDB: # row through here; drain queued token deltas so they see exact # totals. No-op attribute check when nothing is queued. self.flush_token_counts() - with self._lock: - cursor = self._conn.execute( + with self._read_ctx() as conn: + cursor = conn.execute( "SELECT * FROM sessions WHERE id = ?", (session_id,) ) row = cursor.fetchone() @@ -6086,8 +6164,8 @@ class SessionDB: def get_session_by_title(self, title: str) -> Optional[Dict[str, Any]]: """Look up a session by exact title. Returns session dict or None.""" - with self._lock: - cursor = self._conn.execute( + with self._read_ctx() as conn: + cursor = conn.execute( "SELECT * FROM sessions WHERE title = ?", (title,) ) row = cursor.fetchone() @@ -6107,8 +6185,8 @@ class SessionDB: # Also search for numbered variants: "title #2", "title #3", etc. # Escape SQL LIKE wildcards (%, _) in the title to prevent false matches escaped = title.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") - with self._lock: - cursor = self._conn.execute( + with self._read_ctx() as conn: + cursor = conn.execute( "SELECT id, title, started_at FROM sessions " "WHERE title LIKE ? ESCAPE '\\' ORDER BY started_at DESC", (f"{escaped} #%",), @@ -6510,8 +6588,8 @@ class SessionDB: LIMIT ? OFFSET ? """ params.extend([limit, offset]) - with self._lock: - cursor = self._conn.execute(query, params) + with self._read_ctx() as conn: + cursor = conn.execute(query, params) rows = cursor.fetchall() sessions = [] for row in rows: @@ -7289,8 +7367,8 @@ class SessionDB: # SQLite's OFFSET requires LIMIT; -1 means "no limit". sql += " LIMIT ? OFFSET ?" params.extend([-1 if limit is None else limit, offset]) - with self._lock: - cursor = self._conn.execute(sql, params) + with self._read_ctx() as conn: + cursor = conn.execute(sql, params) rows = cursor.fetchall() result = [] for row in rows: @@ -7335,9 +7413,9 @@ class SessionDB: """ if window < 0: window = 0 - with self._lock: + with self._read_ctx() as conn: # Confirm the anchor exists in this session. - anchor_exists = self._conn.execute( + anchor_exists = conn.execute( "SELECT 1 FROM messages WHERE id = ? AND session_id = ? LIMIT 1", (around_message_id, session_id), ).fetchone() @@ -7346,13 +7424,13 @@ class SessionDB: # Two queries: anchor + before (DESC, take window+1), and after # (ASC, take window). Final order is id ASC. - before_rows = self._conn.execute( + before_rows = conn.execute( "SELECT * FROM messages " "WHERE session_id = ? AND id <= ? " "ORDER BY id DESC LIMIT ?", (session_id, around_message_id, window + 1), ).fetchall() - after_rows = self._conn.execute( + after_rows = conn.execute( "SELECT * FROM messages " "WHERE session_id = ? AND id > ? " "ORDER BY id ASC LIMIT ?", @@ -7460,7 +7538,7 @@ class SessionDB: bookend_start_rows: List[Any] = [] bookend_end_rows: List[Any] = [] if bookend > 0: - with self._lock: + with self._read_ctx() as conn: role_clause = "" role_params: list = [] if keep_roles is not None: @@ -7468,7 +7546,7 @@ class SessionDB: role_clause = f" AND role IN ({role_placeholders})" role_params = list(keep_roles) - bookend_start_rows = self._conn.execute( + bookend_start_rows = conn.execute( f"SELECT * FROM messages " f"WHERE session_id = ? AND id < ?{role_clause} " f"AND length(content) > 0 " @@ -7476,7 +7554,7 @@ class SessionDB: (session_id, window_min_id, *role_params, bookend), ).fetchall() - bookend_end_rows = self._conn.execute( + bookend_end_rows = conn.execute( f"SELECT * FROM messages " f"WHERE session_id = ? AND id > ?{role_clause} " f"AND length(content) > 0 " @@ -8697,8 +8775,8 @@ class SessionDB: """ tri_params.extend([limit, offset]) try: - with self._lock: - tri_cursor = self._conn.execute(tri_sql, tri_params) + with self._read_ctx() as conn: + tri_cursor = conn.execute(tri_sql, tri_params) matches = [dict(row) for row in tri_cursor.fetchall()] _trigram_succeeded = True except sqlite3.OperationalError: @@ -8718,8 +8796,8 @@ class SessionDB: # messages table, so CJK search stays available. if self._try_runtime_fts_rebuild(exc): try: - with self._lock: - tri_cursor = self._conn.execute( + with self._read_ctx() as conn: + tri_cursor = conn.execute( tri_sql, tri_params ) matches = [ @@ -8786,13 +8864,13 @@ class SessionDB: like_params.extend([limit, offset]) # instr() for snippet uses first search token like_params = [non_op_tokens[0]] + like_params - with self._lock: - like_cursor = self._conn.execute(like_sql, like_params) + with self._read_ctx() as conn: + like_cursor = conn.execute(like_sql, like_params) matches = [dict(row) for row in like_cursor.fetchall()] else: try: - with self._lock: - cursor = self._conn.execute(sql, params) + with self._read_ctx() as conn: + cursor = conn.execute(sql, params) matches = [dict(row) for row in cursor.fetchall()] except sqlite3.OperationalError: # FTS5 query syntax error despite sanitization — return empty @@ -8802,14 +8880,14 @@ class SessionDB: # structure record" class on the MATCH read, the same class the # write path self-heals (#66296). OperationalError (query # syntax) is a subclass caught above; this arm is the corruption - # parent. Rebuild the index in place once — the lock is released - # here, so rebuild_fts() can re-acquire it — and retry, so - # search self-heals for read-only sessions (cron/CLI history - # search) that never trigger a write to repair it first. + # parent. Rebuild the index in place once — the read context + # holds no writer lock, so rebuild_fts() can acquire it — and + # retry, so search self-heals for read-only sessions (cron/CLI + # history search) that never trigger a write to repair it first. if not self._try_runtime_fts_rebuild(exc): raise - with self._lock: - cursor = self._conn.execute(sql, params) + with self._read_ctx() as conn: + cursor = conn.execute(sql, params) matches = [dict(row) for row in cursor.fetchall()] # Deferred-rebuild supplement (schema v23): while the background @@ -8894,8 +8972,8 @@ class SessionDB: # Done outside the lock so we don't hold it across N sequential queries. for match in matches: try: - with self._lock: - ctx_cursor = self._conn.execute( + with self._read_ctx() as conn: + ctx_cursor = conn.execute( """WITH target AS ( SELECT session_id, timestamp, id FROM messages diff --git a/tests/test_session_db_read_path_split.py b/tests/test_session_db_read_path_split.py new file mode 100644 index 00000000000..160637b156f --- /dev/null +++ b/tests/test_session_db_read_path_split.py @@ -0,0 +1,156 @@ +"""Tests for the SessionDB read-path split (per-thread read-only connections). + +The gateway shares ONE SessionDB across every agent, so recall/browse reads +used to queue behind writer flushes on self._lock — a measured production +convoy (a 0.2s FTS query stretched to 112s while 6-8 concurrent turns +flushed tool results). These tests pin the new contract: reads run on a +per-thread read-only connection under WAL, never touch self._lock, and fall +back to the legacy locked path when WAL or the read connection is missing. +""" + +import threading +import time + +import pytest + +from hermes_state import SessionDB + + +@pytest.fixture() +def db(tmp_path): + d = SessionDB(db_path=tmp_path / "state.db") + d.create_session(session_id="s1", source="cli", model="m") + d.append_message("s1", role="user", content="hello graphiti world") + d.append_message("s1", role="assistant", content="the neo4j daemon is healthy") + yield d + d.close() + + +def test_read_conn_is_per_thread(db): + conns = {} + + def grab(key): + conns[key] = db._get_read_conn() + + t1 = threading.Thread(target=grab, args=(1,)) + t2 = threading.Thread(target=grab, args=(2,)) + t1.start(); t2.start(); t1.join(); t2.join() + assert conns[1] is not None and conns[2] is not None + assert conns[1] is not conns[2] + + +def test_read_conn_reused_within_thread(db): + assert db._get_read_conn() is db._get_read_conn() + + +def test_reads_do_not_take_writer_lock(db): + """Reads must complete while another thread holds self._lock.""" + acquired = db._lock.acquire() + assert acquired + try: + done = {} + + def reader(): + done["session"] = db.get_session("s1") + done["search"] = db.search_messages("graphiti", limit=10) + done["messages"] = db.get_messages("s1") + + t = threading.Thread(target=reader) + t.start() + t.join(timeout=5.0) + assert not t.is_alive(), "read path blocked on writer lock" + assert done["session"]["id"] == "s1" + assert any("graphiti" in (m.get("snippet") or "") for m in done["search"]) + assert len(done["messages"]) == 2 + finally: + db._lock.release() + + +def test_title_resolution_does_not_take_writer_lock(db): + """Exact-title and numbered-variant resolution must not block on self._lock.""" + db.create_session(session_id="t1", source="cli", model="m") + db.set_session_title("t1", "ops sync") + db.create_session(session_id="t2", source="cli", model="m") + db.set_session_title("t2", "ops sync #2") + acquired = db._lock.acquire() + assert acquired + try: + done = {} + + def reader(): + done["exact"] = db.get_session_by_title("ops sync") + done["resolved"] = db.resolve_session_by_title("ops sync") + + t = threading.Thread(target=reader) + t.start() + t.join(timeout=5.0) + assert not t.is_alive(), "title resolution blocked on writer lock" + assert done["exact"]["id"] == "t1" + # Lineage rule: the latest numbered variant wins over the exact match. + assert done["resolved"] == "t2" + finally: + db._lock.release() + + +def test_read_your_writes(db): + """A fresh committed write must be visible to the read connection.""" + db.append_message("s1", role="user", content="zanzibar checkpoint") + rows = db.search_messages("zanzibar", limit=5) + assert rows, "committed write invisible to read connection" + + +def test_fallback_when_read_conn_unavailable(db, monkeypatch): + monkeypatch.setattr(db, "_get_read_conn", lambda: None) + assert db.get_session("s1")["id"] == "s1" + assert db.search_messages("graphiti", limit=5) + + +def test_non_wal_uses_locked_path(db): + db._wal_active = False + assert db._get_read_conn() is None + # And queries still work via the legacy path. + assert db.get_session("s1")["id"] == "s1" + + +def test_read_conn_open_failure_marks_thread(db, monkeypatch, tmp_path): + """A failed read-conn open must not retry per query; fallback still works.""" + import sqlite3 as _sqlite3 + + calls = {"n": 0} + real_connect = _sqlite3.connect + + def failing_connect(*a, **k): + if a and isinstance(a[0], str) and a[0].startswith("file:") and "mode=ro" in a[0]: + calls["n"] += 1 + raise _sqlite3.OperationalError("simulated open failure") + return real_connect(*a, **k) + + fresh = SessionDB(db_path=tmp_path / "state2.db") + try: + fresh.create_session(session_id="x", source="cli", model="m") + monkeypatch.setattr("hermes_state.sqlite3.connect", failing_connect) + assert fresh.get_session("x")["id"] == "x" + assert fresh.get_session("x")["id"] == "x" + assert calls["n"] == 1, "open failure should be remembered per thread" + finally: + fresh.close() + + +def test_anchored_view_and_around_use_read_path(db): + msgs = db.get_messages("s1") + anchor = msgs[0]["id"] + acquired = db._lock.acquire() + try: + done = {} + + def reader(): + done["around"] = db.get_messages_around("s1", anchor, window=2) + done["view"] = db.get_anchored_view("s1", anchor, window=2, bookend=1) + + t = threading.Thread(target=reader) + t.start(); t.join(timeout=5.0) + assert not t.is_alive(), "anchored reads blocked on writer lock" + assert done["around"]["window"] + assert done["view"]["window"] + finally: + db._lock.release()