Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

hermes_state_registry.py

@@ -0,0 +1,345 @@
+"""Process-wide shared SessionDB registry (#90837).
+
+A gateway process opens state.db from many call sites — the runner's
+``AsyncSessionDB``, the ``SessionStore`` per-path cache, per-agent lazy
+recall (``run_agent._get_session_db_for_recall``), per-job cron opens,
+and per-message opens in mirror / channel_directory / slash_commands /
+shutdown_flush / session_search / react_to_message.  Each bare
+``SessionDB()`` mints its own writer connection, ``self._lock``,
+close-time WAL checkpoint, and async token-writer thread.  With N
+independent writer connections on one WAL file, mutual exclusion relies
+only on SQLite's WAL write lock plus each instance's busy_timeout retry
+ladder — and one connection's close-time checkpoint can race another's
+growth, producing the lost/reordered-page-write signature reported
+across 11+ incidents (#90837).
+
+This module owns that boundary: one shared ``SessionDB`` per resolved
+path per process, refcounted, with generation-aware retirement when the
+underlying file is replaced (snapshot restore, recovery swap).
+
+Lifecycle rules:
+
+- ``acquire(path)`` returns the current generation for *path*,
+  incrementing its refcount.  Same path ⇒ same instance ⇒ same writer
+  connection.
+- ``close()`` on a shared instance is a NO-OP.  The registry — not any
+  individual caller — owns the connection lifecycle, so one caller's
+  ``close()`` can never tear down a writer other callers 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 existing holders release.  If a replacement
+  open fails, the registry is left WITHOUT a path entry (never a closed
+  stale object), so the next acquire retries fresh.
+- 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
+
+if TYPE_CHECKING:  # pragma: no cover - import cycle guard, typed only
+    from hermes_state import SessionDB
+
+logger = logging.getLogger(__name__)
+
+
+def _stat_db_file_identity(path: Path) -> Optional[Tuple[int, int]]:
+    """Return ``(st_dev, st_ino)`` for *path*, or None when unavailable.
+
+    Mirrors the hermes_state helper of the same name; kept local so this
+    module has no import-time dependency on hermes_state (which imports
+    this module — the cycle is resolved by deferring SessionDB lookup
+    to call time).
+    """
+    import os
+
+    try:
+        st = os.stat(path)
+    except OSError:
+        return None
+    # Windows volumes (and some network FS) report st_ino=0; a (0, 0)
+    # identity would false-positive every check. Skip the inode half of
+    # the guard there.
+    if not st.st_dev or not st.st_ino:
+        return None
+    return (st.st_dev, st.st_ino)
+
+
+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 (never retired).  A retired generation leaves
+# this table immediately on retirement and lives on in _retired until
+# its last holder releases.
+_generations: Dict[Path, _Generation] = {}
+# Object-keyed retired generations still draining holders.
+_retired: Dict[int, _Generation] = {}  # id(db) → generation
+# Paths whose next generation is currently being constructed.  Construction
+# stays outside _lock because schema reconciliation can take seconds, but peers
+# for the SAME file must wait: otherwise every cold caller opens a writable
+# SQLite connection before the registry chooses one winner.
+_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.
+
+    The same resolved path always returns the same ``SessionDB`` instance
+    within one process, so all long-lived in-process callers share one
+    writer connection, one ``self._lock``, and one token-writer thread.
+
+    If the underlying file was replaced (different inode) since the
+    shared generation was opened — e.g. by ``hermes sessions recover`` or
+    a snapshot restore — the current generation is RETIRED (never lent
+    again) but stays alive for its existing holders, and a fresh
+    generation is opened in its place.
+
+    Raises whatever ``SessionDB.__init__`` raises (malformed, locked,
+    etc.).  On a replacement-open failure the registry holds NO entry for
+    the path, so the next acquire retries fresh rather than handing out
+    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 the live generation (its
+                    # holders keep it until they release) and elect one
+                    # caller to construct 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.  Do not hold the global
+        # registry lock while waiting: unrelated databases continue opening.
+        # A failed opener signals too, so one waiter can retry as the successor.
+        opening.wait()
+
+    # Open a fresh generation OUTSIDE the lock.  The per-path opening marker
+    # prevents redundant writer connections 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: a generation may have been installed by explicit
+            # registry manipulation while this open was in flight.
+            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).
+
+    The instance stays alive — its holders still own references — and is
+    tracked in ``_retired`` keyed by ``id(db)`` so their releases find
+    the right generation even after the path maps to a new one.
+    """
+    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 a shared instance and its refcount was
+    decremented; ``False`` if *db* is not registry-managed (caller owns
+    its own close()).  The final release of a generation tears it down —
+    OUTSIDE the registry lock, so a close-time WAL checkpoint never
+    stalls acquisition for every state.db in the process.
+
+    Object-keyed lookup means an inode replacement cannot strand a
+    still-owned generation: holders of the old generation release into
+    the 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 a shared instance (caller used SessionDB()
+                # directly) — nothing to do; 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: it stops the token writer, checkpoints
+    # the WAL, and drains the read pool — none of which may hold up
+    # acquisition for every other state.db in the process.
+    if needs_teardown:
+        _teardown(db)
+    return True
+
+
+def close_all() -> int:
+    """Close every shared SessionDB in this process, regardless of refcount.
+
+    Called at gateway shutdown (after all agents and cron jobs have
+    finished) to release every WAL write lock and drain every
+    token-writer thread cleanly.  Returns the number of instances
+    closed.  Idempotent.
+    """
+    closed = 0
+    with _lock:
+        generations = list(_generations.values()) + list(_retired.values())
+        _generations.clear()
+        _retired.clear()
+        for generation in generations:
+            generation.retired = True
+    # Teardown outside the lock, one generation at a time.
+    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 in this process.
+
+    For periodic in-process maintenance (the gateway housekeeping tick's
+    deferred-FTS retry). Refcounts are NOT touched: the caller only invokes
+    a method on an instance that some holder already keeps alive; a
+    concurrent final release closes it and 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) ──
+# Kept so call sites and tests can import either from hermes_state
+# (the historical path) or from this module directly.
+
+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.
+
+    The one-line cleanup for call sites that previously did a plain
+    ``db.close()``: shared instances return their refcount to the
+    registry (the registry owns the lifecycle), anything else — read-only
+    opens, CLI one-shots, test fakes — falls 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)