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:
@@ -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:
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user