mirror of
https://github.com/simonw/datasette.git
synced 2026-09-02 22:54:08 +02:00
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:
parent
c748683f2e
commit
96cba70e09
3 changed files with 282 additions and 16 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
211
tests/test_lifespan.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue