270 lines
9.7 KiB
Python
270 lines
9.7 KiB
Python
"""Process-wide shared SessionDB registry.
|
|
|
|
A gateway process opens state.db from many call sites; each bare ``SessionDB()``
|
|
mints its own writer connection, lock, close-time WAL checkpoint and token
|
|
writer thread, and one connection's close-time checkpoint can race another's
|
|
growth (lost/reordered-page-write corruption). This module owns that boundary:
|
|
one shared ``SessionDB`` per resolved path per process, refcounted, with
|
|
generation-aware retirement when the file is replaced (snapshot restore,
|
|
recovery swap).
|
|
|
|
Lifecycle rules:
|
|
|
|
- ``acquire(path)`` returns the current generation for *path* and bumps its
|
|
refcount. Same path ⇒ same instance ⇒ same writer connection.
|
|
- ``close()`` on a shared instance is a NO-OP: the registry owns the connection
|
|
lifecycle, so one caller can never tear down a writer others still hold.
|
|
- ``release(db)`` decrements the generation *db was acquired from* (object-
|
|
keyed, so an inode replacement cannot strand a still-owned generation). The
|
|
final release of a retired generation tears it down.
|
|
- On inode change the old generation is RETIRED (never lent again) but stays
|
|
alive until its holders release. If the replacement open fails the registry
|
|
keeps NO path entry, so the next acquire retries.
|
|
- All teardown happens OUTSIDE the registry lock: a final release's WAL
|
|
checkpoint must never stall acquisition for every state.db.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
|
|
|
|
from hermes_state_common import stat_db_file_identity as _stat_db_file_identity
|
|
|
|
if TYPE_CHECKING: # pragma: no cover - import cycle guard, typed only
|
|
from hermes_state import SessionDB
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class _Generation:
|
|
"""One shared SessionDB generation: instance, refcount, file identity."""
|
|
|
|
__slots__ = ("db", "refcount", "identity", "retired")
|
|
|
|
def __init__(self, db: "SessionDB", identity: Optional[Tuple[int, int]]) -> None:
|
|
self.db = db
|
|
self.refcount = 1
|
|
self.identity = identity
|
|
self.retired = False
|
|
|
|
|
|
_lock = threading.Lock()
|
|
# path → live generation; retired generations move to _retired (keyed by
|
|
# id(db)) until their last holder releases.
|
|
_generations: Dict[Path, _Generation] = {}
|
|
_retired: Dict[int, _Generation] = {}
|
|
# Paths whose next generation is being constructed. Construction runs outside
|
|
# _lock (schema reconciliation can take seconds), but peers for the SAME file
|
|
# must wait or every cold caller opens its own writer before a winner is chosen.
|
|
_opening: Dict[Path, threading.Event] = {}
|
|
|
|
|
|
def _open_session_db(path: Path) -> "SessionDB":
|
|
"""Construct the SessionDB for *path* (call-time import avoids cycles)."""
|
|
from hermes_state import SessionDB
|
|
|
|
return SessionDB(db_path=path)
|
|
|
|
|
|
def _teardown(db: "SessionDB") -> None:
|
|
"""Close a shared instance, clearing its registry-owned flag first."""
|
|
try:
|
|
db._shared_registry_owned = False
|
|
except Exception:
|
|
pass
|
|
try:
|
|
db.close()
|
|
except Exception:
|
|
logger.debug("Error closing shared SessionDB", exc_info=True)
|
|
|
|
|
|
def _finish_opening(path: Path, opening: threading.Event) -> None:
|
|
"""Drop the per-path construction marker and wake waiters (caller holds _lock)."""
|
|
if _opening.get(path) is opening:
|
|
_opening.pop(path, None)
|
|
opening.set()
|
|
|
|
|
|
def acquire(db_path: Optional[Path] = None) -> "SessionDB":
|
|
"""Return the shared SessionDB for *db_path*, incrementing its refcount.
|
|
|
|
If the file was replaced (different inode) since the generation opened,
|
|
that generation is RETIRED but stays alive for its holders, and a fresh one
|
|
is opened in its place. Raises whatever ``SessionDB.__init__`` raises; on
|
|
a replacement-open failure the registry holds NO entry for the path.
|
|
"""
|
|
from hermes_state import _default_db_path
|
|
|
|
raw_path = Path(db_path) if db_path is not None else Path(_default_db_path())
|
|
try:
|
|
path = raw_path.resolve()
|
|
except OSError:
|
|
path = raw_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, then elect one caller to open the replacement.
|
|
_retire_generation_locked(path, generation)
|
|
else:
|
|
generation.refcount += 1
|
|
return generation.db
|
|
opening = _opening.get(path)
|
|
if opening is None:
|
|
opening = _opening[path] = threading.Event()
|
|
break
|
|
# Another caller is constructing this path; wait without holding the
|
|
# global lock. A failed opener signals too, so a waiter can retry.
|
|
opening.wait()
|
|
|
|
# Open OUTSIDE the lock; the per-path marker prevents redundant writers
|
|
# without serialising other files.
|
|
try:
|
|
db = _open_session_db(path)
|
|
db._shared_registry_owned = True
|
|
identity = _stat_db_file_identity(path)
|
|
except BaseException:
|
|
with _lock:
|
|
_finish_opening(path, opening)
|
|
raise
|
|
|
|
with _lock:
|
|
existing = _generations.get(path)
|
|
if existing is not None:
|
|
# Defensive: installed by explicit registry manipulation mid-open.
|
|
existing.refcount += 1
|
|
winner = existing.db
|
|
else:
|
|
_generations[path] = _Generation(db, identity)
|
|
winner = db
|
|
_finish_opening(path, opening)
|
|
if winner is not db:
|
|
_teardown(db)
|
|
return winner
|
|
|
|
|
|
def _retire_generation_locked(path: Path, generation: _Generation) -> None:
|
|
"""Retire *generation* so it is never lent again (caller holds _lock).
|
|
|
|
It stays alive for its holders, tracked in ``_retired`` by ``id(db)`` so
|
|
their releases find it even after the path maps to a new generation.
|
|
"""
|
|
generation.retired = True
|
|
if _generations.get(path) is generation:
|
|
del _generations[path]
|
|
_retired[id(generation.db)] = generation
|
|
|
|
|
|
def release(db: "SessionDB") -> bool:
|
|
"""Decrement the refcount of a shared SessionDB.
|
|
|
|
Returns ``True`` if *db* was shared; ``False`` if it is not registry-managed
|
|
(caller owns close()). The final release tears the generation down OUTSIDE
|
|
the registry lock. Lookup is object-keyed, so holders of an old generation
|
|
release into its retired record, not into whatever the path currently names.
|
|
"""
|
|
if db is None:
|
|
return False
|
|
key = id(db)
|
|
with _lock:
|
|
generation = _retired.get(key)
|
|
if generation is None:
|
|
path = getattr(db, "db_path", None)
|
|
if path is None:
|
|
return False
|
|
try:
|
|
path = Path(path)
|
|
except (TypeError, ValueError):
|
|
return False
|
|
generation = _generations.get(path)
|
|
if generation is None or generation.db is not db:
|
|
# Not shared (bare SessionDB()); the caller owns close().
|
|
return False
|
|
generation.refcount -= 1
|
|
needs_teardown = generation.refcount <= 0
|
|
if needs_teardown:
|
|
if generation.retired:
|
|
_retired.pop(key, None)
|
|
else:
|
|
path = getattr(db, "db_path", None)
|
|
if path is not None:
|
|
try:
|
|
_generations.pop(Path(path), None)
|
|
except (TypeError, ValueError):
|
|
pass
|
|
# Teardown OUTSIDE the lock: stopping the token writer, WAL checkpoint and
|
|
# read-pool drain must not block acquisition for every other state.db.
|
|
if needs_teardown:
|
|
_teardown(db)
|
|
return True
|
|
|
|
|
|
def close_all() -> int:
|
|
"""Close every shared SessionDB regardless of refcount; returns the count.
|
|
|
|
For gateway shutdown, after all agents and cron jobs finished. Idempotent.
|
|
"""
|
|
with _lock:
|
|
generations = list(_generations.values()) + list(_retired.values())
|
|
_generations.clear()
|
|
_retired.clear()
|
|
for generation in generations:
|
|
generation.retired = True
|
|
for generation in generations:
|
|
_teardown(generation.db)
|
|
return len(generations)
|
|
|
|
|
|
def live_shared_session_dbs() -> List["SessionDB"]:
|
|
"""Snapshot of every live (non-retired) shared SessionDB (refcounts untouched).
|
|
|
|
For in-process maintenance (housekeeping deferred-FTS retry). A concurrent
|
|
final release may close an instance, in which case the callee sees
|
|
``_conn is None``.
|
|
"""
|
|
with _lock:
|
|
return [g.db for g in _generations.values() if not g.retired]
|
|
|
|
|
|
def stats() -> Dict[str, int]:
|
|
"""Registry census for tests and diagnostics (no locks held long)."""
|
|
with _lock:
|
|
return {
|
|
"live_generations": len(_generations),
|
|
"retired_generations": len(_retired),
|
|
"total_refcounts": sum(g.refcount for g in _generations.values()),
|
|
}
|
|
|
|
|
|
# ── Backwards-compatible aliases (hermes_state re-exports them) ──
|
|
|
|
def get_shared_session_db(db_path: Optional[Path] = None) -> "SessionDB":
|
|
return acquire(db_path)
|
|
|
|
|
|
def release_shared_session_db(db: "SessionDB") -> bool:
|
|
return release(db)
|
|
|
|
|
|
def close_shared_session_dbs() -> int:
|
|
return close_all()
|
|
|
|
|
|
def release_or_close(db: "SessionDB") -> None:
|
|
"""Release a shared instance, or close it when it is not registry-managed.
|
|
|
|
Drop-in for a plain ``db.close()``: read-only opens, CLI one-shots and
|
|
test fakes fall back to a direct close.
|
|
"""
|
|
if not release(db):
|
|
try:
|
|
db.close()
|
|
except Exception:
|
|
logger.debug("release_or_close fallback close failed", exc_info=True)
|