diff --git a/core/singleton.py b/core/singleton.py index 3adb59b..05aac65 100644 --- a/core/singleton.py +++ b/core/singleton.py @@ -21,6 +21,13 @@ lock (``flock``) on a lockfile: 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 @@ -29,6 +36,31 @@ import os # 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 @@ -48,3 +80,50 @@ def claim(lock_path: "str | None") -> bool: 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")) diff --git a/tests/core/test_singleton.py b/tests/core/test_singleton.py index f0f0ddb..c242178 100644 --- a/tests/core/test_singleton.py +++ b/tests/core/test_singleton.py @@ -3,6 +3,8 @@ import fcntl import importlib import os +import psycopg + from core import singleton @@ -62,3 +64,137 @@ def test_reelection_after_holder_releases(tmp_path): fcntl.flock(other, fcntl.LOCK_UN) os.close(other) assert s.claim(lock) is True + + +# --- claim_pg (Postgres advisory lock) -------------------------------------- +# Hermetic: no real Postgres. IS_POSTGRES/_PG_DSN are set directly on the reloaded +# module (mirroring how tests elsewhere monkeypatch module state), and +# psycopg.connect is stubbed so a "connection" is a bare object carrying only +# what claim_pg touches: .execute(...).fetchone() and .closed/.close(). + +class _FakeCursor: + def __init__(self, won): + self._won = won + + def fetchone(self): + return (self._won,) + + +class _FakeConn: + def __init__(self, won): + self.closed = False + self._won = won + + def execute(self, *a, **kw): + return _FakeCursor(self._won) + + def close(self): + self.closed = True + + +def _postgres(s, monkeypatch): + monkeypatch.setattr(s, "IS_POSTGRES", True) + monkeypatch.setattr(s, "_PG_DSN", "postgresql://fake/dsn") + return s + + +def test_claim_pg_false_when_not_postgres(monkeypatch): + s = _fresh() + monkeypatch.setattr(s, "IS_POSTGRES", False) + assert s.claim_pg(1) is False + assert 1 not in s._pg_held + + +def test_claim_pg_wins_and_holds_the_connection(monkeypatch): + s = _postgres(_fresh(), monkeypatch) + monkeypatch.setattr(psycopg, "connect", lambda *a, **kw: _FakeConn(won=True)) + assert s.claim_pg(1) is True + assert 1 in s._pg_held and s._pg_held[1].closed is False + + +def test_claim_pg_reclaiming_is_idempotent_without_reconnecting(monkeypatch): + s = _postgres(_fresh(), monkeypatch) + held = _FakeConn(won=True) + s._pg_held[1] = held + + def _should_not_be_called(*a, **kw): + raise AssertionError("already holding the lock; must not reconnect") + + monkeypatch.setattr(psycopg, "connect", _should_not_be_called) + assert s.claim_pg(1) is True + assert s._pg_held[1] is held + + +def test_claim_pg_loses_closes_and_does_not_hold(monkeypatch): + s = _postgres(_fresh(), monkeypatch) + lost = _FakeConn(won=False) + monkeypatch.setattr(psycopg, "connect", lambda *a, **kw: lost) + assert s.claim_pg(1) is False + assert lost.closed is True + assert 1 not in s._pg_held + + +def test_claim_pg_reelects_after_a_dead_held_connection(monkeypatch): + s = _postgres(_fresh(), monkeypatch) + dead = _FakeConn(won=True) + dead.closed = True # e.g. a Postgres restart dropped it + s._pg_held[1] = dead + fresh_conn = _FakeConn(won=True) + monkeypatch.setattr(psycopg, "connect", lambda *a, **kw: fresh_conn) + assert s.claim_pg(1) is True + assert s._pg_held[1] is fresh_conn + + +def test_claim_pg_unreachable_postgres_is_not_leader_not_fatal(monkeypatch): + s = _postgres(_fresh(), monkeypatch) + + def _boom(*a, **kw): + raise OSError("connection refused") + + monkeypatch.setattr(psycopg, "connect", _boom) + assert s.claim_pg(1) is False + assert 1 not in s._pg_held + + +# --- claim_leader (env-selected dispatch) ----------------------------------- + +def test_claim_leader_defaults_to_always_leader(monkeypatch): + s = _fresh() + monkeypatch.delenv("THERMOGRAPH_SINGLETON_PG", raising=False) + monkeypatch.delenv("THERMOGRAPH_SINGLETON_LOCK", raising=False) + assert s.claim_leader() is True + + +def test_claim_leader_uses_flock_when_pg_flag_unset(tmp_path, monkeypatch): + s = _fresh() + monkeypatch.delenv("THERMOGRAPH_SINGLETON_PG", raising=False) + lock = str(tmp_path / "notifier.lock") + monkeypatch.setenv("THERMOGRAPH_SINGLETON_LOCK", lock) + assert s.claim_leader() is True + assert lock in s._held + + +def test_claim_leader_ignores_pg_flag_when_not_on_postgres(tmp_path, monkeypatch): + # THERMOGRAPH_SINGLETON_PG is set, but the climate record isn't on Postgres + # (IS_POSTGRES False) -> falls through to the flock, not claim_pg. + s = _fresh() + monkeypatch.setattr(s, "IS_POSTGRES", False) + monkeypatch.setenv("THERMOGRAPH_SINGLETON_PG", "1") + lock = str(tmp_path / "notifier.lock") + monkeypatch.setenv("THERMOGRAPH_SINGLETON_LOCK", lock) + + def _should_not_be_called(key): + raise AssertionError("must not use the pg path when IS_POSTGRES is False") + + monkeypatch.setattr(s, "claim_pg", _should_not_be_called) + assert s.claim_leader() is True + assert lock in s._held + + +def test_claim_leader_dispatches_to_claim_pg_when_configured(monkeypatch): + s = _postgres(_fresh(), monkeypatch) + monkeypatch.setenv("THERMOGRAPH_SINGLETON_PG", "1") + calls = [] + monkeypatch.setattr(s, "claim_pg", lambda key: calls.append(key) or True) + assert s.claim_leader() is True + assert calls == [s.NOTIFIER_LOCK_KEY] diff --git a/web/app.py b/web/app.py index e4a9ec3..3dee27b 100644 --- a/web/app.py +++ b/web/app.py @@ -134,11 +134,12 @@ async def _lifespan(app): places.start_loading() threading.Thread(target=_neighbor_warmer, name="neighbor-warmer", daemon=True).start() # Background subscription evaluator. Its periodic sweep fetches upstream on a timer - # (not per request), so with multiple workers it must run in exactly one — elect a - # leader via a lockfile (THERMOGRAPH_SINGLETON_LOCK). Unset (single worker / dev / - # tests) => always the leader. The neighbor warmer above stays per-worker: each - # drains its own request-fed queue. - if singleton.claim(os.environ.get("THERMOGRAPH_SINGLETON_LOCK")): + # (not per request), so with multiple workers — or multiple *hosts* under Swarm — + # it must run in exactly one: elect a leader (see singleton.claim_leader, which + # picks a host-local flock or a cluster-wide Postgres advisory lock from env). + # Unset (single worker / dev / tests) => always the leader. The neighbor warmer + # above stays per-worker: each drains its own request-fed queue. + if singleton.claim_leader(): notify.start() yield notify.stop()