From c910b8849a99caaaa5ea0f583a727d32d9c11112 Mon Sep 17 00:00:00 2001 From: Alex Garcia Date: Thu, 30 Jul 2026 17:56:42 -0700 Subject: [PATCH] Add datasette.add_background_task() with supervised launch after startup Co-Authored-By: Claude Fable 5 --- datasette/__init__.py | 1 + datasette/app.py | 68 +++++- datasette/background_tasks.py | 259 ++++++++++++++++++++++ datasette/cli.py | 6 + tests/test_background_tasks.py | 382 +++++++++++++++++++++++++++++++++ tests/test_cli_serve_get.py | 50 +++++ 6 files changed, 764 insertions(+), 2 deletions(-) create mode 100644 datasette/background_tasks.py create mode 100644 tests/test_background_tasks.py diff --git a/datasette/__init__.py b/datasette/__init__.py index e0022178..982dcc79 100644 --- a/datasette/__init__.py +++ b/datasette/__init__.py @@ -1,6 +1,7 @@ from datasette.permissions import Permission # noqa from datasette.version import __version_info__, __version__ # noqa from datasette.events import Event # noqa +from datasette.background_tasks import BackgroundTask, BackgroundTaskSupervisor # noqa from datasette.tokens import TokenHandler, TokenInvalid, TokenRestrictions # noqa from datasette.utils.asgi import ( # noqa Forbidden, diff --git a/datasette/app.py b/datasette/app.py index 42be7425..cf79df1f 100644 --- a/datasette/app.py +++ b/datasette/app.py @@ -42,6 +42,7 @@ from jinja2.exceptions import TemplateNotFound from markupsafe import Markup, escape from . import stored_queries, write_sql +from .background_tasks import BackgroundTask, BackgroundTaskSupervisor from .column_types import SQLiteType from .csrf import CrossOriginProtectionMiddleware from .database import Database, QueryInterrupted @@ -454,6 +455,7 @@ class Datasette: self.actions = {} # .invoke_startup() will populate this self._column_types = {} # .invoke_startup() will populate this self._setup_db_done = False + self._suppress_background_tasks = False try: self._refresh_schemas_lock = asyncio.Lock() self._startup_lock = asyncio.Lock() @@ -467,6 +469,7 @@ class Datasette: self._startup_lock = asyncio.Lock() else: raise + self._background_tasks = BackgroundTaskSupervisor(self) self.crossdb = crossdb self.nolock = nolock if memory or crossdb or not self.files: @@ -2836,6 +2839,64 @@ class Datasette: self._setup_db_done = True await self.invoke_startup() + def add_background_task(self, func, name=None) -> BackgroundTask: + """Register a piece of supervised background work, typically from + a plugin's ``startup`` hook. + + ``func`` must be a coroutine function taking one positional + argument, the ``Datasette`` instance - core calls ``func(self)``. + Callable any time after ``__init__``: if background tasks haven't + launched yet (the common case - most callers are ``startup`` hooks, + which run before launch), this buffers the registration until they + do; if they've already launched (e.g. called from a request + handler after the server is up), the task starts immediately. + + Returns a :class:`~datasette.background_tasks.BackgroundTask` + handle (``.name``, ``.state``, ``.task``, ``.exception``, + ``.started_at``, ``.plugin``, ``.cancel()``). + + ``name`` defaults to ``func.__qualname__``; on a name collision a + ``-2``, ``-3``, ... suffix is appended, since names are how + ``/-/tasks`` and log messages identify work. + """ + return self._background_tasks.add(func, name=name) + + async def start_background_tasks(self): + """Run startup (if it hasn't run yet) and launch every registered + background task. + + Public entry point for tests, embedders, and headless CLIs (the + ``datasette-rss``-style ``fetch --due`` shape) that want supervised + background tasks without running a server - equivalent to what + happens automatically via ASGI lifespan / the first-request + fallback in a served deployment. + """ + await self.invoke_startup() + await self._background_tasks.launch_all() + + async def _launch_background_tasks(self): + """Idempotently launch every registered background task. Private: + this is the entry point wired into the lifecycle trigger lists + (the second entry in both ``AsgiLifespan`` and + ``AsgiRunOnFirstRequest``'s ``on_startup``, after + ``_startup_sequence``) - not something plugins or embedders should + call directly; use ``add_background_task`` / + ``start_background_tasks`` instead. + + Positioned after ``_startup_sequence`` in both trigger lists so + launch always happens once every plugin's ``startup`` hook has had + a chance to register work - the ordering guarantee that makes + ``add_background_task`` useful. No-ops when + ``_suppress_background_tasks`` is set (the ``--get`` CLI path: its + one-shot TestClient request flows through the full ASGI stack, + including the first-request fallback, but must never launch + long-lived background work per decision #3 in + ``plans/first-request/04-core-plan.md``). + """ + if self._suppress_background_tasks: + return + await self._background_tasks.launch_all() + def app(self): """Returns an ASGI app function that serves the whole of Datasette""" routes = self._routes() @@ -2848,10 +2909,13 @@ class Datasette: asgi = AsgiTracer(asgi) asgi = AsgiLifespan( asgi, - on_startup=[self._startup_sequence], + on_startup=[self._startup_sequence, self._launch_background_tasks], on_shutdown=[_close_on_shutdown], ) - asgi = AsgiRunOnFirstRequest(asgi, on_startup=[self._startup_sequence]) + asgi = AsgiRunOnFirstRequest( + asgi, + on_startup=[self._startup_sequence, self._launch_background_tasks], + ) for wrapper in pm.hook.asgi_wrapper(datasette=self): asgi = wrapper(asgi) return asgi diff --git a/datasette/background_tasks.py b/datasette/background_tasks.py new file mode 100644 index 00000000..f06bd9b3 --- /dev/null +++ b/datasette/background_tasks.py @@ -0,0 +1,259 @@ +""" +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. + +See ``plans/first-request/04-core-plan.md`` (decisions #2-#5) for the +design rationale. +""" + +from __future__ import annotations + +import asyncio +import datetime +import functools +import inspect +import logging +from typing import Awaitable, Callable, List, Optional + +logger = logging.getLogger("datasette.background_tasks") + + +def _utcnow_iso() -> str: + return datetime.datetime.now(datetime.timezone.utc).isoformat() + + +def _resolve_plugin_name(func: Callable) -> Optional[str]: + """Best-effort, cheap attempt to work out which registered plugin a + background-task function belongs to, for the ``.plugin`` field on + :class:`BackgroundTask` (used by ``/-/tasks`` and logs). + + This matches ``func``'s module against every currently-registered + pluggy plugin's module - the same module a plugin's ``startup`` hook + implementation lives in, in the overwhelmingly common case where + ``add_background_task`` is called directly from (or a couple of + frames below) that hook. It deliberately does *not* walk the call + stack or otherwise try harder: this is a nice-to-have for + introspection, not something worth building heavy machinery for, and + returning ``None`` when it can't tell is a fine fallback. + """ + try: + from .plugins import pm + + module = inspect.getmodule(func) + if module is None: + return None + module_name = getattr(module, "__name__", None) + if not module_name: + return None + for plugin in pm.get_plugins(): + plugin_module = ( + plugin if inspect.ismodule(plugin) else inspect.getmodule(plugin) + ) + if plugin_module is None: + continue + plugin_module_name = getattr(plugin_module, "__name__", None) + if not plugin_module_name: + continue + if module_name == plugin_module_name or module_name.startswith( + plugin_module_name + "." + ): + return pm.get_name(plugin) + except Exception: + # Never let plugin-name resolution break task registration. + return None + return None + + +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]], + plugin: Optional[str] = None, + ): + self.name = name + self.state = "registered" + self.task: Optional[asyncio.Task] = None + self.exception: Optional[BaseException] = None + self.started_at: Optional[str] = None + self.plugin = plugin + self._func = func + self._supervisor: Optional["BackgroundTaskSupervisor"] = 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) + plugin = _resolve_plugin_name(func) + handle = BackgroundTask(actual_name, func, plugin=plugin) + 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) + + +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" diff --git a/datasette/cli.py b/datasette/cli.py index 2694c1f6..7254fe2f 100644 --- a/datasette/cli.py +++ b/datasette/cli.py @@ -678,6 +678,12 @@ def serve( except StartupError as e: raise click.ClickException(e.args[0]) + # --get never launches background tasks (decision #3 in + # plans/first-request/04-core-plan.md): TestClient's request below + # flows through the full ASGI stack, including the + # AsgiRunOnFirstRequest fallback, which would otherwise launch them. + ds._suppress_background_tasks = True + client = TestClient(ds) request_headers = {} if token: diff --git a/tests/test_background_tasks.py b/tests/test_background_tasks.py new file mode 100644 index 00000000..39dba968 --- /dev/null +++ b/tests/test_background_tasks.py @@ -0,0 +1,382 @@ +""" +Tests for datasette.add_background_task() / start_background_tasks() and the +BackgroundTask / BackgroundTaskSupervisor machinery in +datasette/background_tasks.py, per plans/first-request/04-core-plan.md +(decisions #2-#5) and todos/first-request/03-background-tasks-api.md. +""" + +import asyncio +import contextlib +import logging + +import httpx +import pytest + +from datasette import hookimpl +from datasette.app import Datasette +from datasette.plugins import pm + + +async def _drive_lifespan_startup(app): + """Send a single lifespan.startup message into app's ASGI lifespan loop + and return the list of messages sent back, without ever sending + lifespan.shutdown. Copied from tests/test_lifespan.py's helper of the + same name - mirrors what a real server does: after startup completes + it parks waiting for the next event, and we cancel that wait once + we've observed the startup response. + """ + messages_sent = [] + startup_responded = asyncio.Event() + delivered = False + + async def receive(): + nonlocal delivered + if not delivered: + delivered = True + return {"type": "lifespan.startup"} + await asyncio.Event().wait() + + async def send(message): + messages_sent.append(message) + startup_responded.set() + + task = asyncio.create_task(app({"type": "lifespan"}, receive, send)) + try: + await asyncio.wait_for(startup_responded.wait(), timeout=5) + finally: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + return messages_sent + + +@pytest.mark.asyncio +async def test_tasks_registered_in_startup_hook_run_after_lifespan_startup(): + # Two tasks registered by one plugin's startup hook - order preserved, + # both running after lifespan startup completes, and no HTTP request + # of any kind is issued anywhere in this test. + events = [] + + async def task_one(datasette): + events.append("task_one") + await asyncio.Event().wait() + + async def task_two(datasette): + events.append("task_two") + await asyncio.Event().wait() + + class TwoTaskPlugin: + __name__ = "TwoTaskPlugin" + + @hookimpl + def startup(self, datasette): + async def inner(): + datasette.add_background_task(task_one, name="task-one") + datasette.add_background_task(task_two, name="task-two") + + return inner + + ds = Datasette(memory=True) + pm.register(TwoTaskPlugin(), name="two_task_plugin") + try: + app = ds.app() + messages = await _drive_lifespan_startup(app) + assert {"type": "lifespan.startup.complete"} in messages + + handles = ds._background_tasks.tasks() + assert [h.name for h in handles] == ["task-one", "task-two"] + + # Let both tasks run their first line of code. + await asyncio.sleep(0) + assert handles[0].state == "running" + assert handles[1].state == "running" + assert events == ["task_one", "task_two"] + finally: + pm.unregister(name="two_task_plugin") + await ds._background_tasks.cancel_all(grace=1.0) + + +@pytest.mark.asyncio +async def test_launch_waits_for_every_startup_hook_before_running_any_task(): + # PluginA registers a task from its startup hook; PluginB does the + # same from ITS startup hook, which runs after PluginA's (forced with + # tryfirst=True on A). Even though A's registration happens first, + # A's task body must not actually execute until every startup hook - + # including B's - has finished, since launch only happens after + # invoke_startup() completes. This is the ordering guarantee that + # dissolves datasette-cron's tryfirst=True launch hack. + hook_call_order = [] + seen_names_when_a_ran = {} + + async def task_a(datasette): + seen_names_when_a_ran["names"] = [ + h.name for h in datasette._background_tasks.tasks() + ] + + async def task_b(datasette): + pass + + class PluginA: + __name__ = "PluginA" + + @hookimpl(tryfirst=True) + def startup(self, datasette): + async def inner(): + hook_call_order.append("A") + datasette.add_background_task(task_a, name="task-a") + + return inner + + class PluginB: + __name__ = "PluginB" + + @hookimpl + def startup(self, datasette): + async def inner(): + hook_call_order.append("B") + datasette.add_background_task(task_b, name="task-b") + + return inner + + ds = Datasette(memory=True) + pm.register(PluginA(), name="plugin_a") + pm.register(PluginB(), name="plugin_b") + try: + await ds.start_background_tasks() + # Confirm A's startup hook really did run (and register task-a) + # strictly before B's startup hook ran. + assert hook_call_order == ["A", "B"] + + handles = ds._background_tasks.tasks() + await asyncio.wait_for( + asyncio.gather(*[h.task for h in handles]), timeout=5 + ) + # Yet by the time task-a's own body executed (after launch, which + # only happens once every startup hook - including B's - has + # finished), task-b was already registered. + assert "task-b" in seen_names_when_a_ran["names"] + finally: + pm.unregister(name="plugin_a") + pm.unregister(name="plugin_b") + + +@pytest.mark.asyncio +async def test_concurrent_first_requests_launch_background_tasks_exactly_once(): + launch_count = {"n": 0} + + async def counting_task(datasette): + launch_count["n"] += 1 + + class CountingTaskPlugin: + __name__ = "CountingTaskPlugin" + + @hookimpl + def startup(self, datasette): + async def inner(): + datasette.add_background_task(counting_task, name="counting-task") + + return inner + + ds = Datasette(memory=True) + pm.register(CountingTaskPlugin(), name="counting_task_plugin") + try: + app = ds.app() + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://localhost" + ) as client: + responses = await asyncio.gather( + *[client.get("/-/versions.json") for _ in range(10)] + ) + assert all(response.status_code == 200 for response in responses) + + handles = ds._background_tasks.tasks() + assert len(handles) == 1 + await asyncio.wait_for(handles[0].task, timeout=5) + assert launch_count["n"] == 1 + finally: + pm.unregister(name="counting_task_plugin") + + +@pytest.mark.asyncio +async def test_post_launch_registration_starts_immediately_and_cancel_works(): + ds = Datasette(memory=True) + await ds.start_background_tasks() # nothing registered yet, but launched + + started = asyncio.Event() + + async def long_running(datasette): + started.set() + await asyncio.Event().wait() + + handle = ds.add_background_task(long_running, name="dynamic-task") + # Registered after launch: starts immediately rather than sitting in + # "registered" limbo. + assert handle.state == "running" + assert handle.task is not None + + await asyncio.wait_for(started.wait(), timeout=5) + assert handle.state == "running" + + handle.cancel() + with pytest.raises(asyncio.CancelledError): + await handle.task + await asyncio.sleep(0) + assert handle.state == "cancelled" + + +@pytest.mark.asyncio +async def test_pre_launch_registration_starts_as_registered(): + ds = Datasette(memory=True) + + async def task(datasette): + pass + + handle = ds.add_background_task(task, name="buffered-task") + assert handle.state == "registered" + assert handle.task is None + + handle.cancel() # not yet launched: deregisters instead of cancelling + assert handle not in ds._background_tasks.tasks() + + +@pytest.mark.asyncio +async def test_crashing_task_logs_traceback_and_state_is_crashed(caplog): + ds = Datasette(memory=True) + await ds.start_background_tasks() + + survivor_ran = asyncio.Event() + + async def crasher(datasette): + raise RuntimeError("kaboom") + + async def survivor(datasette): + survivor_ran.set() + + with caplog.at_level(logging.ERROR, logger="datasette.background_tasks"): + crash_handle = ds.add_background_task(crasher, name="crasher") + survivor_handle = ds.add_background_task(survivor, name="survivor") + await asyncio.wait_for( + asyncio.gather( + crash_handle.task, survivor_handle.task, return_exceptions=True + ), + timeout=5, + ) + + assert crash_handle.state == "crashed" + assert isinstance(crash_handle.exception, RuntimeError) + assert str(crash_handle.exception) == "kaboom" + + # The crash must not affect any other task. + assert survivor_ran.is_set() + assert survivor_handle.state == "completed" + + assert "crasher" in caplog.text + assert "kaboom" in caplog.text + assert "Traceback" in caplog.text + assert "RuntimeError" in caplog.text + + +def test_name_collisions_get_suffixed_and_explicit_names_are_respected(): + ds = Datasette(memory=True) + + async def noop(datasette): + pass + + async def another_noop(datasette): + pass + + h1 = ds.add_background_task(noop, name="dup") + h2 = ds.add_background_task(another_noop, name="dup") + h3 = ds.add_background_task(noop, name="dup") + assert [h1.name, h2.name, h3.name] == ["dup", "dup-2", "dup-3"] + + h_explicit = ds.add_background_task(noop, name="explicit-name") + assert h_explicit.name == "explicit-name" + + h_default = ds.add_background_task(noop) + assert h_default.name == noop.__qualname__ + + +@pytest.mark.asyncio +async def test_start_background_tasks_on_bare_datasette(): + # The headless-CLI path (datasette-rss's `fetch --due` shape): no + # server, no lifespan, no first HTTP request - just an explicit call. + ran = asyncio.Event() + + async def task(datasette): + ran.set() + + ds = Datasette([]) + assert ds._startup_invoked is False + + handle = ds.add_background_task(task, name="headless-task") + assert handle.state == "registered" + + await ds.start_background_tasks() + + assert ds._startup_invoked is True + await asyncio.wait_for(ran.wait(), timeout=5) + await asyncio.wait_for(handle.task, timeout=5) + # handle.task being done only guarantees the coroutine has returned, + # not that our done-callback (which updates handle.state) has run yet - + # asyncio schedules done-callbacks via call_soon, and awaiting an + # already-done future/task returns immediately without giving the loop + # a chance to drain its ready queue. Yield once to let it run. + await asyncio.sleep(0) + assert handle.state == "completed" + + +@pytest.mark.asyncio +async def test_cancel_all_cancels_running_tasks_and_leaves_completed_alone(): + ds = Datasette(memory=True) + await ds.start_background_tasks() + + async def forever(datasette): + await asyncio.Event().wait() + + async def quick(datasette): + return "done" + + forever_handle = ds.add_background_task(forever, name="forever") + quick_handle = ds.add_background_task(quick, name="quick") + await asyncio.wait_for(quick_handle.task, timeout=5) + assert quick_handle.state == "completed" + + await ds._background_tasks.cancel_all(grace=1.0) + + assert forever_handle.state == "cancelled" + assert quick_handle.state == "completed" + + +@pytest.mark.asyncio +async def test_cancel_all_logs_stragglers_that_outlive_the_grace_period(caplog): + ds = Datasette(memory=True) + await ds.start_background_tasks() + + async def stubborn(datasette): + with contextlib.suppress(asyncio.CancelledError): + await asyncio.sleep(10) + # Swallowing CancelledError above and returning normally simulates + # a task that ignores cancellation for longer than the grace period. + await asyncio.sleep(10) + + handle = ds.add_background_task(stubborn, name="stubborn-task") + # Let the task actually start running and reach its first sleep (inside + # the CancelledError-suppressing block) before cancelling it - a task + # cancelled before it has ever run its first step never enters that + # block at all (the throw happens before the coroutine body starts), + # so it would finish cancelling immediately instead of behaving like a + # straggler. + await asyncio.sleep(0) + + with caplog.at_level(logging.WARNING, logger="datasette.background_tasks"): + await ds._background_tasks.cancel_all(grace=0.1) + + assert "stubborn-task" in caplog.text + + # Clean up: actually cancel it now that the test has made its + # assertion, so it doesn't leak past the end of the test. + handle.task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await handle.task diff --git a/tests/test_cli_serve_get.py b/tests/test_cli_serve_get.py index 01b84f59..c4abcfdb 100644 --- a/tests/test_cli_serve_get.py +++ b/tests/test_cli_serve_get.py @@ -52,6 +52,56 @@ def test_serve_with_get(tmp_path_factory): pm.unregister(to_unregister) +def test_serve_with_get_does_not_launch_background_tasks(tmp_path_factory): + # Per decision #3 in plans/first-request/04-core-plan.md, --get must + # never launch background tasks, even though its TestClient request + # flows through the full ASGI stack (including the AsgiRunOnFirstRequest + # fallback that would otherwise launch them). The plugin's startup hook + # itself still runs (registration happens) - only the launch is + # suppressed, so the sentinel file the background task would write must + # never appear. + plugins_dir = tmp_path_factory.mktemp("plugins_for_get_background_tasks") + sentinel = plugins_dir / "sentinel.txt" + (plugins_dir / "bg_task_for_get.py").write_text( + textwrap.dedent( + """ + from datasette import hookimpl + + @hookimpl + def startup(datasette): + async def inner(): + async def task(datasette): + with open("{sentinel}", "w") as fp: + fp.write("ran") + + datasette.add_background_task(task, name="get-sentinel-task") + + return inner + """.format(sentinel=str(sentinel)), + ), + "utf-8", + ) + runner = CliRunner() + result = runner.invoke( + cli, + [ + "serve", + "--memory", + "--plugins-dir", + str(plugins_dir), + "--get", + "/_memory/-/query.json?sql=select+1", + ], + ) + assert result.exit_code == 0, result.output + assert not sentinel.exists() + + to_unregister = [ + p for p in pm.get_plugins() if p.__name__ == "bg_task_for_get.py" + ][0] + pm.unregister(to_unregister) + + def test_serve_with_get_headers(): runner = CliRunner() result = runner.invoke(