thermograph/backend/core/singleton.py

"""Single-instance leader election for in-process background jobs.

With multiple uvicorn workers each worker process runs its own app lifespan, so a
background thread started there runs once *per worker*. That is correct for the
neighbor warmer (each worker drains its own request-fed queue) but wrong for the
subscription notifier: its periodic sweep scans every subscription and issues
upstream forecast/archive fetches on a timer, independent of any request. Running
it in N workers multiplies the Open-Meteo quota use N-fold (the source of the
overnight 429/503 rate-limit errors after prod went to 3 workers) and does the
same DB scan N times.

``claim`` elects one worker to own such jobs via a non-blocking exclusive advisory
lock (``flock``) on a lockfile:

* the first worker to call it acquires the lock and holds the fd for the process
  lifetime, so the lock is never dropped while it lives -> returns ``True``;
* other workers fail the non-blocking lock and stand down -> return ``False``;
* the OS releases the lock when the leader exits, so a restart (or a crash) lets a
  surviving/new worker win on its next call -> clean re-election, no stale state.

No lockfile configured (single-worker deploy, dev, tests) => always the leader.
This mirrors the metrics store's env-selected coordination: unset keeps the
zero-dependency single-worker behavior; a path opts into cross-worker arbitration.

``claim_pg`` is the multi-*host* analogue: ``claim``'s flock only arbitrates workers
on one machine, so under multi-host Swarm every host would independently elect its
own leader and multiply the Open-Meteo quota use N-fold again. A Postgres advisory
lock is visible to every host talking to the same database, so it elects exactly one
leader cluster-wide. ``claim_leader()`` picks whichever mechanism the deployment has
configured (see its docstring) so callers never choose between them directly.
"""
import fcntl
import os

# lock_path -> held fd. Kept open for the whole process lifetime: closing the fd (or
# letting it be GC'd) would drop the flock and let a second worker also become leader.
_held: dict[str, int] = {}

# Local copy of the Postgres DSN, mirroring core/metrics.py's identical helper
# rather than importing data/climate_store.py: core/ is a lower-level package that
# data/ depends on, never the reverse (see CLAUDE.md's package layering).
DATABASE_URL = os.environ.get("THERMOGRAPH_DATABASE_URL", "").strip()
IS_POSTGRES = DATABASE_URL.startswith("postgresql")


def _libpq_dsn(url: str) -> str:
    """A plain libpq URL psycopg accepts: drop the SQLAlchemy driver suffix
    (``postgresql+asyncpg://`` / ``+psycopg://`` -> ``postgresql://``)."""
    return url.replace("+asyncpg", "").replace("+psycopg", "")


_PG_DSN = _libpq_dsn(DATABASE_URL) if IS_POSTGRES else ""

# key -> the connection holding its advisory lock. pg_try_advisory_lock is
# session-scoped: released the moment its connection closes, so — like _held above
# — the connection must live for the process lifetime, on its own dedicated
# connection (never a pooled/shared one, which could close or be reused mid-hold).
_pg_held: dict[int, object] = {}

# The one lock this module currently arbitrates (the subscription notifier). Any
# int works as an advisory-lock key; bump/extend if a second named lock is added.
NOTIFIER_LOCK_KEY = 1


def claim(lock_path: "str | None") -> bool:
    """Return True if this process should own the single-instance job guarded by
    ``lock_path``. ``None``/empty (no multi-worker coordination configured) is always
    the leader. Idempotent: re-claiming a path this process already holds returns True
    without touching the lock."""
    if not lock_path:
        return True
    if lock_path in _held:
        return True
    fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644)
    try:
        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except OSError:
        # Another process already holds it — we're not the leader.
        os.close(fd)
        return False
    _held[lock_path] = fd
    return True


def claim_pg(key: int) -> bool:
    """Return True if this process holds the cluster-wide Postgres advisory lock
    ``key``. Requires ``THERMOGRAPH_DATABASE_URL`` to be a Postgres URL; returns
    False (never raises) if Postgres is unreachable — the caller treats "not
    leader" and "unavailable" the same way: stand down.

    Idempotent like ``claim``: a process already holding ``key`` returns True
    without opening a second connection. A dead held connection (Postgres restart,
    network blip) is dropped and re-election is attempted on the next call — the
    advisory lock it held was already released the moment the connection died."""
    conn = _pg_held.get(key)
    if conn is not None:
        if not conn.closed:
            return True
        del _pg_held[key]

    if not IS_POSTGRES:
        return False
    try:
        import psycopg  # local import: only the Postgres path needs it
        conn = psycopg.connect(_PG_DSN, autocommit=True)
        won = conn.execute("SELECT pg_try_advisory_lock(%s)", (key,)).fetchone()[0]
    except Exception:  # noqa: BLE001 - unreachable Postgres => not leader, not fatal
        return False
    if not won:
        conn.close()
        return False
    _pg_held[key] = conn
    return True


def claim_leader() -> bool:
    """Elect the leader for single-instance background jobs (today: the
    subscription notifier), picking the coordination mechanism from env:

    * ``THERMOGRAPH_SINGLETON_PG`` set + the climate record on Postgres -> a
      cluster-wide advisory lock (``claim_pg``) — the only mechanism that elects
      one leader across *every host* under multi-host Swarm.
    * else ``THERMOGRAPH_SINGLETON_LOCK`` -> the host-local flock (``claim``),
      today's single-host-multi-worker behavior.
    * else -> always the leader (single-worker deploy, dev, tests) — unchanged.
    """
    if os.environ.get("THERMOGRAPH_SINGLETON_PG") and IS_POSTGRES:
        return claim_pg(NOTIFIER_LOCK_KEY)
    return claim(os.environ.get("THERMOGRAPH_SINGLETON_LOCK"))