shared/session_reaper.py

"""Idle-session reaper for Gradio demo apps.

Gradio's ``demo.unload()`` callback is JavaScript-based and won't fire on
browser crash, network drop, or mobile tab kill.  This module runs a
background thread that closes sessions after a configurable period of
inactivity, preventing orphaned sessions from inflating metrics.
"""

import threading
import time
from collections.abc import Callable


class SessionReaper:
    """Track per-session activity and reap idle sessions.

    Args:
        timeout_s: Seconds of inactivity before a session is reaped.
        on_expired: Called (outside any lock) with the session hash of each
            reaped session.  The callback is responsible for any cleanup
            (decrementing counters, logging ``session_end``, etc.).
        check_interval_s: How often the reaper thread wakes up to scan.
    """

    def __init__(
        self,
        timeout_s: float,
        on_expired: Callable[[str], None],
        check_interval_s: float = 30.0,
    ) -> None:
        self._timeout_s = timeout_s
        self._on_expired = on_expired
        self._check_interval_s = check_interval_s
        self._sessions: dict[str, float] = {}
        self._lock = threading.Lock()
        self._shutting_down = False
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def add(self, session_hash: str) -> None:
        """Register a new session and record its initial activity."""
        with self._lock:
            self._sessions[session_hash] = time.monotonic()

    def remove(self, session_hash: str) -> bool:
        """Unregister a session. Returns ``True`` if it was tracked."""
        with self._lock:
            return self._sessions.pop(session_hash, None) is not None

    def touch(self, session_hash: str) -> bool:
        """Update last-activity timestamp. Returns ``False`` if unknown."""
        with self._lock:
            if session_hash not in self._sessions:
                return False
            self._sessions[session_hash] = time.monotonic()
            return True

    def is_active(self, session_hash: str) -> bool:
        with self._lock:
            return session_hash in self._sessions

    @property
    def active_count(self) -> int:
        with self._lock:
            return len(self._sessions)

    def shutdown(self) -> None:
        """Signal the reaper thread to stop."""
        self._shutting_down = True

    def _run(self) -> None:
        while not self._shutting_down:
            time.sleep(self._check_interval_s)
            now = time.monotonic()
            with self._lock:
                stale = [
                    sh
                    for sh, ts in self._sessions.items()
                    if now - ts >= self._timeout_s
                ]
                for sh in stale:
                    del self._sessions[sh]
            # Callbacks outside the lock to avoid deadlock with the caller.
            for sh in stale:
                self._on_expired(sh)