hermes-agent/hermes_state_registry.py

"""Process-wide shared SessionDB registry.

A gateway process opens state.db from many call sites (runner, SessionStore,
per-agent recall, cron, per-message helpers).  Each bare ``SessionDB()`` mints
its own writer connection, lock, close-time WAL checkpoint and token-writer
thread; N independent writers on one WAL file rely only on SQLite's write lock
plus busy_timeout, 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, not any caller,
  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, not pathname-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 (never a closed stale object) 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 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 —
    ``hermes sessions recover``, snapshot restore — 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, so the next acquire
    retries fresh instead of receiving a closed stale object.
    """
    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 below.
                    _retire_generation_locked(path, generation)
                else:
                    generation.refcount += 1
                    return generation.db

            opening = _opening.get(path)
            if opening is None:
                opening = threading.Event()
                _opening[path] = opening
                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:
            if _opening.get(path) is opening:
                _opening.pop(path, None)
            opening.set()
        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
        if _opening.get(path) is opening:
            _opening.pop(path, None)
        opening.set()
    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 so a close-time WAL checkpoint never stalls acquisition.
    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.
    """
    closed = 0
    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)
        closed += 1
    return closed


def live_shared_session_dbs() -> List["SessionDB"]:
    """Snapshot of every live (non-retired) shared SessionDB.

    For in-process maintenance (housekeeping deferred-FTS retry).  Refcounts
    are NOT touched: 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:
        live = len(_generations)
        retired = len(_retired)
        refs = sum(g.refcount for g in _generations.values())
        return {
            "live_generations": live,
            "retired_generations": retired,
            "total_refcounts": refs,
        }


# ── 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)