""" Supervised background-task registration for Datasette core. Plugins that need long-lived background work (a polling loop, a queue consumer, a scheduled job runner) register it with ``datasette.add_background_task(func, name=None)`` - typically from a ``startup`` plugin hook - instead of fire-and-forgetting their own ``asyncio.create_task()``. Core owns: - **references**: every launched ``asyncio.Task`` is kept alive on a :class:`BackgroundTaskSupervisor`, so it can never be silently garbage collected the way an unreferenced ``create_task()`` call can be; - **launch timing**: registered work is buffered until :meth:`BackgroundTaskSupervisor.launch_all` runs, which core arranges to happen only after *every* plugin's ``startup`` hook has finished - so a task that depends on another plugin having registered something first doesn't need ``tryfirst=True`` ordering tricks; - **crash surfacing**: an unhandled exception in a background task is logged with its full traceback to the ``datasette.background_tasks`` logger and recorded on the handle, instead of becoming an "Task exception was never retrieved" warning nobody sees; - **cancellation**: :meth:`BackgroundTaskSupervisor.cancel_all` cancels every task still running and waits (with a grace period) for them to actually stop. """ from __future__ import annotations import asyncio import datetime import functools import logging from collections.abc import Awaitable, Callable logger = logging.getLogger("datasette.background_tasks") def _utcnow_iso() -> str: return datetime.datetime.now(datetime.timezone.utc).isoformat() def _function_path(func: Callable) -> str: """Describe the callable without guessing which plugin registered it.""" while isinstance(func, functools.partial): func = func.func if not hasattr(func, "__qualname__"): func = type(func).__call__ return f"{func.__module__}.{func.__qualname__}" class BackgroundTask: """A handle to a single piece of supervised background work. States: ``registered`` (added but not yet launched) -> ``running`` -> one of ``completed`` (returned cleanly), ``crashed`` (raised an exception other than ``CancelledError`` - see ``.exception``), or ``cancelled`` (``.cancel()`` was called, or it was still running at shutdown). """ def __init__( self, name: str, func: Callable[[object], Awaitable[None]], ): self.name = name self.state = "registered" self.task: asyncio.Task | None = None self.exception: BaseException | None = None self.started_at: str | None = None self.function = _function_path(func) self._func = func self._supervisor: BackgroundTaskSupervisor | None = None def cancel(self) -> None: """Cancel this task. If it has already been launched, cancels the underlying ``asyncio.Task`` - its state becomes ``cancelled`` once the cancellation is observed (asynchronously, via the task's done callback). If it has not been launched yet, this is a no-op as far as asyncio is concerned (there's no task to cancel) but it deregisters the handle from its supervisor so it never runs. """ if self.task is not None: self.task.cancel() elif self._supervisor is not None: self._supervisor._deregister(self) def __repr__(self) -> str: return f"" class BackgroundTaskSupervisor: """Owns registration and launch of every :class:`BackgroundTask` for a single ``Datasette`` instance. Registration (:meth:`add`) is separate from launch (:meth:`launch_all`): plugins register work whenever convenient (typically from a ``startup`` hook, but request handlers can register dynamic per-job work too), and it either sits buffered until :meth:`launch_all` runs, or - if :meth:`launch_all` has already run - starts immediately. Strong references to every :class:`BackgroundTask` (and its ``asyncio.Task``) are kept for the life of the instance, by design - that's what makes the enrichments-style "fire-and-forget task gets garbage collected mid-flight" bug impossible here. There is currently no pruning of completed/crashed/cancelled tasks, so a plugin that dynamically registers many short-lived tasks over a long process lifetime (a per-job registration pattern, e.g. one task per queued job) will grow this list without bound. That's an accepted v1 trade-off in favour of full introspection (``/-/tasks``); revisit with a pruning or capping policy if unbounded growth is reported in practice. """ def __init__(self, datasette): self._datasette = datasette self._tasks: list[BackgroundTask] = [] self._names = set() self._launched = False self._lock = asyncio.Lock() def add(self, func, name=None) -> BackgroundTask: base_name = name or getattr(func, "__qualname__", None) or repr(func) actual_name = self._unique_name(base_name) handle = BackgroundTask(actual_name, func) handle._supervisor = self self._tasks.append(handle) self._names.add(actual_name) if self._launched: self._launch_one(handle) return handle def _unique_name(self, base_name: str) -> str: if base_name not in self._names: return base_name n = 2 while f"{base_name}-{n}" in self._names: n += 1 return f"{base_name}-{n}" def _deregister(self, handle: BackgroundTask) -> None: try: self._tasks.remove(handle) except ValueError: pass self._names.discard(handle.name) def _launch_one(self, handle: BackgroundTask) -> None: handle.state = "running" handle.started_at = _utcnow_iso() handle.task = asyncio.create_task( handle._func(self._datasette), name=handle.name ) handle.task.add_done_callback(functools.partial(_on_task_done, handle)) async def launch_all(self) -> None: """Launch every currently-registered task that hasn't launched yet. Idempotent and safe to call concurrently: subsequent (or racing) calls are no-ops once the first has set ``self._launched``. """ if self._launched: return async with self._lock: if self._launched: return self._launched = True for handle in list(self._tasks): if handle.task is None: self._launch_one(handle) async def cancel_all(self, grace: float = 5.0) -> None: """Cancel every task that isn't already done, then wait up to ``grace`` seconds for them to actually finish. Stragglers still running after that are logged by name (but left to finish or not on their own - this does not forcibly kill them, asyncio has no mechanism for that). """ handles_by_task = { handle.task: handle for handle in self._tasks if handle.task is not None } pending = [task for task in handles_by_task if not task.done()] for task in pending: task.cancel() if not pending: return _done, not_done = await asyncio.wait(pending, timeout=grace) if not_done: names = sorted(handles_by_task[task].name for task in not_done) logger.warning( "%d background task(s) did not finish within the %.1fs grace " "period after cancellation: %s", len(names), grace, ", ".join(names), ) def tasks(self) -> list[BackgroundTask]: """Return every registered :class:`BackgroundTask`, launched or not, in registration order. Used by the ``/-/tasks`` debug endpoint. """ return list(self._tasks) @property def launched(self) -> bool: """Whether :meth:`launch_all` has run yet - lets ``/-/tasks`` distinguish "no tasks registered" from "tasks registered but nothing has armed the launch yet" without reaching for the private ``_launched`` attribute. """ return self._launched def _on_task_done(handle: BackgroundTask, task: asyncio.Task) -> None: if task.cancelled(): handle.state = "cancelled" return exc = task.exception() if exc is not None: handle.state = "crashed" handle.exception = exc logger.error("Background task %r crashed", handle.name, exc_info=exc) return handle.state = "completed"