postgres-mutex 0.1.0__py3-none-any.whl

postgres_mutex/_sync.py

@@ -0,0 +1,298 @@
+"""Sync implementation: connection handling, acquire/release/heartbeat over psycopg 3."""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import threading
+import time
+from collections.abc import Iterator
+from typing import Any, Protocol
+
+import psycopg
+from psycopg import sql
+
+from . import _sql
+from ._config import MutexConfig
+from .exceptions import LockAcquisitionTimeout, LockNotHeldError
+from .metrics import MetricsHook, NullMetrics, safe_call
+
+logger = logging.getLogger("postgres_mutex")
+
+
+class SyncConnectionPool(Protocol):
+    """Shape of ``psycopg_pool.ConnectionPool`` — matched structurally so postgres_mutex
+    doesn't need psycopg-pool as a hard dependency; any object with this method works.
+    """
+
+    def connection(self, timeout: float | None = None) -> Any: ...  # context manager
+
+
+class SyncMutexCore:
+    """Does the actual SQL work, either over one dedicated connection (``dsn`` mode)
+    or by borrowing short-lived connections from a caller-supplied pool (``pool``
+    mode) — see ``Mutex.__init__``.
+
+    A background thread sends heartbeats while the lock is held. In ``dsn`` mode, all
+    queries against the single shared connection (from the main thread or the
+    heartbeat thread) go through ``_db_lock`` — psycopg connections are not safe for
+    concurrent use from multiple threads. In ``pool`` mode the lock is unnecessary
+    (each borrow gets its own connection) but harmless, so the same code path is used
+    for both to avoid maintaining two.
+    """
+
+    def __init__(
+        self,
+        config: MutexConfig,
+        metrics: MetricsHook | None,
+        *,
+        dsn: str | None = None,
+        pool: SyncConnectionPool | None = None,
+    ) -> None:
+        self.dsn = dsn
+        self.pool = pool
+        self.config = config
+        self.metrics = metrics or NullMetrics()
+        self._conn: psycopg.Connection[Any] | None = None
+        self._db_lock = threading.Lock()
+        self._heartbeat_thread: threading.Thread | None = None
+        self._stop_heartbeat = threading.Event()
+        self._held = False
+        self._lost = threading.Event()
+        self._last_alert_at = 0.0
+
+    # -- connection / schema -------------------------------------------------
+
+    def _connection(self) -> psycopg.Connection[Any]:
+        """The dedicated dsn-mode connection, opened lazily. Not used in pool mode."""
+        if self._conn is None or self._conn.closed:
+            assert self.dsn is not None
+            self._conn = psycopg.connect(self.dsn, autocommit=True)
+        return self._conn
+
+    @contextlib.contextmanager
+    def _borrow(self) -> Iterator[psycopg.Connection[Any]]:
+        """Yield a connection to run one operation against.
+
+        dsn mode: yields the long-lived dedicated connection, unchanged after use.
+        pool mode: borrows one from the pool and returns it on exit — postgres_mutex never
+        owns or closes a caller-supplied pool.
+        """
+        if self.pool is not None:
+            with self.pool.connection() as conn:
+                yield conn
+        else:
+            yield self._connection()
+
+    def ensure_schema(self) -> None:
+        query = sql.SQL(_sql.CREATE_TABLE).format(
+            table=sql.Identifier(self.config.table),
+            table_uq=sql.Identifier(f"{self.config.table}_uq"),
+            table_chk=sql.Identifier(f"{self.config.table}_chk_locked"),
+        )
+        with self._db_lock, self._borrow() as conn, conn.cursor() as cur:
+            cur.execute(query)
+
+    def close(self) -> None:
+        """Close the dedicated dsn-mode connection, if one was opened. In pool mode,
+        this is a no-op — the pool belongs to the caller.
+        """
+        self.stop_heartbeat()
+        with self._db_lock:
+            if self._conn is not None and not self._conn.closed:
+                self._conn.close()
+
+    # -- core operations -------------------------------------------------
+
+    def try_acquire(self) -> bool:
+        q_select = sql.SQL(_sql.SELECT_FOR_ACQUIRE).format(table=sql.Identifier(self.config.table))
+        q_delete = sql.SQL(_sql.DELETE_ROW).format(table=sql.Identifier(self.config.table))
+        q_insert = sql.SQL(_sql.INSERT_ROW).format(table=sql.Identifier(self.config.table))
+        lock_params = {"lock_name": self.config.lock_name}
+        insert_params = {"lock_name": self.config.lock_name, "instance_id": self.config.instance_id}
+
+        with self._db_lock, self._borrow() as conn:
+            try:
+                with conn.transaction(), conn.cursor() as cur:
+                    cur.execute(q_select, lock_params)
+                    row = cur.fetchone()
+
+                    if row is None:
+                        cur.execute(q_insert, insert_params)
+                        acquired, reclaimed_from, holder, seconds_since = True, None, None, None
+                    else:
+                        holder, seconds_since = row
+                        if (
+                            seconds_since is not None
+                            and seconds_since >= self.config.stale_threshold
+                        ):
+                            cur.execute(q_delete, lock_params)
+                            cur.execute(q_insert, insert_params)
+                            acquired, reclaimed_from = True, holder
+                        else:
+                            acquired, reclaimed_from = False, None
+            except psycopg.errors.UniqueViolation:
+                # SELECT ... FOR UPDATE only serializes against a row that already
+                # exists — when the lock is free, two contenders can both see "no
+                # row" and both attempt INSERT. The UNIQUE constraint is still the
+                # source of truth: exactly one wins, and the loser just lost a race,
+                # not an error. Treat it the same as ordinary contention.
+                acquired, reclaimed_from, holder, seconds_since = False, None, None, None
+
+        if not acquired:
+            safe_call(
+                self.metrics,
+                "contended",
+                lock_name=self.config.lock_name,
+                instance_id=self.config.instance_id,
+            )
+            if holder is not None and seconds_since is not None:  # noqa: SIM102 (mypy narrowing)
+                if seconds_since >= self.config.alert_threshold:
+                    self._fire_alert(holder, seconds_since)
+            return False
+
+        self._held = True
+        self._lost.clear()
+        if reclaimed_from is not None:
+            safe_call(
+                self.metrics,
+                "stale_reclaimed",
+                lock_name=self.config.lock_name,
+                instance_id=self.config.instance_id,
+                previous_holder=reclaimed_from,
+            )
+            logger.warning(
+                "postgres_mutex: reclaimed lock %r from stale holder %r (instance %r)",
+                self.config.lock_name,
+                reclaimed_from,
+                self.config.instance_id,
+            )
+        safe_call(
+            self.metrics,
+            "acquired",
+            lock_name=self.config.lock_name,
+            instance_id=self.config.instance_id,
+        )
+        return True
+
+    def _fire_alert(self, holder: str, seconds_since: float) -> None:
+        """Page on-call: the current holder's heartbeat is stale enough to investigate
+        (past alert_threshold) but not yet stale enough to auto-reclaim.
+        """
+        now = time.monotonic()
+        if now - self._last_alert_at < self.config.alert_threshold:
+            return  # already alerted recently; don't page repeatedly every poll
+        self._last_alert_at = now
+
+        logger.warning(
+            "postgres_mutex: lock %r held by %r with no heartbeat for %.1fs "
+            "(alert threshold %.1fs)",
+            self.config.lock_name,
+            holder,
+            seconds_since,
+            self.config.alert_threshold,
+        )
+        safe_call(
+            self.metrics,
+            "alert",
+            lock_name=self.config.lock_name,
+            holder=holder,
+            seconds_since_heartbeat=seconds_since,
+        )
+
+    def heartbeat(self) -> None:
+        query = sql.SQL(_sql.HEARTBEAT).format(table=sql.Identifier(self.config.table))
+        params = {"lock_name": self.config.lock_name, "instance_id": self.config.instance_id}
+        start = time.monotonic()
+        with self._db_lock, self._borrow() as conn, conn.cursor() as cur:
+            cur.execute(query, params)
+            row = cur.fetchone()
+        elapsed = time.monotonic() - start
+
+        if row is None:
+            self._held = False
+            self._lost.set()
+            raise LockNotHeldError(self.config.lock_name)
+
+        safe_call(
+            self.metrics,
+            "heartbeat_latency",
+            lock_name=self.config.lock_name,
+            instance_id=self.config.instance_id,
+            seconds=elapsed,
+        )
+
+    def release(self) -> None:
+        self.stop_heartbeat()
+        query = sql.SQL(_sql.RELEASE).format(table=sql.Identifier(self.config.table))
+        params = {"lock_name": self.config.lock_name, "instance_id": self.config.instance_id}
+        with self._db_lock, self._borrow() as conn, conn.cursor() as cur:
+            cur.execute(query, params)
+            row = cur.fetchone()
+
+        self._held = False
+        if row is None:
+            raise LockNotHeldError(self.config.lock_name)
+
+        safe_call(
+            self.metrics,
+            "released",
+            lock_name=self.config.lock_name,
+            instance_id=self.config.instance_id,
+        )
+
+    # -- background heartbeat thread -------------------------------------------------
+
+    def start_heartbeat(self) -> None:
+        if self._heartbeat_thread is not None:
+            return
+        self._stop_heartbeat.clear()
+        self._heartbeat_thread = threading.Thread(
+            target=self._heartbeat_loop,
+            name=f"postgres-mutex-heartbeat-{self.config.lock_name}",
+            daemon=True,
+        )
+        self._heartbeat_thread.start()
+
+    def stop_heartbeat(self) -> None:
+        self._stop_heartbeat.set()
+        thread, self._heartbeat_thread = self._heartbeat_thread, None
+        if thread is not None and thread.is_alive() and thread is not threading.current_thread():
+            thread.join(timeout=self.config.heartbeat_interval)
+
+    def _heartbeat_loop(self) -> None:
+        while not self._stop_heartbeat.wait(self.config.heartbeat_interval):
+            try:
+                self.heartbeat()
+            except LockNotHeldError:
+                logger.error(
+                    "postgres_mutex: lost lock %r (instance %r) — heartbeat missed the "
+                    "stale window and another instance reclaimed it",
+                    self.config.lock_name,
+                    self.config.instance_id,
+                )
+                return
+            except Exception:
+                logger.exception(
+                    "postgres_mutex: heartbeat failed for lock %r (instance %r); will retry",
+                    self.config.lock_name,
+                    self.config.instance_id,
+                )
+
+    # -- blocking wait -------------------------------------------------
+
+    def acquire(self, blocking: bool, timeout: float | None) -> bool:
+        if self.try_acquire():
+            self.start_heartbeat()
+            return True
+        if not blocking:
+            return False
+
+        deadline = None if timeout is None else time.monotonic() + timeout
+        while True:
+            if deadline is not None and time.monotonic() >= deadline:
+                raise LockAcquisitionTimeout(self.config.lock_name, timeout)  # type: ignore[arg-type]
+            time.sleep(self.config.poll_interval)
+            if self.try_acquire():
+                self.start_heartbeat()
+                return True

postgres_mutex/exceptions.py

@@ -0,0 +1,32 @@
+"""Exceptions raised by postgres_mutex."""
+
+from __future__ import annotations
+
+
+class PgMutexError(Exception):
+    """Base class for all postgres_mutex errors."""
+
+
+class LockAcquisitionTimeout(PgMutexError):
+    """Raised when a blocking ``acquire()`` times out before the lock frees up."""
+
+    def __init__(self, lock_name: str, timeout: float) -> None:
+        self.lock_name = lock_name
+        self.timeout = timeout
+        super().__init__(f"timed out after {timeout}s waiting to acquire lock {lock_name!r}")
+
+
+class LockNotHeldError(PgMutexError):
+    """Raised by ``release()`` or ``heartbeat()`` when the caller is not the current holder.
+
+    This happens when the lock was stolen out from under an instance because its
+    heartbeat went stale for longer than ``stale_threshold`` — i.e. another instance
+    already decided this one had crashed.
+    """
+
+    def __init__(self, lock_name: str) -> None:
+        self.lock_name = lock_name
+        super().__init__(
+            f"lock {lock_name!r} is not held by this instance "
+            "(it may have been reclaimed after a missed heartbeat)"
+        )

postgres_mutex/metrics.py

@@ -0,0 +1,77 @@
+"""Metrics hook interface.
+
+postgres_mutex has no hard dependency on any metrics backend. Instead it defines a small
+protocol that Prometheus (``prometheus_client``), OpenTelemetry, statsd, or a
+homegrown metrics object can all satisfy without postgres_mutex knowing about any of them.
+
+Every method is called synchronously and must not raise — a broken metrics backend
+should never take down the lock. ``Mutex`` wraps every call in a best-effort try/except.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class MetricsHook(Protocol):
+    """Implement whichever subset of these you care about; unset methods are no-ops.
+
+    A minimal Prometheus adapter, for example, might implement ``acquired`` as
+    ``counter.labels(lock_name=lock_name).inc()`` and ``heartbeat_latency`` as
+    ``histogram.labels(lock_name=lock_name).observe(seconds)``.
+    """
+
+    def acquired(self, lock_name: str, *, instance_id: str) -> None:
+        """Called every time this instance successfully acquires the lock."""
+
+    def released(self, lock_name: str, *, instance_id: str) -> None:
+        """Called every time this instance cleanly releases the lock."""
+
+    def contended(self, lock_name: str, *, instance_id: str) -> None:
+        """Called every time an acquire attempt finds the lock already held."""
+
+    def stale_reclaimed(self, lock_name: str, *, instance_id: str, previous_holder: str) -> None:
+        """Called when this instance steals a lock whose holder went stale (>10min)."""
+
+    def alert(self, lock_name: str, *, holder: str, seconds_since_heartbeat: float) -> None:
+        """Called when a holder's heartbeat is stale enough to page on-call (>30s)."""
+
+    def heartbeat_latency(self, lock_name: str, *, instance_id: str, seconds: float) -> None:
+        """Called after each heartbeat UPDATE with the round-trip latency."""
+
+
+class NullMetrics:
+    """Default no-op metrics hook used when the caller doesn't supply one."""
+
+    def acquired(self, lock_name: str, *, instance_id: str) -> None:
+        pass
+
+    def released(self, lock_name: str, *, instance_id: str) -> None:
+        pass
+
+    def contended(self, lock_name: str, *, instance_id: str) -> None:
+        pass
+
+    def stale_reclaimed(self, lock_name: str, *, instance_id: str, previous_holder: str) -> None:
+        pass
+
+    def alert(self, lock_name: str, *, holder: str, seconds_since_heartbeat: float) -> None:
+        pass
+
+    def heartbeat_latency(self, lock_name: str, *, instance_id: str, seconds: float) -> None:
+        pass
+
+
+def safe_call(hook: object, method: str, /, **kwargs: object) -> None:
+    """Call ``hook.<method>(**kwargs)`` if present, swallowing any exception it raises.
+
+    A user-supplied metrics backend must never be able to break lock acquisition or
+    release. If it raises, we drop the metric on the floor rather than propagate.
+    """
+    fn = getattr(hook, method, None)
+    if fn is None:
+        return
+    with contextlib.suppress(Exception):  # metrics must never break the mutex
+        fn(**kwargs)