From 4417317335d2563395575dbeea6e396e5e25455f Mon Sep 17 00:00:00 2001 From: Paul Osborne Date: Mon, 23 Mar 2026 10:57:01 -0500 Subject: [PATCH] Read request bodies lazily, fix wsgi performance regressions With a large maximum (4GB) specified for http_body.read, we ended up with a very large allocation taking place in the host. In the case of this hostcall, at least, we do need to read the body in chunks for reasonable performance. This change fixes that by ensuring that we perform reads in chunks; additionally, we now use a BufferedReader over our own RawIOBase which allows us to lazily read the bodies until later in execution (or not at all if the request body isn't read by the handler). --- fastly_compute/requests/response.py | 9 +-- fastly_compute/utils.py | 98 +++++++++++++++++++++-------- fastly_compute/wsgi.py | 28 +++------ 3 files changed, 85 insertions(+), 50 deletions(-) diff --git a/fastly_compute/requests/response.py b/fastly_compute/requests/response.py index 4963c5d..04deb91 100644 --- a/fastly_compute/requests/response.py +++ b/fastly_compute/requests/response.py @@ -6,7 +6,7 @@ from wit_world.imports import async_io, http_resp -from ..utils import read_response_body +from ..utils import create_body_reader from .exceptions import HTTPError @@ -85,7 +85,8 @@ def headers(self) -> dict[str, str]: def content(self) -> bytes: """Response body as bytes.""" if self._content is None: - self._content = self._read_body() + reader = create_body_reader(self._response_body) + self._content = reader.read() return self._content @property @@ -172,10 +173,6 @@ def encoding(self) -> str | None: """Response encoding.""" return self._parse_charset() - def _read_body(self) -> bytes: - """Read the complete response body from WIT.""" - return read_response_body(self._response_body) - def __bool__(self) -> bool: """Boolean evaluation returns ok status.""" return self.ok diff --git a/fastly_compute/utils.py b/fastly_compute/utils.py index ce29e92..acb6641 100644 --- a/fastly_compute/utils.py +++ b/fastly_compute/utils.py @@ -1,36 +1,82 @@ """Utility functions for fastly_compute package.""" +from io import BufferedReader, RawIOBase + from wit_world.imports import async_io, http_body -from fastly_compute.exceptions.types.error import Error -from fastly_compute.requests.exceptions import RequestException +class _RawBodyReader(RawIOBase): + """Raw I/O implementation that reads from a Fastly HTTP body. + + This class provides a streaming interface to read from Fastly HTTP bodies + without buffering the entire body in memory upfront. + """ + + def __init__(self, body_handle: async_io.Pollable, chunk_size: int = 4096): + """Initialize the reader. + + :param body_handle: Fastly HTTP body handle to read from + :param chunk_size: Size of chunks to read at a time (default: 4096) + """ + self._body = body_handle + self._chunk_size = chunk_size + self._closed = False + + def readable(self) -> bool: + """Return whether the stream is readable. + + :return: True if the stream is readable, False otherwise + """ + return not self._closed + + def readinto(self, b) -> int: + """Read up to len(b) bytes into writable buffer b. + + :param b: Writable buffer to read into + :return: Number of bytes read, or 0 on EOF + """ + if self._closed: + return 0 + + # Read a chunk from host into provided buffer + chunk_size = min(len(b), self._chunk_size) + chunk = http_body.read(self._body, chunk_size) + + if not chunk: + self._closed = True + return 0 + + # Copy chunk into the provided buffer + n = len(chunk) + b[:n] = chunk + return n + + +def create_body_reader( + body: async_io.Pollable, chunk_size: int = 4096 +) -> BufferedReader: + """Create a file-like reader for streaming a Fastly HTTP body. + + This function returns a BufferedReader that streams the body on-demand, + avoiding the need to buffer the entire body in memory. The returned reader + is compatible with the WSGI InputStream Protocol (PEP 3333). -def read_response_body( - response_body: async_io.Pollable, chunk_size: int = 4096 -) -> bytes: - """Read the complete response body from a WIT response body object. + The returned reader supports all standard file operations: read(), readline(), + readlines(), and iteration. Note that seeking is not supported as the body + is a forward-only stream. - Args: - response_body: WIT response body object to read from - chunk_size: Size of chunks to read at a time (default: 4096) + :param body: Fastly HTTP body handle to read from + :param chunk_size: Size of chunks to read at a time (default: 4096) + :return: A BufferedReader that streams the body content (WSGI compatible) - Returns: - Complete response body as bytes + Example:: - Raises: - RequestException: If there is a problem reading the response body. + def handle(request, body): + reader = create_body_reader(body) + # Read entire body + full_body = reader.read() + # Or read line by line + for line in reader: + process(line) """ - body_data: bytes = b"" - while True: - try: - chunk = http_body.read(response_body, chunk_size) - except Error as e: - raise RequestException.from_fastly_error(e, "http_body.read") from e - - if len(chunk) == 0: - break - else: - body_data += chunk - - return body_data + return BufferedReader(_RawBodyReader(body, chunk_size)) diff --git a/fastly_compute/wsgi.py b/fastly_compute/wsgi.py index 265aff2..cca949b 100644 --- a/fastly_compute/wsgi.py +++ b/fastly_compute/wsgi.py @@ -11,9 +11,9 @@ import sys import traceback from collections.abc import Callable -from io import BytesIO, Reader from typing import Any from urllib.parse import urlparse +from wsgiref.types import InputStream from wit_world.exports import HttpIncoming as WitHttpIncoming from wit_world.imports import async_io, http_body, http_req, http_resp @@ -25,11 +25,12 @@ from wit_world.imports.http_resp import send_downstream from fastly_compute.exceptions.types.error import CannotRead +from fastly_compute.utils import create_body_reader def serve_wsgi_request( req: http_req.Request, - body: Reader[bytes], + body: InputStream, app: Callable, handle_errors: bool = False, ) -> None: @@ -39,13 +40,11 @@ def serve_wsgi_request( specification, allowing any WSGI-compatible web framework to run on Fastly Compute. - Args: - req: Fastly HTTP request object from WIT bindings - body: Fastly HTTP body object from WIT bindings - app: WSGI application callable - handle_errors: If True, the wrapper will log exceptions and return 500; if not, - then it will be handled by the server (or will be handled by - the WSGI app/framework itself). + :param req: Fastly HTTP request object from WIT bindings + :param body: WSGI input stream containing the request body (PEP 3333 compliant) + :param app: WSGI application callable + :param handle_errors: If True, log exceptions and return 500; otherwise let + the server or WSGI app handle them """ response = http_resp.Response.new() response_body = http_body.new() @@ -208,17 +207,10 @@ def __call__(self): def handle(self, request: http_req.Request, body: async_io.Pollable) -> None: """Handle incoming HTTP requests by serving them through the WSGI app.""" - - def body_reader(body: async_io.Pollable) -> Reader[bytes]: - """Given a Fastly HTTP body object, return a file-like object - containing the body's content. - """ - return BytesIO(http_body.read(body, 2**32 - 1)) - with request: # Ensure dropping of request resource before trying to get another one. This dodges a crash. serve_wsgi_request( request, - body_reader(body), + create_body_reader(body), self.wsgi_app, handle_errors=self.handle_errors, ) @@ -241,7 +233,7 @@ def body_reader(body: async_io.Pollable) -> Reader[bytes]: with request: serve_wsgi_request( request, - body_reader(body), + create_body_reader(body), self.wsgi_app, handle_errors=self.handle_errors, )