perf(state): read-path split — per-thread read-only connections for recall reads

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.
This commit is contained in:
Soju06 2026-07-16 08:32:27 +00:00 committed by kshitij
parent 2e9559adf0
commit 6623ee9bb2
2 changed files with 268 additions and 34 deletions

View file

@ -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

View file

@ -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()