Read request bodies lazily, fix wsgi performance regressions

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

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))

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,
                     )