backend/pg_restore_watchdog.py
Rama master · réplica de lectura del repositorio original.
"""Cancel a pg_restore that waits too long on a table lock.
``pg_restore --clean`` drops every table and index, which needs an exclusive
lock on each. A connection holding a lock elsewhere made the restore wait out
its whole five-minute timeout while app queries queued behind it. A plain
``lock_timeout`` cannot bound that: pg_dump writes ``SET lock_timeout = 0`` into
every dump, and pg_restore replays it (VM test of 1.14.4-rc.2 waited 47.7 s on a
lock held for 45 s despite ``PGOPTIONS=-c lock_timeout=30s``).
This watchdog polls ``pg_stat_activity`` for the restore's own backend, found by
the unique ``application_name`` that restore was started with, and cancels it
once **one statement** has been waiting on a lock for ``limit_s`` without a
break. The restore runs in one transaction, so a cancel rolls it back
completely.
"""
from __future__ import annotations
import logging
import threading
import time
from collections.abc import Callable
logger = logging.getLogger(__name__)
POLL_INTERVAL_S = 1.0
#: A sample is (pid, wait_event_type, query).
Sample = tuple[int, str | None, str]
class LockWatchdog:
"""Decides when to cancel; the polling and the cancel call are injected."""
def __init__(
self,
*,
query: Callable[[], list[Sample]],
cancel: Callable[[int], bool],
limit_s: float,
clock: Callable[[], float] = time.monotonic,
) -> None:
self._query = query
self._cancel = cancel
self._limit_s = limit_s
self._clock = clock
# (pid, statement) -> first time we saw THIS statement waiting on a lock.
# Keyed by statement as well: separate statements each waiting briefly
# must not add up to one cancellation.
self._waiting_since: dict[tuple[int, str], float] = {}
self.cancelled = False
self.sample_failures = 0
def tick(self) -> None:
"""Take one sample and cancel a backend that has waited past the limit.
A failing sample is counted and ignored: losing sight of the restore for
a moment must not silently end the watch.
"""
if self.cancelled:
return
try:
samples = self._query()
except Exception as exc: # noqa: BLE001 — the watch outlives a bad sample
self.sample_failures += 1
logger.debug("pg_restore watchdog could not read pg_stat_activity: %s", exc)
return
now = self._clock()
seen: set[tuple[int, str]] = set()
for pid, wait_event_type, query in samples:
key = (pid, query or "")
if wait_event_type != "Lock":
continue
seen.add(key)
since = self._waiting_since.setdefault(key, now)
if now - since < self._limit_s:
continue
logger.warning(
"pg_restore backend %s waited %.0f s on a lock — cancelling the restore (%s)",
pid,
now - since,
(query or "")[:80],
)
if self._cancel(pid):
self.cancelled = True
else:
logger.warning("pg_restore backend %s did not accept the cancel request", pid)
self._waiting_since[key] = now # try again after another limit
return
# Anything not waiting on a lock right now starts over next time.
for key in set(self._waiting_since) - seen:
del self._waiting_since[key]
class RunningWatchdog:
"""A LockWatchdog ticking on a background thread until ``stop()``."""
def __init__(self, dog: LockWatchdog, dispose: Callable[[], None]) -> None:
self._dog = dog
self._dispose = dispose
self._stop = threading.Event()
self._thread = threading.Thread(target=self._run, name="pg-restore-watchdog", daemon=True)
self._thread.start()
@property
def cancelled(self) -> bool:
return self._dog.cancelled
def _run(self) -> None:
try:
while not self._stop.wait(POLL_INTERVAL_S):
self._dog.tick()
except Exception as exc: # noqa: BLE001 — a broken watchdog must not break the restore
logger.warning("pg_restore lock watchdog stopped: %s", exc)
finally:
# Owned by the thread: a stop() that times out still releases it.
try:
self._dispose()
except Exception as exc: # noqa: BLE001
logger.debug("pg_restore watchdog cleanup failed: %s", exc)
def stop(self) -> None:
self._stop.set()
self._thread.join(timeout=5)
if self._thread.is_alive():
logger.warning("pg_restore lock watchdog did not stop within 5 s; leaving it to finish")
def start(database_url: str, app_name: str, limit_s: float) -> RunningWatchdog:
"""Watch backends named ``app_name`` on the database behind ``database_url``."""
from sqlalchemy import create_engine, text
from sqlalchemy.pool import NullPool
# Own engine, autocommit, no pool: the watchdog must never hold a
# transaction (or a lock) of its own while the restore runs. The connect
# timeout keeps a blackholed server from pinning the thread.
engine = create_engine(
database_url,
poolclass=NullPool,
isolation_level="AUTOCOMMIT",
connect_args={"connect_timeout": 5},
)
def query() -> list[Sample]:
with engine.connect() as conn:
conn.execute(text("SET statement_timeout = '5s'"))
rows = conn.execute(
text(
"SELECT pid, wait_event_type, query FROM pg_stat_activity "
"WHERE application_name = :name AND datname = current_database() "
"AND pid <> pg_backend_pid()"
),
{"name": app_name},
).fetchall()
return [(int(pid), wait, sql or "") for pid, wait, sql in rows]
def cancel(pid: int) -> bool:
with engine.connect() as conn:
return bool(conn.execute(text("SELECT pg_cancel_backend(:pid)"), {"pid": pid}).scalar())
dog = LockWatchdog(query=query, cancel=cancel, limit_s=limit_s)
return RunningWatchdog(dog, dispose=engine.dispose)