feat(runtime-sdk): Add serialize option in WSGI entrypoint

packages/runtime-sdk/src/workers/_serialize.py

@@ -0,0 +1,60 @@
+import asyncio
+from weakref import ReferenceType, WeakKeyDictionary, ref
+
+# Use weakref so that event loops can be garbage collected
+# when they are no longer in use.
+_locks: WeakKeyDictionary[asyncio.AbstractEventLoop, ReferenceType[asyncio.Lock]] = (
+    WeakKeyDictionary()
+)
+
+
+def _get_lock() -> asyncio.Lock:
+    """Return the serialization lock for the currently running event loop."""
+    loop = asyncio.get_running_loop()
+    lock_ref = _locks.get(loop)
+    lock = lock_ref() if lock_ref is not None else None
+    if lock is None:
+        lock = asyncio.Lock()
+        _locks[loop] = ref(lock)
+    return lock
+
+
+class RequestLock:
+    """Manage one request's ownership of the shared serialization lock."""
+
+    def __init__(self, enabled: bool) -> None:
+        self._enabled = enabled
+        self._lock: asyncio.Lock | None = None
+        self._release_on_exit = True
+
+    async def __aenter__(self) -> "RequestLock":
+        if self._enabled:
+            self._lock = _get_lock()
+            await self._lock.acquire()
+        return self
+
+    async def __aexit__(self, _exc_type, _exc, _traceback) -> None:
+        if not self._enabled:
+            return
+        if self._release_on_exit:
+            self.release()
+
+    def defer_release(self) -> None:
+        """
+        When there are background streams that need to be processed,
+        we need to defer the release of the lock until the streams are done.
+        """
+        if not self._enabled:
+            return
+
+        self._release_on_exit = False
+
+    def release(self) -> None:
+        """Release this request's lock once."""
+        if not self._enabled:
+            return
+
+        lock = self._lock
+        if lock is not None:
+            self._lock = None
+            lock.release()

packages/runtime-sdk/src/workers/wsgi.py

@@ -8,6 +8,7 @@
 import js
 
 from workers import Context, Request, WorkerEntrypoint
+from workers._serialize import RequestLock
 
 logger = logging.getLogger("wsgi")
 NULL_BODY_STATUSES = frozenset({101, 103, 204, 205, 304})
@@ -291,6 +292,7 @@ def process_request(
     req: "Request | js.Request",
     env: Any,
     body: "bytes | io.IOBase",
+    on_close: "Callable[[], None] | None" = None,
 ) -> js.Response:
     environ = build_environ(req, env, body)
 
@@ -325,11 +327,16 @@ def start_response(status, response_headers, exc_info=None):
     result_iter = iter(result)
 
     def close_all() -> None:
-        _close_iterable(result)
         try:
-            environ["wsgi.input"].close()
-        except Exception:  # noqa: BLE001 - best-effort cleanup
-            logger.exception("Failed to close wsgi.input")
+            _close_iterable(result)
+        finally:
+            try:
+                environ["wsgi.input"].close()
+            except Exception:  # noqa: BLE001 - best-effort cleanup
+                logger.exception("Failed to close wsgi.input")
+            finally:
+                if on_close is not None:
+                    on_close()
 
     # WSGI apps must call start_response before yielding the first body chunk,
     # but some defer it until the first non-empty chunk is produced. Pull that
@@ -367,25 +374,36 @@ async def fetch(
     env: Any,
     # Accepted for parity with asgi.fetch; WSGI has no use for it.
     ctx: Context | None = None,
+    *,
+    serialize: bool = False,
 ) -> js.Response:
     logger.debug("WSGI request: %s %s", req.method, req.url)
-    # Prefer lazily streaming the body through `wsgi.input` (no full buffering);
-    # fall back to pre-buffering when `run_sync`/JSPI isn't available.
-    body: bytes | io.IOBase | None = _make_wsgi_input(req)
-    if body is None:
-        body = await _read_body(req)
-    try:
-        return process_request(app, req, env, body)
-    except Exception:
-        logger.exception("WSGI request failed")
-        raise
-
-
-def entrypoint(app: Any) -> type[WorkerEntrypoint]:
+    async with RequestLock(serialize) as request_lock:
+        try:
+            # Prefer lazily streaming the body through `wsgi.input` (no full buffering);
+            # fall back to pre-buffering when `run_sync`/JSPI isn't available.
+            body: bytes | io.IOBase | None = _make_wsgi_input(req)
+            if body is None:
+                body = await _read_body(req)
+            response = process_request(
+                app,
+                req,
+                env,
+                body,
+                on_close=request_lock.release,
+            )
+            request_lock.defer_release()
+            return response
+        except Exception:
+            logger.exception("WSGI request failed")
+            raise
+
+
+def entrypoint(app: Any, *, serialize: bool = False) -> type[WorkerEntrypoint]:
     """Create the default Worker entrypoint for a WSGI application."""
 
     class Default(WorkerEntrypoint):
         async def fetch(self, request):
-            return await fetch(app, request, self.env)
+            return await fetch(app, request, self.env, serialize=serialize)
 
     return Default

packages/runtime-sdk/tests/workerd-test/wsgi/tests/test_wsgi.py

@@ -1,8 +1,9 @@
+import asyncio
 import json
 
 import js
 import pytest
-from pyodide.ffi import to_js
+from pyodide.ffi import run_sync, to_js
 from worker import (
     STREAMING_CHUNK_SIZE,
     STREAMING_NUM_CHUNKS,
@@ -119,3 +120,124 @@ def test_build_environ_handles_js_and_python_requests():
     py_env = wsgi.build_environ(py_request, env, b"")
     assert js_env["HTTP_HEADER1"] == py_env["HTTP_HEADER1"] == "Value1"
     assert js_env["HTTP_HEADER2"] == py_env["HTTP_HEADER2"] == "Value2"
+
+
+class _ConcurrentApp:
+    def __init__(self):
+        self.active = 0
+        self.max_active = 0
+
+    def __call__(self, environ, start_response):
+        self.active += 1
+        self.max_active = max(self.max_active, self.active)
+        run_sync(asyncio.sleep(0.05))
+        self.active -= 1
+        start_response("200 OK", [])
+        return [b"ok"]
+
+
+async def _fetch_wsgi_text(app, *, serialize=False):
+    response = await wsgi.fetch(
+        app, js.Request.new("http://example.com/"), env, serialize=serialize
+    )
+    return await response.text()
+
+
+@pytest.mark.asyncio
+async def test_serialized_wsgi_requests_do_not_overlap():
+    app = _ConcurrentApp()
+    await asyncio.wait_for(
+        asyncio.gather(
+            _fetch_wsgi_text(app, serialize=True),
+            _fetch_wsgi_text(app, serialize=True),
+            _fetch_wsgi_text(app, serialize=True),
+            _fetch_wsgi_text(app, serialize=True),
+        ),
+        timeout=5,
+    )
+    assert app.max_active == 1
+
+
+@pytest.mark.asyncio
+async def test_default_wsgi_requests_can_overlap():
+    app = _ConcurrentApp()
+    await asyncio.wait_for(
+        asyncio.gather(
+            _fetch_wsgi_text(app),
+            _fetch_wsgi_text(app),
+        ),
+        timeout=5,
+    )
+    assert app.max_active == 2
+
+
+@pytest.mark.asyncio
+async def test_serialized_wsgi_streams_and_closes_original_iterable():
+    class Iterable:
+        def __init__(self):
+            self.closed = False
+
+        def __iter__(self):
+            yield b"first"
+            yield b"second"
+
+        def close(self):
+            self.closed = True
+
+    result = Iterable()
+    seen_input = None
+    seen_body = None
+
+    def app(environ, start_response):
+        nonlocal seen_body, seen_input
+        seen_input = environ["wsgi.input"]
+        seen_body = seen_input.read()
+        write = start_response(
+            "200 Everything", [("Set-Cookie", "a=1"), ("Set-Cookie", "b=2")]
+        )
+        write(b"before")
+        return result
+
+    response = await wsgi.fetch(
+        app,
+        js.Request.new("http://example.com/", method="POST", body="request"),
+        env,
+        serialize=True,
+    )
+    assert seen_input is not None
+    assert seen_body == b"request"
+    assert response.status == 200
+    assert response.statusText == "Everything"
+    assert not result.closed
+    assert not seen_input.closed
+    assert await response.text() == "beforefirstsecond"
+    assert result.closed
+    assert seen_input.closed
+    assert response.headers.get("set-cookie") == "a=1, b=2"
+
+
+@pytest.mark.asyncio
+async def test_serialized_wsgi_iteration_error_releases_lock():
+    def failing_app(environ, start_response):
+        start_response("200 OK", [])
+
+        def fail():
+            raise RuntimeError("iteration failed")
+            yield b"unreachable"
+
+        return fail()
+
+    with pytest.raises(RuntimeError, match="iteration failed"):
+        await asyncio.wait_for(
+            wsgi.fetch(
+                failing_app, js.Request.new("http://example.com/"), env, serialize=True
+            ),
+            timeout=5,
+        )
+    response = await asyncio.wait_for(
+        wsgi.fetch(
+            header_echo_app, js.Request.new("http://example.com/"), env, serialize=True
+        ),
+        timeout=5,
+    )
+    assert await response.text() == "Hello, World"