From 3e018bb1b571cef87c67ae718a5472f23d6b3c6f Mon Sep 17 00:00:00 2001 From: Alex Garcia Date: Tue, 1 Sep 2026 09:39:25 -0700 Subject: [PATCH] Run startup via ASGI lifespan instead of waiting for the first request (#2887) * Run startup via ASGI lifespan instead of waiting for the first request * Ensure immutable table counts still precompute when startup ran first --- datasette/app.py | 47 ++++++-- datasette/cli.py | 7 +- datasette/utils/asgi.py | 42 +++++-- tests/test_lifespan.py | 259 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 337 insertions(+), 18 deletions(-) create mode 100644 tests/test_lifespan.py diff --git a/datasette/app.py b/datasette/app.py index c82ea075..42be7425 100644 --- a/datasette/app.py +++ b/datasette/app.py @@ -453,8 +453,10 @@ class Datasette: self.databases = collections.OrderedDict() self.actions = {} # .invoke_startup() will populate this self._column_types = {} # .invoke_startup() will populate this + self._setup_db_done = False try: self._refresh_schemas_lock = asyncio.Lock() + self._startup_lock = asyncio.Lock() except RuntimeError as rex: # Workaround for intermittent test failure, see: # https://github.com/simonw/datasette/issues/1802 @@ -462,6 +464,7 @@ class Datasette: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) self._refresh_schemas_lock = asyncio.Lock() + self._startup_lock = asyncio.Lock() else: raise self.crossdb = crossdb @@ -2803,24 +2806,52 @@ class Datasette: raise RowNotFound(db.name, table_name, pk_values) return ResolvedRow(db, table_name, sql, params, pks, pk_values, results.first()) + async def _startup_sequence(self): + """Idempotently run the full startup sequence: table counts for + immutable databases, then invoke_startup(). Safe to call more than + once and safe to call concurrently - callers block until whichever + call got there first has finished. + + This is the single entry point used by both AsgiLifespan (so + real deployments finish startup before accepting requests) and + AsgiRunOnFirstRequest (the fallback for hosts that never send + lifespan events, e.g. DatasetteClient's httpx.ASGITransport), and + `datasette serve` (cli.py) calls it too. The fast path below checks + both `_startup_invoked` and `_setup_db_done` - not just the former - + so that a bare `await ds.invoke_startup()` made by a caller ahead of + `_startup_sequence()` (which only sets `_startup_invoked`) can't + make this method skip the immutable-database table-count precompute. + """ + if self._startup_invoked and self._setup_db_done: + return + async with self._startup_lock: + if self._startup_invoked and self._setup_db_done: + return + if not self._setup_db_done: + # First time server starts up, calculate table counts for + # immutable databases + for database in self.databases.values(): + if not database.is_mutable: + await database.table_counts(limit=60 * 60 * 1000) + self._setup_db_done = True + await self.invoke_startup() + def app(self): """Returns an ASGI app function that serves the whole of Datasette""" routes = self._routes() - async def setup_db(): - # First time server starts up, calculate table counts for immutable databases - for database in self.databases.values(): - if not database.is_mutable: - await database.table_counts(limit=60 * 60 * 1000) - 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_shutdown=[_close_on_shutdown]) - asgi = AsgiRunOnFirstRequest(asgi, on_startup=[setup_db, self.invoke_startup]) + asgi = AsgiLifespan( + asgi, + on_startup=[self._startup_sequence], + on_shutdown=[_close_on_shutdown], + ) + asgi = AsgiRunOnFirstRequest(asgi, on_startup=[self._startup_sequence]) for wrapper in pm.hook.asgi_wrapper(datasette=self): asgi = wrapper(asgi) return asgi diff --git a/datasette/cli.py b/datasette/cli.py index 12024a14..2694c1f6 100644 --- a/datasette/cli.py +++ b/datasette/cli.py @@ -710,9 +710,12 @@ def serve( # Populate internal catalog tables before invoke_startup await check_databases(ds) - # Run the "startup" plugin hooks + # Run the full startup sequence (immutable-database table-count + # precompute + the "startup" plugin hooks) via the same entry point + # AsgiLifespan/AsgiRunOnFirstRequest use, so it's not skipped when + # uvicorn's lifespan.startup fires moments later. try: - await ds.invoke_startup() + await ds._startup_sequence() except StartupError as e: raise click.ClickException(e.args[0]) diff --git a/datasette/utils/asgi.py b/datasette/utils/asgi.py index 812194fd..2614ad02 100644 --- a/datasette/utils/asgi.py +++ b/datasette/utils/asgi.py @@ -1,3 +1,4 @@ +import asyncio import json import re from http.cookies import Morsel, SimpleCookie @@ -300,12 +301,24 @@ class AsgiLifespan: while True: message = await receive() if message["type"] == "lifespan.startup": - for fn in self.on_startup: - await fn() + try: + for fn in self.on_startup: + await fn() + except Exception as e: # noqa: BLE001 + await send( + {"type": "lifespan.startup.failed", "message": str(e)} + ) + return await send({"type": "lifespan.startup.complete"}) elif message["type"] == "lifespan.shutdown": - for fn in self.on_shutdown: - await fn() + try: + for fn in self.on_shutdown: + await fn() + except Exception as e: # noqa: BLE001 + await send( + {"type": "lifespan.shutdown.failed", "message": str(e)} + ) + return await send({"type": "lifespan.shutdown.complete"}) return else: @@ -624,10 +637,23 @@ class AsgiRunOnFirstRequest: self.asgi = asgi self.on_startup = on_startup self._started = False + # Guards against concurrent early requests interleaving with startup: + # without this, several requests could all observe `_started is + # False` and proceed before any of them finish running the hooks. + self._lock = asyncio.Lock() async def __call__(self, scope, receive, send): - if not self._started: - self._started = True - for hook in self.on_startup: - await hook() + # Leave "lifespan" scope events alone - this shim only exists as a + # fallback for hosts that never send them. It wraps AsgiLifespan, so + # if it ran on_startup here too, a startup exception would escape + # before AsgiLifespan's own try/except got a chance to turn it into + # a lifespan.startup.failed message. + if scope["type"] != "lifespan" and not self._started: + async with self._lock: + # Re-check: another request may have finished startup while + # we were waiting for the lock. + if not self._started: + for hook in self.on_startup: + await hook() + self._started = True return await self.asgi(scope, receive, send) diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py new file mode 100644 index 00000000..3655285e --- /dev/null +++ b/tests/test_lifespan.py @@ -0,0 +1,259 @@ +""" +Tests for wiring Datasette startup (setup_db table counts + invoke_startup) +into the ASGI lifespan protocol. + +These exercise Datasette._startup_sequence() via three different callers: +- AsgiLifespan, by hand-driving lifespan.startup messages (no HTTP request) +- AsgiRunOnFirstRequest, the fallback for hosts that never send lifespan + events (this is what DatasetteClient / plain httpx.ASGITransport uses) +- Both at once, to prove startup hooks run at most once +""" + +import asyncio +import contextlib +import sqlite3 + +import httpx +import pytest + +from datasette import hookimpl +from datasette.app import Datasette +from datasette.database import Database +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. Mirrors what a real server does: after startup + completes it parks waiting for the next event. We cancel that wait + once we've observed the startup response, rather than closing the + Datasette instance down with a shutdown message. + """ + messages_sent = [] + startup_responded = asyncio.Event() + delivered = False + + async def receive(): + nonlocal delivered + if not delivered: + delivered = True + return {"type": "lifespan.startup"} + # No further messages: block until the task is cancelled below, + # same as a real server parked waiting for lifespan.shutdown. + 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_lifespan_startup_runs_before_any_request(): + ds = Datasette(memory=True) + assert ds._startup_invoked is False + app = ds.app() + + messages = await _drive_lifespan_startup(app) + + assert {"type": "lifespan.startup.complete"} in messages + assert ds._startup_invoked is True + # Internal catalog tables should be populated too, entirely without an + # HTTP request having been made. + internal_db = ds.get_internal_database() + databases = await internal_db.execute("select * from catalog_databases") + assert len(databases.rows) >= 1 + + +@pytest.mark.asyncio +async def test_lifespan_startup_failure_reports_lifespan_startup_failed(): + class RaisingStartupPlugin: + __name__ = "RaisingStartupPlugin" + + @hookimpl + def startup(self, datasette): + async def inner(): + raise RuntimeError("boom from startup hook") + + return inner + + ds = Datasette(memory=True) + pm.register(RaisingStartupPlugin(), name="raising_startup_plugin") + try: + app = ds.app() + messages = await _drive_lifespan_startup(app) + finally: + pm.unregister(name="raising_startup_plugin") + + assert messages == [ + {"type": "lifespan.startup.failed", "message": "boom from startup hook"} + ] + # The exception happened before invoke_startup() got to the end of its + # body, so startup is not considered to have completed. + assert ds._startup_invoked is False + + +@pytest.mark.asyncio +async def test_startup_runs_exactly_once_across_lifespan_and_first_request(): + call_count = {"n": 0} + + class CountingStartupPlugin: + __name__ = "CountingStartupPlugin" + + @hookimpl + def startup(self, datasette): + async def inner(): + call_count["n"] += 1 + + return inner + + ds = Datasette(memory=True) + pm.register(CountingStartupPlugin(), name="counting_startup_plugin") + try: + # Build the ASGI app once, the way a real deployment does - and + # reuse the SAME app instance for both the lifespan drive and the + # HTTP requests below, since a fresh ds.app() call would reset the + # AsgiRunOnFirstRequest fallback's state. + app = ds.app() + + messages = await _drive_lifespan_startup(app) + assert {"type": "lifespan.startup.complete"} in messages + assert call_count["n"] == 1 + + # A first HTTP request (as if the host never sent lifespan events, + # or lifespan already ran) should not run the hook again. + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://localhost" + ) as client: + response1 = await client.get("/-/versions.json") + assert response1.status_code == 200 + # ... nor should a second, repeat request. + response2 = await client.get("/-/versions.json") + assert response2.status_code == 200 + finally: + pm.unregister(name="counting_startup_plugin") + + assert call_count["n"] == 1 + + +@pytest.mark.asyncio +async def test_no_lifespan_first_request_still_triggers_startup(): + # Pin today's behavior: a client that never drives ASGI lifespan events + # at all (like httpx.ASGITransport, which DatasetteClient uses) still + # gets startup armed by the AsgiRunOnFirstRequest fallback. + ds = Datasette(memory=True) + assert ds._startup_invoked is False + app = ds.app() + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://localhost" + ) as client: + response = await client.get("/-/versions.json") + assert response.status_code == 200 + + assert ds._startup_invoked is True + internal_db = ds.get_internal_database() + databases = await internal_db.execute("select * from catalog_databases") + assert len(databases.rows) >= 1 + + +@pytest.mark.asyncio +async def test_datasette_client_first_request_triggers_startup(): + # Same as above, but through the real DatasetteClient (ds.client) that + # plugins and tests actually use, to confirm nothing regressed there. + ds = Datasette(memory=True) + assert ds._startup_invoked is False + response = await ds.client.get("/-/versions.json") + assert response.status_code == 200 + assert ds._startup_invoked is True + + +@pytest.mark.asyncio +async def test_concurrent_first_requests_all_wait_for_slow_startup(): + call_count = {"n": 0} + + class SlowStartupPlugin: + __name__ = "SlowStartupPlugin" + + @hookimpl + def startup(self, datasette): + async def inner(): + call_count["n"] += 1 + await asyncio.sleep(0.2) + + return inner + + ds = Datasette(memory=True) + pm.register(SlowStartupPlugin(), name="slow_startup_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)] + ) + finally: + pm.unregister(name="slow_startup_plugin") + + # Every one of the 10 simultaneous first requests must have blocked + # until startup actually finished, not raced ahead of it. + assert all(response.status_code == 200 for response in responses) + assert call_count["n"] == 1 + assert ds._startup_invoked is True + + +@pytest.mark.asyncio +async def test_setup_db_still_runs_when_invoke_startup_ran_first(tmp_path, monkeypatch): + # Regression test: `datasette serve` (cli.py _serve_async) calls + # ds.invoke_startup() directly, before uvicorn ever sends a + # lifespan.startup event that drives _startup_sequence(). If + # _startup_sequence()'s fast path only checked `_startup_invoked`, it + # would see startup already done and skip the immutable-database + # table-count precompute (setup_db) entirely - a silent regression + # versus main, where AsgiRunOnFirstRequest ran setup_db unconditionally + # on request #1. + db_path = tmp_path / "immutable.db" + conn = sqlite3.connect(str(db_path)) + conn.execute("create table t (id integer primary key)") + conn.commit() + conn.close() + + ds = Datasette([], immutables=[str(db_path)]) + + call_count = {"n": 0} + original_table_counts = Database.table_counts + + async def counting_table_counts(self, *args, **kwargs): + call_count["n"] += 1 + return await original_table_counts(self, *args, **kwargs) + + monkeypatch.setattr(Database, "table_counts", counting_table_counts) + + # Simulate the CLI path: invoke_startup() runs directly and completes + # BEFORE _startup_sequence() ever gets a chance to run setup_db. + await ds.invoke_startup() + assert ds._startup_invoked is True + assert call_count["n"] == 0 + + # The lifespan/first-request path (or the CLI itself, per the fix) + # calling the shared entry point afterwards must still precompute + # table counts for immutable databases. + await ds._startup_sequence() + assert call_count["n"] == 1 + assert ds._setup_db_done is True + + # Idempotency: a second call must not recompute. + await ds._startup_sequence() + assert call_count["n"] == 1