245 lines
9.2 KiB
Python
245 lines
9.2 KiB
Python
"""Process-wide shared SessionDB registry.
|
|
|
|
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.
|
|
- ``close()`` on a shared instance is a NO-OP: the registry owns the connection lifecycle.
|
|
- ``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.
|
|
- 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 _db_path_of(db: "SessionDB") -> Optional[Path]:
|
|
"""``Path(db.db_path)`` or None when absent/unconvertible."""
|
|
path = getattr(db, "db_path", None)
|
|
if path is None:
|
|
return None
|
|
try:
|
|
return Path(path)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
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. ``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 names."""
|
|
if db is None:
|
|
return False
|
|
key = id(db)
|
|
with _lock:
|
|
generation = _retired.get(key)
|
|
if generation is None:
|
|
path = _db_path_of(db)
|
|
if path is None:
|
|
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 = _db_path_of(db)
|
|
if path is not None:
|
|
_generations.pop(path, None)
|
|
# 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. 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).
|
|
get_shared_session_db = acquire
|
|
release_shared_session_db = release
|
|
close_shared_session_dbs = 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."""
|
|
if not release(db):
|
|
try:
|
|
db.close()
|
|
except Exception:
|
|
logger.debug("release_or_close fallback close failed", exc_info=True)
|