feat(runtime-sdk): add lifespan="worker" to asgi.fetch for a worker-lifetime lifespan

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:

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"

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