mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-31 19:16:29 +00:00
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:
parent
2e9559adf0
commit
6623ee9bb2
2 changed files with 268 additions and 34 deletions
146
hermes_state.py
146
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
|
||||
|
|
|
|||
156
tests/test_session_db_read_path_split.py
Normal file
156
tests/test_session_db_read_path_split.py
Normal 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()
|
||||
Loading…
Add table
Add a link
Reference in a new issue