Add a Postgres advisory-lock leader election for multi-host deploys (#229)

The subscription notifier elects one leader via a host-local flock
(THERMOGRAPH_SINGLETON_LOCK) so multiple uvicorn workers on one host don't
each run it. Under multi-host Swarm that guard is insufficient: each host
would independently elect its own leader, multiplying Open-Meteo quota use
N-fold again.

Add claim_pg(key) alongside the existing claim(lock_path): a cluster-wide
Postgres advisory lock, visible to every host talking to the same database.
The holding connection is dedicated and kept for the process lifetime
(advisory locks are session-scoped); a dead connection is dropped and
re-election retried on the next call.

claim_leader() dispatches between the two mechanisms from env:
THERMOGRAPH_SINGLETON_PG (+ Postgres) -> claim_pg; else
THERMOGRAPH_SINGLETON_LOCK -> claim; else always leader, unchanged. Wired in
at web/app.py in place of the direct claim() call. Off by default, so
today's single-host behavior is unaffected.
This commit is contained in:
Emi Griffith 2026-07-20 17:07:52 -07:00 committed by GitHub
parent bf889681f6
commit 86a6e0aaac
3 changed files with 221 additions and 5 deletions

View file

@ -21,6 +21,13 @@ lock (``flock``) on a lockfile:
No lockfile configured (single-worker deploy, dev, tests) => always the leader. No lockfile configured (single-worker deploy, dev, tests) => always the leader.
This mirrors the metrics store's env-selected coordination: unset keeps the This mirrors the metrics store's env-selected coordination: unset keeps the
zero-dependency single-worker behavior; a path opts into cross-worker arbitration. 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 fcntl
import os 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. # letting it be GC'd) would drop the flock and let a second worker also become leader.
_held: dict[str, int] = {} _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: def claim(lock_path: "str | None") -> bool:
"""Return True if this process should own the single-instance job guarded by """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 return False
_held[lock_path] = fd _held[lock_path] = fd
return True 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"))

View file

@ -3,6 +3,8 @@ import fcntl
import importlib import importlib
import os import os
import psycopg
from core import singleton from core import singleton
@ -62,3 +64,137 @@ def test_reelection_after_holder_releases(tmp_path):
fcntl.flock(other, fcntl.LOCK_UN) fcntl.flock(other, fcntl.LOCK_UN)
os.close(other) os.close(other)
assert s.claim(lock) is True 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]

View file

@ -134,11 +134,12 @@ async def _lifespan(app):
places.start_loading() places.start_loading()
threading.Thread(target=_neighbor_warmer, name="neighbor-warmer", daemon=True).start() threading.Thread(target=_neighbor_warmer, name="neighbor-warmer", daemon=True).start()
# Background subscription evaluator. Its periodic sweep fetches upstream on a timer # 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 # (not per request), so with multiple workers — or multiple *hosts* under Swarm —
# leader via a lockfile (THERMOGRAPH_SINGLETON_LOCK). Unset (single worker / dev / # it must run in exactly one: elect a leader (see singleton.claim_leader, which
# tests) => always the leader. The neighbor warmer above stays per-worker: each # picks a host-local flock or a cluster-wide Postgres advisory lock from env).
# drains its own request-fed queue. # Unset (single worker / dev / tests) => always the leader. The neighbor warmer
if singleton.claim(os.environ.get("THERMOGRAPH_SINGLETON_LOCK")): # above stays per-worker: each drains its own request-fed queue.
if singleton.claim_leader():
notify.start() notify.start()
yield yield
notify.stop() notify.stop()