datasette/datasette/background_tasks.py

"""
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"<BackgroundTask name={self.name!r} state={self.state!r}>"


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"