Run startup via ASGI lifespan instead of waiting for the first request

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Alex Garcia 2026-07-30 17:37:39 -07:00
commit 96cba70e09
3 changed files with 282 additions and 16 deletions

View file

@ -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,50 @@ 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). It
may also race an explicit `await ds.invoke_startup()` call made by
`datasette serve` before the server starts serving - that's fine,
`invoke_startup()` has its own `_startup_invoked` guard.
"""
if self._startup_invoked:
return
async with self._startup_lock:
if self._startup_invoked:
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

View file

@ -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:
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:
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)

211
tests/test_lifespan.py Normal file
View file

@ -0,0 +1,211 @@
"""
Tests for wiring Datasette startup (setup_db table counts + invoke_startup)
into the ASGI lifespan protocol, per plans/first-request/02-lifespan-startup.
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 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. 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