diff --git a/packages/runtime-sdk/src/workers/asgi.py b/packages/runtime-sdk/src/workers/asgi.py index 76d2a5fb..e4de944b 100644 --- a/packages/runtime-sdk/src/workers/asgi.py +++ b/packages/runtime-sdk/src/workers/asgi.py @@ -1,8 +1,9 @@ +import asyncio import logging from asyncio import Event, Future, Queue, create_task, ensure_future from collections.abc import Awaitable from contextlib import contextmanager -from typing import Any +from typing import Any, Literal from urllib.parse import unquote import js @@ -93,6 +94,70 @@ def request_to_scope(req, env, ws=False, state=None): return scope +async def _no_shutdown() -> None: + return + + +# How long a request waits on a startup another request began, before giving up +# and letting the next request re-drive it. See _ensure_worker_lifespan. +WORKER_LIFESPAN_STARTUP_TIMEOUT = 30.0 + +# Keyed by id(app), holding the app itself alongside the value: an id is only +# unique among live objects, so a collected app could hand its id, and another +# app's lifespan state, to an unrelated object. Keeping it alive costs nothing +# here, since worker mode means the app lives as long as the isolate. +_worker_lifespans: dict[int, tuple[Any, Future]] = {} +_worker_lifespan_states: dict[int, tuple[Any, dict[str, Any]]] = {} + + +def _forget_worker_lifespan(app: Any, fut: Future) -> None: + cached = _worker_lifespans.get(id(app)) + if cached is not None and cached[1] is fut: + del _worker_lifespans[id(app)] + + +async def _ensure_worker_lifespan(app: Any) -> dict[str, Any]: + cached_state = _worker_lifespan_states.get(id(app)) + if cached_state is not None: + # workerd drops promise continuations scheduled onto a dead IoContext + # (compat flag handle_cross_request_promise_resolution, enabled by date + # since 2024-10-14), so once startup is done the state is read directly + # rather than by awaiting a future an earlier request created. + # Concurrent cold starts below still take that risk. + return cached_state[1] + + # Concurrent cold-start requests share one startup instead of each driving + # its own lifespan cycle. + cached = _worker_lifespans.get(id(app)) + if cached is None: + fut = ensure_future(start_application(app)) + _worker_lifespans[id(app)] = (app, fut) + else: + fut = cached[1] + try: + # Shielded so that one waiter timing out or being cancelled does not + # cancel the startup the other waiters are relying on. Bounded because + # a waiter here is on the dropped-continuation path described above, so + # without a deadline it could wait forever, and so would every request + # after it. + _shutdown, state = await asyncio.wait_for( + asyncio.shield(fut), WORKER_LIFESPAN_STARTUP_TIMEOUT + ) + except TimeoutError as exc: + _forget_worker_lifespan(app, fut) + raise RuntimeError( + "ASGI lifespan startup did not complete within " + f"{WORKER_LIFESPAN_STARTUP_TIMEOUT}s; the next request will retry it" + ) from exc + except BaseException: + # Drop the failed startup so the next request retries, rather than + # turning one transient cold-start error into permanent failures. + _forget_worker_lifespan(app, fut) + raise + _worker_lifespan_states[id(app)] = (app, state) + return state + + async def start_application(app): # Drives one ASGI lifespan startup/shutdown cycle before/after serving the # request. The lifespan protocol is optional, so we must tolerate apps that @@ -427,10 +492,20 @@ def _close_transport(finished): async def fetch( - app: Any, req: "Request | js.Request", env: Any, ctx: Context | None = None + app: Any, + req: "Request | js.Request", + env: Any, + ctx: Context | None = None, + *, + lifespan: Literal["request", "worker"] = "request", ) -> js.Response: logger.debug("ASGI request: %s %s", req.method, req.url) - shutdown, state = await start_application(app) + if lifespan == "worker": + shutdown, state = _no_shutdown, await _ensure_worker_lifespan(app) + elif lifespan == "request": + shutdown, state = await start_application(app) + else: + raise ValueError(f"Unsupported lifespan mode: {lifespan!r}") try: result, request_task = await process_request(app, req, env, ctx, state=state) except Exception: diff --git a/packages/runtime-sdk/tests/workerd-test/asgi/tests/test_asgi.py b/packages/runtime-sdk/tests/workerd-test/asgi/tests/test_asgi.py index 28c4479d..bdc7c2da 100644 --- a/packages/runtime-sdk/tests/workerd-test/asgi/tests/test_asgi.py +++ b/packages/runtime-sdk/tests/workerd-test/asgi/tests/test_asgi.py @@ -428,3 +428,38 @@ async def test_multiple_set_cookie_headers_survive(): "first=1; Path=/", "second=2; Path=/", ] + + +async def _startup_count(path): + response = await asyncio.wait_for( + env.SELF.fetch(f"http://example.com{path}"), timeout=5 + ) + return await response.text() + + +@pytest.mark.asyncio +async def test_worker_lifespan_starts_the_app_once(): + assert await _startup_count("/lifespan-worker") == "1" + assert await _startup_count("/lifespan-worker") == "1" + + +@pytest.mark.asyncio +async def test_request_lifespan_starts_the_app_every_time(): + first = int(await _startup_count("/lifespan-request")) + assert int(await _startup_count("/lifespan-request")) == first + 1 + + +@pytest.mark.asyncio +async def test_worker_lifespan_streams_the_whole_body(): + # A streaming response returns before its app task finishes. + response = await asyncio.wait_for( + env.SELF.fetch("http://example.com/lifespan-worker-stream"), timeout=5 + ) + body = await asyncio.wait_for(response.bytes(), timeout=5) + assert len(body) == STREAMING_CHUNK_SIZE * STREAMING_NUM_CHUNKS + + +@pytest.mark.asyncio +async def test_worker_lifespan_accepts_an_unhashable_app(): + assert await _startup_count("/lifespan-worker-unhashable") == "1" + assert await _startup_count("/lifespan-worker-unhashable") == "1" diff --git a/packages/runtime-sdk/tests/workerd-test/asgi/worker.py b/packages/runtime-sdk/tests/workerd-test/asgi/worker.py index 87c7bc21..de143358 100644 --- a/packages/runtime-sdk/tests/workerd-test/asgi/worker.py +++ b/packages/runtime-sdk/tests/workerd-test/asgi/worker.py @@ -205,6 +205,67 @@ async def __call__(self, scope, receive, send): raise RuntimeError("app failed mid-stream") +class UnhashableLifespanApp: + """Unhashable and not weak-referenceable, which worker mode must accept.""" + + __slots__ = ("startup_count",) + __hash__ = None + + def __init__(self): + self.startup_count = 0 + + async def __call__(self, scope, receive, send): + if scope["type"] == "lifespan": + while True: + message = await receive() + if message["type"] == "lifespan.startup": + self.startup_count += 1 + await send({"type": "lifespan.startup.complete"}) + elif message["type"] == "lifespan.shutdown": + await send({"type": "lifespan.shutdown.complete"}) + return + await receive() + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send( + {"type": "http.response.body", "body": str(self.startup_count).encode()} + ) + + +class CountingLifespanApp: + """Reports how many times its lifespan has been started.""" + + def __init__(self): + self.startup_count = 0 + + async def __call__(self, scope, receive, send): + if scope["type"] == "lifespan": + while True: + message = await receive() + if message["type"] == "lifespan.startup": + self.startup_count += 1 + scope["state"]["startup_count"] = self.startup_count + await send({"type": "lifespan.startup.complete"}) + elif message["type"] == "lifespan.shutdown": + await send({"type": "lifespan.shutdown.complete"}) + return + await receive() + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"text/plain")], + } + ) + await send( + { + "type": "http.response.body", + # Read back through the scope, so the test covers the lifespan + # state reaching a later request rather than just the instance. + "body": str(scope["state"]["startup_count"]).encode(), + } + ) + + class ScopeEchoApp: """Echoes selected scope fields as JSON so tests can inspect them.""" @@ -260,6 +321,9 @@ async def __call__(self, scope, receive, send): sse_app = SSEApp() streaming_app = StreamingApp() scope_echo_app = ScopeEchoApp() +worker_lifespan_app = CountingLifespanApp() +request_lifespan_app = CountingLifespanApp() +unhashable_lifespan_app = UnhashableLifespanApp() late_failure_stream_app = LateFailureStreamApp() multi_cookie_app = MultiCookieApp() @@ -279,6 +343,20 @@ async def fetch(self, request): return await asgi.fetch(streaming_app, request, self.env, self.ctx) elif path.startswith("/scope"): return await asgi.fetch(scope_echo_app, request, self.env, self.ctx) + elif path == "/lifespan-worker": + return await asgi.fetch( + worker_lifespan_app, request, self.env, self.ctx, lifespan="worker" + ) + elif path == "/lifespan-worker-unhashable": + return await asgi.fetch( + unhashable_lifespan_app, request, self.env, self.ctx, lifespan="worker" + ) + elif path == "/lifespan-worker-stream": + return await asgi.fetch( + streaming_app, request, self.env, self.ctx, lifespan="worker" + ) + elif path == "/lifespan-request": + return await asgi.fetch(request_lifespan_app, request, self.env, self.ctx) elif path == "/stream-late-failure": return await asgi.fetch( late_failure_stream_app, request, self.env, self.ctx