396 lines
16 KiB
Python
396 lines
16 KiB
Python
 
|
"""The worker's scheduler: APScheduler 3 in site time, settings read on every fire, "run now" over NOTIFY.
|
||
|
|
|
||
|
|
Each job has a SettingsTrigger whose next fire time comes from the job's cadence and the current feed settings,
|
||
|
|
so a changed interval applies on the next fire. A save NOTIFYs `ps_settings` and the listener calls
|
||
|
|
`reschedule()` so it applies at once. A "run now" NOTIFYs `ps_job_request` and the listener moves the job to now.
|
||
|
|
After a run the job's next fire is measured from the end of the run.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import contextlib
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
import os
|
||
|
|
import socket
|
||
|
|
import threading
|
||
|
|
from collections.abc import Callable, Iterable
|
||
|
|
from datetime import datetime, timedelta
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
from zoneinfo import ZoneInfo
|
||
|
|
|
||
|
|
from apscheduler.executors.pool import ThreadPoolExecutor
|
||
|
|
from apscheduler.schedulers.background import BackgroundScheduler
|
||
|
|
from apscheduler.triggers.base import BaseTrigger
|
||
|
|
|
||
|
|
from .. import db
|
||
|
|
from ..feed import config
|
||
|
|
from ..feed.config import FeedSettings
|
||
|
|
from ..settings import Settings
|
||
|
|
from . import heartbeat, runs
|
||
|
|
from .jobs import JobSpec, Stack, load_jobs
|
||
|
|
from .redaction import redact
|
||
|
|
|
||
|
|
log = logging.getLogger("packetstation.worker")
|
||
|
|
|
||
|
|
CHANNELS = ("ps_settings", "ps_job_request")
|
||
|
|
IDLE_RECHECK_SECONDS = 3600 # a job with no schedule wakes this often and does nothing
|
||
|
|
BACKOFF_MAX = 30.0
|
||
|
|
THREADS = 4
|
||
|
|
MISFIRE_GRACE = 300
|
||
|
|
|
||
|
|
|
||
|
|
class SettingsTrigger(BaseTrigger):
|
||
|
|
"""Fires `seconds(now)` after now. `seconds` reads the current settings, so a change applies on the next fire.
|
||
|
|
None means "no schedule": the trigger wakes again after an hour and the fire does nothing."""
|
||
|
|
|
||
|
|
__slots__ = ("seconds",)
|
||
|
|
|
||
|
|
def __init__(self, seconds: Callable[[datetime], int | None]):
|
||
|
|
self.seconds = seconds
|
||
|
|
|
||
|
|
def get_next_fire_time(self, previous_fire_time: datetime | None, now: datetime) -> datetime:
|
||
|
|
secs = self.seconds(now)
|
||
|
|
return now + timedelta(seconds=max(1, IDLE_RECHECK_SECONDS if secs is None else int(secs)))
|
||
|
|
|
||
|
|
def __str__(self) -> str:
|
||
|
|
return "settings"
|
||
|
|
|
||
|
|
|
||
|
|
def next_due(
|
||
|
|
spec: JobSpec, feed: FeedSettings, now_site: datetime, last_end_site: datetime | None = None
|
||
|
|
) -> datetime | None:
|
||
|
|
"""When the job is next due: its cadence measured from the end of its last run (or from now when it never ran),
|
||
|
|
never before now. None when the cadence says "no schedule"."""
|
||
|
|
base = last_end_site if last_end_site is not None and last_end_site <= now_site else now_site
|
||
|
|
secs = spec.cadence(feed, base)
|
||
|
|
if secs is None:
|
||
|
|
return None
|
||
|
|
due = base + timedelta(seconds=max(1, int(secs)))
|
||
|
|
return due if due > now_site else now_site
|
||
|
|
|
||
|
|
|
||
|
|
class Scheduler:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
settings: Settings,
|
||
|
|
*,
|
||
|
|
build: str | None = None,
|
||
|
|
jobs: dict[str, JobSpec] | None = None,
|
||
|
|
connect: Callable[[], Any] | None = None,
|
||
|
|
):
|
||
|
|
self.settings = settings
|
||
|
|
self.build = build or settings.ps_build
|
||
|
|
self.jobs = dict(jobs) if jobs is not None else dict(load_jobs())
|
||
|
|
self.tz = ZoneInfo(settings.site_tz)
|
||
|
|
self._connect = connect or (lambda: db.connect_as("worker", settings, autocommit=True))
|
||
|
|
self.stack = Stack()
|
||
|
|
self.host = socket.gethostname()
|
||
|
|
self.pid = os.getpid()
|
||
|
|
self.feed = FeedSettings()
|
||
|
|
self._lock = threading.Lock()
|
||
|
|
self._running: set[str] = set()
|
||
|
|
self._next_trigger: dict[str, str] = {}
|
||
|
|
self._again: dict[str, str] = {}
|
||
|
|
self._last_end: dict[str, datetime] = {}
|
||
|
|
self._db_ready = False
|
||
|
|
self._stop = threading.Event()
|
||
|
|
self._listener: threading.Thread | None = None
|
||
|
|
self.listening = False
|
||
|
|
self.aps = BackgroundScheduler(
|
||
|
|
executors={"default": ThreadPoolExecutor(THREADS)},
|
||
|
|
job_defaults={"coalesce": True, "max_instances": 1, "misfire_grace_time": MISFIRE_GRACE},
|
||
|
|
timezone=self.tz,
|
||
|
|
)
|
||
|
|
self.started = False
|
||
|
|
|
||
|
|
# ------------------------------------------------ small helpers used by runs.run_job
|
||
|
|
|
||
|
|
def connect(self) -> Any:
|
||
|
|
return self._connect()
|
||
|
|
|
||
|
|
def now(self) -> datetime:
|
||
|
|
return datetime.now(self.tz)
|
||
|
|
|
||
|
|
def set_feed(self, feed: FeedSettings) -> None:
|
||
|
|
self.feed = feed
|
||
|
|
|
||
|
|
def running(self) -> list[str]:
|
||
|
|
with self._lock:
|
||
|
|
return sorted(self._running)
|
||
|
|
|
||
|
|
def running_add(self, name: str) -> None:
|
||
|
|
with self._lock:
|
||
|
|
self._running.add(name)
|
||
|
|
|
||
|
|
def running_discard(self, name: str) -> None:
|
||
|
|
with self._lock:
|
||
|
|
self._running.discard(name)
|
||
|
|
|
||
|
|
# ------------------------------------------------ lifecycle
|
||
|
|
|
||
|
|
def start(self) -> None:
|
||
|
|
"""Set up the database side when it is reachable, add the jobs and start the scheduler and listener."""
|
||
|
|
self._setup_db()
|
||
|
|
now = self.now()
|
||
|
|
for name, spec in self.jobs.items():
|
||
|
|
self.aps.add_job(
|
||
|
|
self._fire,
|
||
|
|
trigger=SettingsTrigger(self._seconds_for(spec)),
|
||
|
|
args=(name,),
|
||
|
|
id=name,
|
||
|
|
name=name,
|
||
|
|
next_run_time=self._next_time(spec, now),
|
||
|
|
replace_existing=True,
|
||
|
|
)
|
||
|
|
self.aps.start()
|
||
|
|
self.started = True
|
||
|
|
self._write_schedule(self.jobs)
|
||
|
|
self._stop.clear()
|
||
|
|
self._listener = threading.Thread(target=self._listen, name="ps-worker-listen", daemon=True)
|
||
|
|
self._listener.start()
|
||
|
|
log.info("scheduler started", extra={"ps": {"jobs": sorted(self.jobs), "host": self.host}})
|
||
|
|
|
||
|
|
def shutdown(self, wait: bool = True) -> None:
|
||
|
|
self._stop.set()
|
||
|
|
if self.started:
|
||
|
|
self.aps.shutdown(wait=wait)
|
||
|
|
self.started = False
|
||
|
|
if self._listener is not None:
|
||
|
|
self._listener.join(5.0)
|
||
|
|
self._listener = None
|
||
|
|
|
||
|
|
def _setup_db(self) -> bool:
|
||
|
|
"""Seed the feed settings, mark this host's leftover runs abandoned, load the settings and the last run
|
||
|
|
ends. Safe to call again; returns False when the database is not reachable."""
|
||
|
|
if self._db_ready:
|
||
|
|
return True
|
||
|
|
try:
|
||
|
|
with self.connect() as conn:
|
||
|
|
with conn.transaction():
|
||
|
|
if config.seed(conn):
|
||
|
|
log.info("feed settings seeded with the defaults")
|
||
|
|
self._abandon(conn)
|
||
|
|
self.feed = config.load(conn)
|
||
|
|
for job, end in conn.execute(
|
||
|
|
"SELECT job, max(finished_at) FROM job_runs WHERE finished_at IS NOT NULL"
|
||
|
|
" AND status <> 'skipped' GROUP BY job"
|
||
|
|
).fetchall():
|
||
|
|
if end is not None:
|
||
|
|
self._last_end[job] = end.astimezone(self.tz)
|
||
|
|
except Exception as e:
|
||
|
|
log.warning("worker database setup deferred: %s", redact(f"{type(e).__name__}: {e}"))
|
||
|
|
return False
|
||
|
|
self._db_ready = True
|
||
|
|
return True
|
||
|
|
|
||
|
|
def _abandon(self, conn: Any) -> None:
|
||
|
|
"""Runs left `running` by a dead worker become `abandoned`: this host's always, another host's when no
|
||
|
|
session holds that job's advisory lock (a container recreated with a new host name)."""
|
||
|
|
rows = conn.execute("SELECT id, job, host FROM job_runs WHERE status = 'running'").fetchall()
|
||
|
|
for rid, job, host in rows:
|
||
|
|
if host != self.host:
|
||
|
|
got = conn.execute("SELECT pg_try_advisory_lock(%s, hashtext(%s))", (runs.LOCK_SPACE, job)).fetchone()[
|
||
|
|
0
|
||
|
|
]
|
||
|
|
if not got:
|
||
|
|
continue
|
||
|
|
conn.execute("SELECT pg_advisory_unlock(%s, hashtext(%s))", (runs.LOCK_SPACE, job))
|
||
|
|
conn.execute(
|
||
|
|
"UPDATE job_runs SET status = 'abandoned', finished_at = now(),"
|
||
|
|
" errors = errors || '[\"The worker stopped during this run.\"]'::jsonb"
|
||
|
|
" WHERE id = %s AND status = 'running'",
|
||
|
|
(rid,),
|
||
|
|
)
|
||
|
|
log.warning("job run %s (%s) marked abandoned", rid, job)
|
||
|
|
|
||
|
|
# ------------------------------------------------ timing
|
||
|
|
|
||
|
|
def _seconds_for(self, spec: JobSpec) -> Callable[[datetime], int | None]:
|
||
|
|
def seconds(now: datetime) -> int | None:
|
||
|
|
return spec.cadence(self.feed, now.astimezone(self.tz))
|
||
|
|
|
||
|
|
return seconds
|
||
|
|
|
||
|
|
def _next_time(self, spec: JobSpec, now: datetime) -> datetime:
|
||
|
|
due = next_due(spec, self.feed, now, self._last_end.get(spec.name))
|
||
|
|
return due if due is not None else now + timedelta(seconds=IDLE_RECHECK_SECONDS)
|
||
|
|
|
||
|
|
def reschedule(self, names: Iterable[str] | None = None) -> None:
|
||
|
|
"""Recompute the next run of every job (or of `names`) from the current settings."""
|
||
|
|
if not self.started:
|
||
|
|
return
|
||
|
|
now = self.now()
|
||
|
|
todo = list(names) if names is not None else list(self.jobs)
|
||
|
|
for name in todo:
|
||
|
|
spec = self.jobs.get(name)
|
||
|
|
if spec is None:
|
||
|
|
continue
|
||
|
|
with contextlib.suppress(Exception): # JobLookupError during shutdown
|
||
|
|
self.aps.modify_job(name, next_run_time=self._next_time(spec, now))
|
||
|
|
self._write_schedule(todo)
|
||
|
|
|
||
|
|
def run_soon(self, name: str, trigger: str = "request") -> None:
|
||
|
|
"""Move a job to now. A job that is running now runs again right after."""
|
||
|
|
if name not in self.jobs:
|
||
|
|
log.warning("ignoring a run request for unknown job %r", name[:60])
|
||
|
|
return
|
||
|
|
with self._lock:
|
||
|
|
if name in self._running:
|
||
|
|
self._again[name] = trigger
|
||
|
|
return
|
||
|
|
self._next_trigger[name] = trigger
|
||
|
|
if self.started:
|
||
|
|
with contextlib.suppress(Exception):
|
||
|
|
self.aps.modify_job(name, next_run_time=self.now())
|
||
|
|
|
||
|
|
def after_run(self, name: str, status: str, pending: bool) -> None:
|
||
|
|
"""Called by run_job when a run ends: chain, run again for new requests, else time the next run from now."""
|
||
|
|
self._last_end[name] = self.now()
|
||
|
|
spec = self.jobs[name]
|
||
|
|
with self._lock:
|
||
|
|
again = self._again.pop(name, None)
|
||
|
|
if status in ("ok", "degraded"):
|
||
|
|
for c in spec.chain:
|
||
|
|
self.run_soon(c, "chain")
|
||
|
|
if pending or again:
|
||
|
|
self.run_soon(name, again or "request")
|
||
|
|
self._write_schedule([name])
|
||
|
|
else:
|
||
|
|
self.reschedule([name])
|
||
|
|
|
||
|
|
def _fire(self, name: str) -> None:
|
||
|
|
with self._lock:
|
||
|
|
trigger = self._next_trigger.pop(name, "schedule")
|
||
|
|
spec = self.jobs[name]
|
||
|
|
if not self._db_ready and not self._setup_db():
|
||
|
|
log.warning("job %s not run: database not reachable", name)
|
||
|
|
return
|
||
|
|
if trigger == "schedule" and spec.cadence(self.feed, self.now()) is None:
|
||
|
|
self._write_schedule([name])
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
runs.run_job(self, name, trigger)
|
||
|
|
except Exception as e:
|
||
|
|
log.error("job %s could not run: %s", name, redact(f"{type(e).__name__}: {e}"))
|
||
|
|
|
||
|
|
def _write_schedule(self, names: Iterable[str]) -> None:
|
||
|
|
"""job_schedule rows for the Admin and Clerical screens. Best effort."""
|
||
|
|
if not self.started:
|
||
|
|
return
|
||
|
|
rows = []
|
||
|
|
now = self.now()
|
||
|
|
for name in names:
|
||
|
|
spec = self.jobs.get(name)
|
||
|
|
if spec is None:
|
||
|
|
continue
|
||
|
|
job = self.aps.get_job(name)
|
||
|
|
nxt = job.next_run_time if job is not None else None
|
||
|
|
secs = spec.cadence(self.feed, now)
|
||
|
|
if secs is None:
|
||
|
|
nxt = None
|
||
|
|
rows.append((name, nxt, secs, bool(spec.pausable and self.feed.paused)))
|
||
|
|
if not rows:
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
with self.connect() as conn:
|
||
|
|
for r in rows:
|
||
|
|
conn.execute(
|
||
|
|
"INSERT INTO job_schedule (job, next_run_at, interval_seconds, paused, updated_at)"
|
||
|
|
" VALUES (%s, %s, %s, %s, now()) ON CONFLICT (job) DO UPDATE SET"
|
||
|
|
" next_run_at = EXCLUDED.next_run_at, interval_seconds = EXCLUDED.interval_seconds,"
|
||
|
|
" paused = EXCLUDED.paused, updated_at = now()",
|
||
|
|
r,
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
log.warning("job_schedule not written: %s", type(e).__name__)
|
||
|
|
|
||
|
|
# ------------------------------------------------ settings
|
||
|
|
|
||
|
|
def reload_settings(self, conn: Any | None = None) -> None:
|
||
|
|
try:
|
||
|
|
if conn is not None:
|
||
|
|
self.feed = config.load(conn)
|
||
|
|
else:
|
||
|
|
with self.connect() as c:
|
||
|
|
self.feed = config.load(c)
|
||
|
|
except Exception as e:
|
||
|
|
log.warning("feed settings not reloaded: %s", type(e).__name__)
|
||
|
|
return
|
||
|
|
self.reschedule()
|
||
|
|
|
||
|
|
# ------------------------------------------------ listener
|
||
|
|
|
||
|
|
def _on_connect(self, conn: Any) -> None:
|
||
|
|
"""After every (re)connect: catch up on what was missed while not listening."""
|
||
|
|
self._setup_db()
|
||
|
|
self.reload_settings(conn)
|
||
|
|
for (job,) in conn.execute("SELECT DISTINCT job FROM job_requests WHERE claimed_at IS NULL").fetchall():
|
||
|
|
self.run_soon(job, "request")
|
||
|
|
|
||
|
|
def _notified(self, channel: str, payload: str) -> None:
|
||
|
|
if channel == "ps_settings":
|
||
|
|
self.reload_settings()
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
data = json.loads(payload)
|
||
|
|
except ValueError:
|
||
|
|
log.warning("ignoring a non-JSON ps_job_request payload")
|
||
|
|
return
|
||
|
|
job = data.get("job") if isinstance(data, dict) else None
|
||
|
|
if isinstance(job, str):
|
||
|
|
self.run_soon(job, "request")
|
||
|
|
|
||
|
|
def _listen(self) -> None:
|
||
|
|
delay = 1.0
|
||
|
|
while not self._stop.is_set():
|
||
|
|
conn = None
|
||
|
|
try:
|
||
|
|
conn = self.connect()
|
||
|
|
for ch in CHANNELS:
|
||
|
|
conn.execute(f"LISTEN {ch}")
|
||
|
|
self.listening = True
|
||
|
|
delay = 1.0
|
||
|
|
log.info("worker listening on %s", ", ".join(CHANNELS))
|
||
|
|
self._on_connect(conn)
|
||
|
|
while not self._stop.is_set():
|
||
|
|
for n in conn.notifies(timeout=1.0):
|
||
|
|
self._notified(n.channel, n.payload)
|
||
|
|
except Exception as e:
|
||
|
|
if not self._stop.is_set():
|
||
|
|
log.warning("worker listener: %s; retrying in %.0f s", type(e).__name__, delay)
|
||
|
|
finally:
|
||
|
|
self.listening = False
|
||
|
|
if conn is not None:
|
||
|
|
with contextlib.suppress(Exception):
|
||
|
|
conn.close()
|
||
|
|
if self._stop.wait(delay):
|
||
|
|
break
|
||
|
|
delay = min(delay * 2, BACKOFF_MAX)
|
||
|
|
|
||
|
|
# ------------------------------------------------ heartbeat loop (main thread)
|
||
|
|
|
||
|
|
def serve(self, beat_file: Path, stop: threading.Event, interval: float) -> None:
|
||
|
|
"""Beat until `stop`: the beat file and the worker_heartbeat row."""
|
||
|
|
conn = None
|
||
|
|
while not stop.is_set():
|
||
|
|
try:
|
||
|
|
heartbeat.touch(beat_file)
|
||
|
|
except OSError as e:
|
||
|
|
log.error("can't touch beat file %s: %s", beat_file, e.strerror)
|
||
|
|
try:
|
||
|
|
if conn is None or conn.closed:
|
||
|
|
conn = self.connect()
|
||
|
|
heartbeat.write_row(conn, build=self.build, host=self.host, pid=self.pid, running=self.running())
|
||
|
|
except Exception as e:
|
||
|
|
log.warning("heartbeat row not written: %s", type(e).__name__)
|
||
|
|
if conn is not None:
|
||
|
|
with contextlib.suppress(Exception):
|
||
|
|
conn.close()
|
||
|
|
conn = None
|
||
|
|
stop.wait(interval)
|
||
|
|
if conn is not None:
|
||
|
|
with contextlib.suppress(Exception):
|
||
|
|
conn.close()
|