diff --git a/packages/runtime-sdk/src/workers/_serialize.py b/packages/runtime-sdk/src/workers/_serialize.py new file mode 100644 index 00000000..dee3c8dd --- /dev/null +++ b/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() diff --git a/packages/runtime-sdk/src/workers/wsgi.py b/packages/runtime-sdk/src/workers/wsgi.py index b9af4043..0d7155a9 100644 --- a/packages/runtime-sdk/src/workers/wsgi.py +++ b/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 diff --git a/packages/runtime-sdk/tests/workerd-test/wsgi/tests/test_wsgi.py b/packages/runtime-sdk/tests/workerd-test/wsgi/tests/test_wsgi.py index bfcec792..bb26a8f0 100644 --- a/packages/runtime-sdk/tests/workerd-test/wsgi/tests/test_wsgi.py +++ b/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"