postgres-mutex 0.1.0__py3-none-any.whl

postgres_mutex/__init__.py

@@ -0,0 +1,22 @@
+"""postgres_mutex: a zero-dependency distributed mutex for Postgres.
+
+One table, heartbeat-based liveness, automatic recovery from crashed lock holders.
+No ZooKeeper, no Redis, no Consul, no etcd — if you already have Postgres, you
+shouldn't need any of those just to run a job on exactly one instance at a time.
+"""
+
+from .exceptions import LockAcquisitionTimeout, LockNotHeldError, PgMutexError
+from .metrics import MetricsHook, NullMetrics
+from .mutex import Mutex
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Mutex",
+    "PgMutexError",
+    "LockAcquisitionTimeout",
+    "LockNotHeldError",
+    "MetricsHook",
+    "NullMetrics",
+    "__version__",
+]

postgres_mutex/_async.py

@@ -0,0 +1,290 @@
+"""Async implementation: mirrors _sync.py but over psycopg's AsyncConnection and asyncio."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import logging
+import time
+from collections.abc import AsyncIterator
+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 AsyncConnectionPool(Protocol):
+    """Shape of ``psycopg_pool.AsyncConnectionPool`` — matched structurally so
+    postgres_mutex doesn't need psycopg-pool as a hard dependency.
+    """
+
+    def connection(self, timeout: float | None = None) -> Any: ...  # async context manager
+
+
+class AsyncMutexCore:
+    """Async counterpart to ``SyncMutexCore``.
+
+    Either owns one dedicated connection (``dsn`` mode) or borrows short-lived
+    connections from a caller-supplied pool (``pool`` mode) — see ``Mutex.__init__``.
+    In ``dsn`` mode, a single asyncio task holds the connection and both issues
+    heartbeats and serves ad hoc queries (acquire/release); the asyncio.Lock below
+    just prevents two coroutines from interleaving queries on that one connection.
+    In ``pool`` mode the lock is unnecessary (each borrow gets its own connection)
+    but harmless, so the same code path is used for both.
+    """
+
+    def __init__(
+        self,
+        config: MutexConfig,
+        metrics: MetricsHook | None,
+        *,
+        dsn: str | None = None,
+        pool: AsyncConnectionPool | None = None,
+    ) -> None:
+        self.dsn = dsn
+        self.pool = pool
+        self.config = config
+        self.metrics = metrics or NullMetrics()
+        self._conn: psycopg.AsyncConnection[Any] | None = None
+        self._db_lock = asyncio.Lock()
+        self._heartbeat_task: asyncio.Task[None] | None = None
+        self._held = False
+        self._lost = asyncio.Event()
+        self._last_alert_at = 0.0
+
+    async def _connection(self) -> psycopg.AsyncConnection[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 = await psycopg.AsyncConnection.connect(self.dsn, autocommit=True)
+        return self._conn
+
+    @contextlib.asynccontextmanager
+    async def _borrow(self) -> AsyncIterator[psycopg.AsyncConnection[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:
+            async with self.pool.connection() as conn:
+                yield conn
+        else:
+            yield await self._connection()
+
+    async 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"),
+        )
+        async with self._db_lock, self._borrow() as conn, conn.cursor() as cur:
+            await cur.execute(query)
+
+    async 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.
+        """
+        await self.stop_heartbeat()
+        async with self._db_lock:
+            if self._conn is not None and not self._conn.closed:
+                await self._conn.close()
+
+    async 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}
+
+        async with self._db_lock, self._borrow() as conn:
+            try:
+                async with conn.transaction(), conn.cursor() as cur:
+                    await cur.execute(q_select, lock_params)
+                    row = await cur.fetchone()
+
+                    if row is None:
+                        await 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
+                        ):
+                            await cur.execute(q_delete, lock_params)
+                            await cur.execute(q_insert, insert_params)
+                            acquired, reclaimed_from = True, holder
+                        else:
+                            acquired, reclaimed_from = False, None
+            except psycopg.errors.UniqueViolation:
+                # See SyncMutexCore.try_acquire: SELECT ... FOR UPDATE can't serialize
+                # two contenders racing to INSERT a lock row that doesn't exist yet.
+                # The UNIQUE constraint still guarantees exactly one wins; the loser
+                # just lost a race, not an error.
+                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 = asyncio.Event()
+        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,
+        )
+
+    async 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()
+        async with self._db_lock, self._borrow() as conn, conn.cursor() as cur:
+            await cur.execute(query, params)
+            row = await 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,
+        )
+
+    async def release(self) -> None:
+        await 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}
+        async with self._db_lock, self._borrow() as conn, conn.cursor() as cur:
+            await cur.execute(query, params)
+            row = await 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 task -------------------------------------------------
+
+    def start_heartbeat(self) -> None:
+        if self._heartbeat_task is not None:
+            return
+        self._heartbeat_task = asyncio.ensure_future(self._heartbeat_loop())
+
+    async def stop_heartbeat(self) -> None:
+        task, self._heartbeat_task = self._heartbeat_task, None
+        if task is not None and not task.done():
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+
+    async def _heartbeat_loop(self) -> None:
+        try:
+            while True:
+                await asyncio.sleep(self.config.heartbeat_interval)
+                try:
+                    await 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,
+                    )
+        except asyncio.CancelledError:
+            raise
+
+    # -- blocking wait -------------------------------------------------
+
+    async def acquire(self, blocking: bool, timeout: float | None) -> bool:
+        if await 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]
+            await asyncio.sleep(self.config.poll_interval)
+            if await self.try_acquire():
+                self.start_heartbeat()
+                return True

postgres_mutex/_config.py

@@ -0,0 +1,52 @@
+"""Configuration and validation for a Mutex instance."""
+
+from __future__ import annotations
+
+import re
+import socket
+import uuid
+from dataclasses import dataclass, field
+
+_VALID_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
+
+def default_instance_id() -> str:
+    """A reasonably unique, human-traceable default: ``<hostname>-<8 hex chars>``."""
+    return f"{socket.gethostname()}-{uuid.uuid4().hex[:8]}"
+
+
+@dataclass(frozen=True)
+class MutexConfig:
+    """Validated, immutable configuration for one lock.
+
+    The dual-threshold defaults (30s alert / 600s auto-release) match the design
+    rationale in the project README: a single timeout forces a choice between
+    releasing too eagerly (handing the lock to another instance that hits the same
+    slow-processing problem) or too slowly (leaving a crashed holder's job undone
+    for ages). The alert window gives on-call a chance to intervene before the
+    system self-heals.
+    """
+
+    lock_name: str
+    instance_id: str = field(default_factory=default_instance_id)
+    table: str = "mutex_lock"
+    heartbeat_interval: float = 10.0
+    alert_threshold: float = 30.0
+    stale_threshold: float = 600.0
+    poll_interval: float = 0.5
+
+    def __post_init__(self) -> None:
+        if not self.lock_name or len(self.lock_name) > 64:
+            raise ValueError("lock_name must be non-empty and <= 64 characters")
+        if not self.instance_id or len(self.instance_id) > 64:
+            raise ValueError("instance_id must be non-empty and <= 64 characters")
+        if not _VALID_IDENTIFIER.match(self.table):
+            raise ValueError(f"invalid table name: {self.table!r}")
+        if not (0 < self.heartbeat_interval < self.alert_threshold < self.stale_threshold):
+            raise ValueError(
+                "require 0 < heartbeat_interval < alert_threshold < stale_threshold "
+                f"(got {self.heartbeat_interval}, {self.alert_threshold}, "
+                f"{self.stale_threshold})"
+            )
+        if self.poll_interval <= 0:
+            raise ValueError("poll_interval must be > 0")

postgres_mutex/_sql.py

@@ -0,0 +1,69 @@
+"""Raw SQL for the mutex_lock table.
+
+Kept in one place so the acquire/steal logic — the one part of this project that has
+to be exactly right — is easy to audit independently of the Python control flow
+around it.
+"""
+
+from __future__ import annotations
+
+CREATE_TABLE = """
+CREATE TABLE IF NOT EXISTS {table} (
+    lock_name      VARCHAR(64) NOT NULL,
+    locked         INT NOT NULL DEFAULT 1,
+    instance_id    VARCHAR(64) NOT NULL,
+    last_heartbeat TIMESTAMPTZ NOT NULL,
+    acquired_at    TIMESTAMPTZ NOT NULL,
+    CONSTRAINT {table_uq} UNIQUE (locked, lock_name),
+    CONSTRAINT {table_chk} CHECK (locked = 1)
+)
+"""
+
+# Acquire is a locate-then-act sequence run inside one explicit transaction (see
+# SyncMutexCore.try_acquire / AsyncMutexCore.try_acquire):
+#
+#   1. SELECT_FOR_ACQUIRE ... FOR UPDATE locks the row (if any) so two concurrent
+#      acquirers for the same lock_name serialize on it instead of racing.
+#   2. If no row: INSERT_ROW.
+#   3. If the row is stale (past stale_threshold): DELETE_ROW then INSERT_ROW.
+#   4. Otherwise: no writes; the transaction commits as a no-op.
+#
+# Earlier this was a single statement with a data-modifying DELETE ... CTE feeding an
+# INSERT ... ON CONFLICT. That looked atomic but isn't: sibling CTEs in one statement
+# all execute against the *same* snapshot, so the INSERT's conflict check can't see
+# rows the CTE's DELETE just removed — the delete would apply but the insert would
+# still no-op, silently emptying the table. An explicit transaction with row-level
+# locking avoids that pitfall and, as a bonus, lets us read the previous holder's
+# instance_id directly instead of smuggling it out of a CTE.
+SELECT_FOR_ACQUIRE = """
+SELECT instance_id, EXTRACT(EPOCH FROM (now() - last_heartbeat)) AS seconds_since_heartbeat
+FROM {table}
+WHERE lock_name = %(lock_name)s
+FOR UPDATE
+"""
+
+DELETE_ROW = """
+DELETE FROM {table}
+WHERE lock_name = %(lock_name)s
+"""
+
+INSERT_ROW = """
+INSERT INTO {table} (lock_name, locked, instance_id, last_heartbeat, acquired_at)
+VALUES (%(lock_name)s, 1, %(instance_id)s, now(), now())
+"""
+
+HEARTBEAT = """
+UPDATE {table}
+SET last_heartbeat = now()
+WHERE lock_name = %(lock_name)s AND instance_id = %(instance_id)s
+RETURNING 1
+"""
+
+# Clean shutdown: delete the row outright so the next instance doesn't wait out the
+# stale_threshold. Scoped to instance_id so a holder that already lost the lock to a
+# stale-reclaim can't accidentally delete the new holder's row.
+RELEASE = """
+DELETE FROM {table}
+WHERE lock_name = %(lock_name)s AND instance_id = %(instance_id)s
+RETURNING 1
+"""