refactor(state): flatten registry acquire flow; restore mirror cleanup assertion

Post-merge review follow-ups on #100201:

- acquire(): drop the never-iterating 'while True' and the redundant
  'existing is not generation' half of the race check — after the retire
  path, generation is always None on the fresh-open leg, so 'existing is
  not None' is the complete condition. Same behavior, flat flow.
- test_mirror: the conversion to the shared registry dropped the
  cleanup assertion entirely; restore it by patching
  hermes_state.release_or_close and asserting the handle is released
  exactly once after _append_to_sqlite.
This commit is contained in:
kshitijk4poor
2026-09-01 21:07:37 +05:30
committed by kshitij
parent 28044757aa
commit 58472d803a
2 changed files with 47 additions and 46 deletions

View File

@@ -134,45 +134,43 @@ def acquire(db_path: Optional[Path] = None) -> "SessionDB":
path = Path(db_path) if db_path is not None else Path(_default_db_path())
while True:
with _lock:
generation = _generations.get(path)
if generation is not None:
current = _stat_db_file_identity(path)
if (
current is not None
and generation.identity is not None
and current != generation.identity
):
# File replaced: retire the live generation (its
# holders keep it until they release) and fall
# through to opening a fresh one below.
_retire_generation_locked(path, generation)
generation = None
else:
generation.refcount += 1
return generation.db
# Open a fresh generation OUTSIDE the lock: construction can
# take seconds (write-lock patience) and must not block every
# other state.db acquisition in the process.
db = _open_session_db(path)
db._shared_registry_owned = True
identity = _stat_db_file_identity(path)
with _lock:
existing = _generations.get(path)
if existing is not None and existing is not generation:
# Someone else opened a generation while we were
# constructing. Ours loses — close it (outside the
# lock) and use theirs.
existing.refcount += 1
winner = existing.db
with _lock:
generation = _generations.get(path)
if generation is not None:
current = _stat_db_file_identity(path)
if (
current is not None
and generation.identity is not None
and current != generation.identity
):
# File replaced: retire the live generation (its
# holders keep it until they release) and fall
# through to opening a fresh one below.
_retire_generation_locked(path, generation)
else:
_generations[path] = _Generation(db, identity)
winner = db
if winner is not db:
_teardown(db)
return winner
generation.refcount += 1
return generation.db
# Open a fresh generation OUTSIDE the lock: construction can
# take seconds (write-lock patience) and must not block every
# other state.db acquisition in the process.
db = _open_session_db(path)
db._shared_registry_owned = True
identity = _stat_db_file_identity(path)
with _lock:
existing = _generations.get(path)
if existing is not None:
# Someone else opened a generation while we were
# constructing (or retired ours and installed a new one).
# Ours loses — close it (outside the lock) and use theirs.
existing.refcount += 1
winner = existing.db
else:
_generations[path] = _Generation(db, identity)
winner = db
if winner is not db:
_teardown(db)
return winner
def _retire_generation_locked(path: Path, generation: _Generation) -> None:

View File

@@ -119,18 +119,21 @@ class TestMirrorToSession:
class TestAppendToSqlite:
def test_connection_is_closed_after_use(self, tmp_path):
"""Verify _append_to_sqlite releases the shared SessionDB handle."""
def test_connection_is_released_after_use(self, tmp_path):
"""Verify _append_to_sqlite returns the shared SessionDB reference."""
from gateway.mirror import _append_to_sqlite
mock_db = MagicMock()
released = []
with patch("hermes_state.get_shared_session_db", return_value=mock_db):
with patch("hermes_state.get_shared_session_db", return_value=mock_db), \
patch(
"hermes_state.release_or_close",
side_effect=lambda db: released.append(db),
):
_append_to_sqlite("sess_1", {"role": "assistant", "content": "hello"})
mock_db.append_message.assert_called_once()
# Shared instances are released (not closed) — the registry owns close().
# release_shared_session_db is a module-level function, so verify the
# mock was passed to it by checking that append_message was called
# (the real release_shared_session_db on a MagicMock is a no-op since
# the mock isn't in the registry).
assert released == [mock_db], (
"the shared handle must be released exactly once after use"
)