131 lines
4.4 KiB
Python
131 lines
4.4 KiB
Python
"""Plugin background jobs on the shared arq worker (SPEC §11.2).
|
|||
|
|
|
||
|
|
Each plugin ``JobSpec`` is registered as an arq function under a namespaced name
|
||
|
|
``plg.{plugin}.{func}`` (enqueue on demand from hooks/routes via ``sdk.enqueue``),
|
||
|
|
and, when it has a ``schedule``, also as an arq cron job. ``singleton=True`` wraps
|
||
|
|
the job in a Redis NX lock keyed by (plugin, func, user) so at most one instance
|
||
|
|
runs per user (or globally for scheduleds with no user). Jobs share the worker
|
||
|
|
pool + retry semantics with core jobs and log under the ``plg`` logger.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from arq import cron
|
||
|
|
from arq.cron import CronJob
|
||
|
|
from arq.worker import Function, func
|
||
|
|
|
||
|
|
from app.plugins import loader
|
||
|
|
from app.plugins.types import JobSpec, Plugin
|
||
|
|
from app.redis_client import get_redis
|
||
|
|
|
||
|
|
logger = logging.getLogger("plg")
|
||
|
|
|
||
|
|
# Safety expiry on the singleton lock so a crashed job can't wedge a user forever.
|
||
|
|
# A job that runs longer than this could overlap; acceptable for M3.
|
||
|
|
_LOCK_TTL = 3600
|
||
|
|
|
||
|
|
# Standard 5-field cron order and each field's inclusive value bounds.
|
||
|
|
_FIELD_ORDER = ("minute", "hour", "day", "month", "weekday")
|
||
|
|
_FIELD_BOUNDS = {
|
||
|
|
"minute": (0, 59),
|
||
|
|
"hour": (0, 23),
|
||
|
|
"day": (1, 31),
|
||
|
|
"month": (1, 12),
|
||
|
|
"weekday": (0, 6),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def job_name(plugin_name: str, spec: JobSpec) -> str:
|
||
|
|
"""The arq function name a plugin job is registered under."""
|
||
|
|
return f"plg.{plugin_name}.{spec.func.__name__}"
|
||
|
|
|
||
|
|
|
||
|
|
def _lock_key(plugin_name: str, func_name: str, user: object) -> str:
|
||
|
|
return f"plglock:{plugin_name}:{func_name}:{user}"
|
||
|
|
|
||
|
|
|
||
|
|
def _make_runner(plugin: Plugin, spec: JobSpec) -> tuple[Any, str]:
|
||
|
|
"""Wrap a job func with the singleton lock (when requested) + namespaced name."""
|
||
|
|
func_name = spec.func.__name__
|
||
|
|
name = job_name(plugin.name, spec)
|
||
|
|
|
||
|
|
async def runner(ctx: dict[str, Any], *args: Any, **kwargs: Any) -> Any:
|
||
|
|
if not spec.singleton:
|
||
|
|
return await spec.func(ctx, *args, **kwargs)
|
||
|
|
user = kwargs.get("user_id", "global")
|
||
|
|
key = _lock_key(plugin.name, func_name, user)
|
||
|
|
acquired = await get_redis().set(key, b"1", nx=True, ex=_LOCK_TTL)
|
||
|
|
if not acquired:
|
||
|
|
logger.info("%s skipped: singleton already running (user=%s)", name, user)
|
||
|
|
return None
|
||
|
|
try:
|
||
|
|
return await spec.func(ctx, *args, **kwargs)
|
||
|
|
finally:
|
||
|
|
await get_redis().delete(key)
|
||
|
|
|
||
|
|
return runner, name
|
||
|
|
|
||
|
|
|
||
|
|
def plugin_arq_functions() -> list[Function]:
|
||
|
|
"""Every plugin job as an arq function (enqueue-on-demand), namespaced."""
|
||
|
|
funcs: list[Function] = []
|
||
|
|
for plugin in loader.get_plugins():
|
||
|
|
for spec in plugin.jobs:
|
||
|
|
runner, name = _make_runner(plugin, spec)
|
||
|
|
funcs.append(func(runner, name=name))
|
||
|
|
return funcs
|
||
|
|
|
||
|
|
|
||
|
|
def _parse_field(token: str, lo: int, hi: int) -> set[int] | None:
|
||
|
|
"""Parse one cron field to a value set (``*`` -> None = every)."""
|
||
|
|
if token == "*":
|
||
|
|
return None
|
||
|
|
values: set[int] = set()
|
||
|
|
for part in token.split(","):
|
||
|
|
if part.startswith("*/"):
|
||
|
|
values.update(range(lo, hi + 1, int(part[2:])))
|
||
|
|
elif "-" in part:
|
||
|
|
start, end = part.split("-", 1)
|
||
|
|
values.update(range(int(start), int(end) + 1))
|
||
|
|
else:
|
||
|
|
values.add(int(part))
|
||
|
|
return values
|
||
|
|
|
||
|
|
|
||
|
|
def parse_cron(expr: str) -> dict[str, set[int] | None]:
|
||
|
|
"""Parse a standard 5-field cron string into arq ``cron`` field kwargs."""
|
||
|
|
parts = expr.split()
|
||
|
|
if len(parts) != 5:
|
||
|
|
raise ValueError(f"cron expression must have 5 fields (min hour dom mon dow): {expr!r}")
|
||
|
|
return {
|
||
|
|
name: _parse_field(token, *_FIELD_BOUNDS[name])
|
||
|
|
for name, token in zip(_FIELD_ORDER, parts, strict=True)
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def plugin_cron_jobs() -> list[CronJob]:
|
||
|
|
"""Cron-scheduled plugin jobs as arq ``CronJob``s (run on the shared worker)."""
|
||
|
|
jobs: list[CronJob] = []
|
||
|
|
for plugin in loader.get_plugins():
|
||
|
|
for spec in plugin.jobs:
|
||
|
|
if spec.schedule is None:
|
||
|
|
continue
|
||
|
|
runner, name = _make_runner(plugin, spec)
|
||
|
|
f = parse_cron(spec.schedule)
|
||
|
|
jobs.append(
|
||
|
|
cron(
|
||
|
|
runner,
|
||
|
|
name=f"{name}.cron",
|
||
|
|
second=0,
|
||
|
|
minute=f["minute"],
|
||
|
|
hour=f["hour"],
|
||
|
|
day=f["day"],
|
||
|
|
month=f["month"],
|
||
|
|
weekday=f["weekday"],
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return jobs
|