diff --git a/datasette/app.py b/datasette/app.py index e5d55b42..9f6c81d3 100644 --- a/datasette/app.py +++ b/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, diff --git a/datasette/hookspecs.py b/datasette/hookspecs.py index f89f2f36..0b807f8c 100644 --- a/datasette/hookspecs.py +++ b/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""" diff --git a/docs/internals.rst b/docs/internals.rst index 03a5ec15..b3c217bf 100644 --- a/docs/internals.rst +++ b/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 diff --git a/docs/plugin_hooks.rst b/docs/plugin_hooks.rst index fa24c7c7..a1537ef1 100644 --- a/docs/plugin_hooks.rst +++ b/docs/plugin_hooks.rst @@ -1212,6 +1212,31 @@ Potential use-cases: Examples: `datasette-saved-queries `__, `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) diff --git a/tests/test_cli_serve_server.py b/tests/test_cli_serve_server.py index b76180fd..f6f63b0a 100644 --- a/tests/test_cli_serve_server.py +++ b/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" diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 734f0fc2..a5d09b27 100644 --- a/tests/test_plugins.py +++ b/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( diff --git a/tests/test_shutdown.py b/tests/test_shutdown.py new file mode 100644 index 00000000..b6fb02c9 --- /dev/null +++ b/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