Merge f614a781bb into cae36803bd

datasette/app.py

@@ -423,6 +423,7 @@ class Datasette:
         default_deny=False,
     ):
         self._startup_invoked = False
+        self._shutdown_invoked = False
         self._closed = False
         assert config_dir is None or isinstance(
             config_dir, Path
@@ -2896,20 +2897,48 @@ class Datasette:
             return
         await self._background_tasks.launch_all()
 
+    async def invoke_shutdown(self):
+        """Run the graceful teardown sequence: plugin ``shutdown`` hooks,
+        then cancel and drain supervised background tasks, then close
+        every database.
+
+        Idempotent (guarded by ``_shutdown_invoked``) and safe to call more
+        than once - a second ``lifespan.shutdown`` message from a
+        misbehaving ASGI host, or any future caller, must not re-run
+        teardown. A ``shutdown`` hook that raises is logged and swallowed
+        rather than propagated, so one broken plugin can't skip another
+        plugin's cleanup, or skip task cancellation / ``close()``
+        altogether.
+
+        Order matters: hooks run first, while background tasks are still
+        alive, so a plugin can coordinate with its own task (e.g. tell a
+        queue consumer to stop pulling new work) before that task gets
+        cancelled; ``close()`` runs last so both the hooks and the
+        cancelled tasks still have working database connections to write
+        any final state.
+        """
+        if self._shutdown_invoked:
+            return
+        self._shutdown_invoked = True
+        for hook in pm.hook.shutdown(datasette=self):
+            try:
+                await await_me_maybe(hook)
+            except Exception:
+                logging.getLogger("datasette").exception("shutdown hook failed")
+        await self._background_tasks.cancel_all(grace=5.0)
+        self.close()
+
     def app(self):
         """Returns an ASGI app function that serves the whole of Datasette"""
         routes = self._routes()
 
-        async def _close_on_shutdown():
-            self.close()
-
         asgi = CrossOriginProtectionMiddleware(DatasetteRouter(self, routes), self)
         if self.setting("trace_debug"):
             asgi = AsgiTracer(asgi)
         asgi = AsgiLifespan(
             asgi,
             on_startup=[self._startup_sequence, self._launch_background_tasks],
-            on_shutdown=[_close_on_shutdown],
+            on_shutdown=[self.invoke_shutdown],
         )
         asgi = AsgiRunOnFirstRequest(
             asgi,

datasette/hookspecs.py

@@ -9,6 +9,11 @@ def startup(datasette):
     """Fires directly after Datasette first starts running"""
 
 
+@hookspec
+def shutdown(datasette):
+    """Called once when the Datasette server is shutting down"""
+
+
 @hookspec
 def asgi_wrapper(datasette):
     """Returns an ASGI middleware callable to wrap our ASGI application with"""

docs/internals.rst

@@ -1416,7 +1416,7 @@ Datasette guarantees a fixed sequence of events between the moment a ``Datasette
 2. **Startup** — ``await datasette.invoke_startup()`` runs once: it populates the internal database's catalog of table schemas (:ref:`internals_internal`), loads canned queries and column type configuration, then calls every registered :ref:`plugin_hook_startup` hook, in plugin registration order. When Datasette is being served, table-count precomputation for immutable databases runs immediately before this, as part of the same startup sequence.
 3. **Background-task launch** — once *every* ``startup`` hook has finished (not before), every task registered with :ref:`datasette_add_background_task` — by any plugin — is launched. A task registered by one plugin's ``startup`` hook can safely depend on state set up by another plugin's ``startup`` hook, because launch only happens after the whole round of hooks completes.
 4. **Serving** — the instance handles requests (or, for headless or CLI use, does whatever the embedding program does with it).
-5. **Shutdown** — triggered by the ASGI ``lifespan.shutdown`` event (Ctrl-C, ``SIGTERM``) or the end of a ``datasette serve`` process: every still-running background task is cancelled and given a five-second grace period to actually stop; finally every database connection is released via :ref:`datasette_close`.
+5. **Shutdown** — triggered by the ASGI ``lifespan.shutdown`` event (Ctrl-C, ``SIGTERM``) or the end of a ``datasette serve`` process: every :ref:`plugin_hook_shutdown` hook runs first, while background tasks are still alive, so a plugin can tell its own task to wind down gracefully; every still-running background task is then cancelled and given a five-second grace period to actually stop; finally every database connection is released via :ref:`datasette_close`.
 
 .. admonition:: Startup hooks run on the event loop that serves requests
 

docs/plugin_hooks.rst

@@ -1212,6 +1212,31 @@ Potential use-cases:
 
 Examples: `datasette-saved-queries <https://datasette.io/plugins/datasette-saved-queries>`__, `datasette-init <https://datasette.io/plugins/datasette-init>`__
 
+.. _plugin_hook_shutdown:
+
+shutdown(datasette)
+-------------------
+
+This hook fires once, when the Datasette application server is shutting down gracefully - triggered by the ASGI ``lifespan.shutdown`` event, which includes pressing Ctrl-C or sending ``SIGTERM`` to a ``datasette serve`` process. It is not called on a hard kill (``SIGKILL``), since there is no opportunity to run any code in that case.
+
+Like ``startup()``, this can be a regular function or it can return an async function to be awaited.
+
+It runs before Datasette cancels any background tasks it is supervising (see :ref:`datasette_add_background_task`) and before it closes its database connections, so you can use it to tell your plugin's own background work to stop gracefully while a database connection is still available to write out any final state. See :ref:`datasette_lifecycle` for exactly where this fits into the full startup-to-shutdown sequence:
+
+.. code-block:: python
+
+    @hookimpl
+    def shutdown(datasette):
+        async def inner():
+            db = datasette.get_database()
+            await db.execute_write(
+                "insert into shutdown_log (at) values (datetime('now'))"
+            )
+
+        return inner
+
+If your ``shutdown()`` hook raises an exception it will be logged but not re-raised, so one plugin's broken shutdown code cannot prevent other plugins - or Datasette itself - from finishing their own teardown.
+
 .. _plugin_hook_actor_from_request:
 
 actor_from_request(datasette, request)

tests/test_cli_serve_server.py

@@ -1,4 +1,6 @@
+import signal
 import socket
+import subprocess
 import time
 
 import httpx
@@ -145,3 +147,83 @@ def test_startup_error_fails_fast_before_port_binds(serve_with_plugins):
         socket.create_connection(("127.0.0.1", port), timeout=0.2),
     ):
         pass
+
+
+# Proves the ticket-01 serve path (asyncio.run(_serve_async()) wrapping
+# uvicorn.Server.serve()) actually delivers a graceful signal through to
+# uvicorn's lifespan.shutdown, which now runs Datasette.invoke_shutdown()
+# and therefore every plugin's `shutdown` hook. The plugin below writes a
+# sentinel file from inside its shutdown hook so this can be checked from
+# outside the subprocess after it exits.
+SHUTDOWN_SENTINEL_PLUGIN_TEMPLATE = """
+import pathlib
+from datasette import hookimpl
+
+SENTINEL_PATH = {sentinel_path!r}
+
+
+@hookimpl
+def shutdown(datasette):
+    pathlib.Path(SENTINEL_PATH).write_text("shutdown ran", "utf-8")
+"""
+
+
+def _start_serve_with_shutdown_sentinel(serve_with_plugins, tmp_path):
+    sentinel_path = tmp_path / "shutdown-sentinel.txt"
+    proc, _ = serve_with_plugins(
+        {
+            "shutdown_sentinel_plugin": SHUTDOWN_SENTINEL_PLUGIN_TEMPLATE.format(
+                sentinel_path=str(sentinel_path)
+            )
+        }
+    )
+    return proc, sentinel_path
+
+
+@pytest.mark.serial
+def test_sigterm_runs_shutdown_hooks(serve_with_plugins, tmp_path):
+    ds_proc, sentinel_path = _start_serve_with_shutdown_sentinel(
+        serve_with_plugins, tmp_path
+    )
+    assert not sentinel_path.exists()
+    ds_proc.send_signal(signal.SIGTERM)
+    try:
+        ds_proc.wait(timeout=10)
+    except subprocess.TimeoutExpired:
+        ds_proc.kill()
+        ds_proc.wait()
+        raise AssertionError(
+            "datasette serve did not exit within 10s of SIGTERM\n"
+            + ds_proc.stdout.read().decode("utf-8")
+        )
+    output = ds_proc.stdout.read().decode("utf-8")
+    assert sentinel_path.exists(), (
+        "shutdown hook never wrote its sentinel file after SIGTERM\n" + output
+    )
+    assert sentinel_path.read_text("utf-8") == "shutdown ran"
+
+
+@pytest.mark.serial
+@pytest.mark.skipif(
+    not hasattr(signal, "SIGINT"), reason="Requires signal.SIGINT support"
+)
+def test_sigint_runs_shutdown_hooks(serve_with_plugins, tmp_path):
+    ds_proc, sentinel_path = _start_serve_with_shutdown_sentinel(
+        serve_with_plugins, tmp_path
+    )
+    assert not sentinel_path.exists()
+    ds_proc.send_signal(signal.SIGINT)
+    try:
+        ds_proc.wait(timeout=10)
+    except subprocess.TimeoutExpired:
+        ds_proc.kill()
+        ds_proc.wait()
+        raise AssertionError(
+            "datasette serve did not exit within 10s of SIGINT\n"
+            + ds_proc.stdout.read().decode("utf-8")
+        )
+    output = ds_proc.stdout.read().decode("utf-8")
+    assert sentinel_path.exists(), (
+        "shutdown hook never wrote its sentinel file after SIGINT\n" + output
+    )
+    assert sentinel_path.read_text("utf-8") == "shutdown ran"

tests/test_plugins.py

@@ -53,6 +53,13 @@ def test_hook_jump_items_sql():
     assert "jump_items_sql" in dir(pm.hook)
 
 
+def test_hook_shutdown():
+    # Detailed behavior (ordering against background-task cancellation and
+    # close(), idempotency, exception handling, sync vs async support) is
+    # covered in tests/test_shutdown.py.
+    assert "shutdown" in dir(pm.hook)
+
+
 @pytest.mark.asyncio
 async def test_hook_plugins_dir_plugin_prepare_connection(ds_client):
     response = await ds_client.get(

tests/test_shutdown.py

@@ -0,0 +1,318 @@
+"""
+Tests for the shutdown(datasette) plugin hook and Datasette.invoke_shutdown().
+
+Order under test: plugin `shutdown` hooks run first (while background tasks
+are still alive) -> supervised background tasks are cancelled and drained
+(fixed 5s grace) -> databases are closed. Hook exceptions are logged, never
+propagated, and never skip a later step. The whole sequence is idempotent,
+so a second lifespan.shutdown message (or any other second caller) is a
+no-op.
+"""
+
+import asyncio
+import contextlib
+import logging
+
+import pytest
+
+from datasette import hookimpl
+from datasette.app import Datasette
+from datasette.plugins import pm
+
+
+async def _drive_lifespan(app, messages):
+    """Drive a single ASGI lifespan connection against `app`, feeding
+    `messages` to it in order via receive(). Returns once `app` itself
+    returns, which - per AsgiLifespan - happens after it has processed a
+    lifespan.shutdown message and sent back .complete or .failed. Returns
+    the list of messages `app` sent via send().
+
+    `messages` must end with a lifespan.shutdown (or a startup that fails)
+    or this will hang forever waiting for a message that never comes,
+    since AsgiLifespan only returns after handling shutdown.
+    """
+    sent = []
+    idx = 0
+
+    async def receive():
+        nonlocal idx
+        if idx < len(messages):
+            message = messages[idx]
+            idx += 1
+            return message
+        # Real servers park here waiting for lifespan.shutdown; nothing
+        # left to deliver in this test, so just block - the caller is
+        # expected to have already gotten what it needs via a preceding
+        # lifespan.shutdown in `messages`.
+        await asyncio.Event().wait()
+
+    async def send(message):
+        sent.append(message)
+
+    await app({"type": "lifespan"}, receive, send)
+    return sent
+
+
+@pytest.mark.asyncio
+async def test_shutdown_hook_runs_before_task_cancellation_then_closes():
+    # Ordering proof: the shutdown hook observes the background task still
+    # "running" (hooks run BEFORE cancellation); after the whole lifespan
+    # drive completes, the task is "cancelled" and the Datasette instance
+    # is closed.
+    events = []
+    task_state_seen_in_hook = {}
+
+    async def bg_task(datasette):
+        await asyncio.Event().wait()
+
+    class LifecyclePlugin:
+        __name__ = "LifecyclePlugin"
+
+        @hookimpl
+        def startup(self, datasette):
+            async def inner():
+                datasette.add_background_task(bg_task, name="bg-task")
+
+            return inner
+
+        @hookimpl
+        def shutdown(self, datasette):
+            async def inner():
+                events.append("shutdown-hook-ran")
+                handle = datasette._background_tasks.tasks()[0]
+                task_state_seen_in_hook["state"] = handle.state
+
+            return inner
+
+    ds = Datasette(memory=True)
+    pm.register(LifecyclePlugin(), name="lifecycle_plugin")
+    try:
+        app = ds.app()
+        messages = await _drive_lifespan(
+            app, [{"type": "lifespan.startup"}, {"type": "lifespan.shutdown"}]
+        )
+    finally:
+        pm.unregister(name="lifecycle_plugin")
+
+    assert {"type": "lifespan.startup.complete"} in messages
+    assert {"type": "lifespan.shutdown.complete"} in messages
+    assert events == ["shutdown-hook-ran"]
+    assert task_state_seen_in_hook["state"] == "running"
+
+    handle = ds._background_tasks.tasks()[0]
+    assert handle.state == "cancelled"
+    assert ds._closed is True
+    assert ds._shutdown_invoked is True
+
+
+@pytest.mark.asyncio
+async def test_second_lifespan_shutdown_does_not_double_invoke():
+    call_count = {"n": 0}
+
+    class CountingShutdownPlugin:
+        __name__ = "CountingShutdownPlugin"
+
+        @hookimpl
+        def shutdown(self, datasette):
+            async def inner():
+                call_count["n"] += 1
+
+            return inner
+
+    ds = Datasette(memory=True)
+    pm.register(CountingShutdownPlugin(), name="counting_shutdown_plugin")
+    try:
+        app = ds.app()
+        messages = await _drive_lifespan(
+            app, [{"type": "lifespan.startup"}, {"type": "lifespan.shutdown"}]
+        )
+        assert {"type": "lifespan.shutdown.complete"} in messages
+        assert call_count["n"] == 1
+
+        # A second, separate lifespan connection (a misbehaving host, or a
+        # second embedder driving the same Datasette instance) sends
+        # lifespan.shutdown again. AsgiLifespan itself has no memory of
+        # the earlier connection, so this exercises invoke_shutdown()'s
+        # own idempotency guard, not anything ASGI-layer.
+        messages2 = await _drive_lifespan(app, [{"type": "lifespan.shutdown"}])
+        assert {"type": "lifespan.shutdown.complete"} in messages2
+        assert call_count["n"] == 1
+    finally:
+        pm.unregister(name="counting_shutdown_plugin")
+
+    assert ds._closed is True
+
+
+@pytest.mark.asyncio
+async def test_raising_shutdown_hook_is_logged_and_does_not_block_the_rest(caplog):
+    events = []
+
+    async def bg_task(datasette):
+        await asyncio.Event().wait()
+
+    class RaisingShutdownPlugin:
+        __name__ = "RaisingShutdownPlugin"
+
+        @hookimpl
+        def shutdown(self, datasette):
+            async def inner():
+                raise RuntimeError("boom from shutdown hook")
+
+            return inner
+
+    class WellBehavedPlugin:
+        __name__ = "WellBehavedPlugin"
+
+        @hookimpl
+        def startup(self, datasette):
+            async def inner():
+                datasette.add_background_task(bg_task, name="bg-task")
+
+            return inner
+
+        @hookimpl
+        def shutdown(self, datasette):
+            async def inner():
+                events.append("well-behaved-ran")
+
+            return inner
+
+    ds = Datasette(memory=True)
+    pm.register(RaisingShutdownPlugin(), name="raising_shutdown_plugin")
+    pm.register(WellBehavedPlugin(), name="well_behaved_plugin")
+    try:
+        app = ds.app()
+        with caplog.at_level(logging.ERROR, logger="datasette"):
+            messages = await _drive_lifespan(
+                app, [{"type": "lifespan.startup"}, {"type": "lifespan.shutdown"}]
+            )
+    finally:
+        pm.unregister(name="raising_shutdown_plugin")
+        pm.unregister(name="well_behaved_plugin")
+
+    # The exception must never turn into lifespan.shutdown.failed - it's
+    # swallowed inside invoke_shutdown(), logged, and teardown continues.
+    assert {"type": "lifespan.shutdown.complete"} in messages
+    assert events == ["well-behaved-ran"]
+    assert "shutdown hook failed" in caplog.text
+    assert "boom from shutdown hook" in caplog.text
+
+    handle = ds._background_tasks.tasks()[0]
+    assert handle.state == "cancelled"
+    assert ds._closed is True
+
+
+@pytest.mark.asyncio
+async def test_sync_shutdown_hook_variant_works():
+    events = []
+
+    class SyncShutdownPlugin:
+        __name__ = "SyncShutdownPlugin"
+
+        @hookimpl
+        def shutdown(self, datasette):
+            # Deliberately not returning a coroutine/callable - a plain
+            # sync hookimpl, same as `def startup(datasette): ...` is
+            # supported via await_me_maybe.
+            events.append("sync-shutdown-ran")
+
+    ds = Datasette(memory=True)
+    pm.register(SyncShutdownPlugin(), name="sync_shutdown_plugin")
+    try:
+        app = ds.app()
+        messages = await _drive_lifespan(
+            app, [{"type": "lifespan.startup"}, {"type": "lifespan.shutdown"}]
+        )
+    finally:
+        pm.unregister(name="sync_shutdown_plugin")
+
+    assert {"type": "lifespan.shutdown.complete"} in messages
+    assert events == ["sync-shutdown-ran"]
+    assert ds._closed is True
+
+
+@pytest.mark.asyncio
+async def test_shutdown_logs_stragglers_that_outlive_the_grace_period(
+    caplog, monkeypatch
+):
+    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 - same shape as test_background_tasks.py's equivalent
+        # test.
+        await asyncio.sleep(10)
+
+    class StubbornTaskPlugin:
+        __name__ = "StubbornTaskPlugin"
+
+        @hookimpl
+        def startup(self, datasette):
+            async def inner():
+                datasette.add_background_task(stubborn, name="stubborn-task")
+
+            return inner
+
+    ds = Datasette(memory=True)
+    pm.register(StubbornTaskPlugin(), name="stubborn_task_plugin")
+
+    # invoke_shutdown() calls self._background_tasks.cancel_all(grace=5.0)
+    # with a grace hardcoded in app.py, per the ticket. Rather than
+    # sleeping for a real 5s in this test, monkeypatch the
+    # BackgroundTaskSupervisor *instance's* cancel_all to a wrapper that
+    # ignores the caller-supplied grace and substitutes a small one - this
+    # is the cleanest seam because it requires no production-code changes
+    # (no grace= setting/attribute to add) and leaves invoke_shutdown's
+    # own code under test untouched.
+    real_cancel_all = ds._background_tasks.cancel_all
+
+    async def fast_cancel_all(grace=5.0):
+        return await real_cancel_all(grace=0.1)
+
+    monkeypatch.setattr(ds._background_tasks, "cancel_all", fast_cancel_all)
+
+    try:
+        app = ds.app()
+
+        sent = []
+        queue = asyncio.Queue()
+        startup_complete = asyncio.Event()
+
+        async def receive():
+            return await queue.get()
+
+        async def send(message):
+            sent.append(message)
+            if message.get("type") == "lifespan.startup.complete":
+                startup_complete.set()
+
+        task = asyncio.create_task(app({"type": "lifespan"}, receive, send))
+        await queue.put({"type": "lifespan.startup"})
+        await asyncio.wait_for(startup_complete.wait(), timeout=5)
+
+        # Let the stubborn background task actually start running and
+        # reach its CancelledError-suppressing sleep before shutdown
+        # cancels it - a task cancelled before it has ever run its first
+        # step never enters that block at all, so it would finish
+        # cancelling immediately instead of behaving like a straggler.
+        await asyncio.sleep(0.05)
+
+        with caplog.at_level(logging.WARNING, logger="datasette.background_tasks"):
+            await queue.put({"type": "lifespan.shutdown"})
+            await asyncio.wait_for(task, timeout=5)
+
+        assert {"type": "lifespan.shutdown.complete"} in sent
+        assert "stubborn-task" in caplog.text
+        assert ds._closed is True
+    finally:
+        pm.unregister(name="stubborn_task_plugin")
+        # Clean up: the stubborn task ignored cancellation and is still
+        # sleeping past the shrunk grace period; actually cancel and await
+        # it now that the test has made its assertions, so it doesn't leak
+        # past the end of the test.
+        handles = ds._background_tasks.tasks()
+        if handles and handles[0].task and not handles[0].task.done():
+            handles[0].task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await handles[0].task