Decode compressed response bodies incrementally (#1126)

Author: Kludex

src/httpx2/httpx2/_client.py

@@ -160,8 +160,13 @@ def __init__(self, stream: AsyncByteStream, start: float) -> None:
         self.elapsed: datetime.timedelta | None = None
 
     async def __aiter__(self) -> typing.AsyncIterator[bytes]:
-        async for chunk in self._stream:
-            yield chunk
+        stream = self._stream.__aiter__()
+        try:
+            async for chunk in stream:
+                yield chunk
+        finally:
+            if isinstance(stream, AsyncGenerator):
+                await stream.aclose()
 
     async def aclose(self) -> None:
         self.elapsed = datetime.timedelta(seconds=time.perf_counter() - self._start)

src/httpx2/httpx2/_content.py

@@ -2,7 +2,7 @@
 
 import inspect
 import warnings
-from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator, Mapping
+from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Iterable, Iterator, Mapping
 from json import dumps as json_dumps
 from typing import (
     Any,
@@ -79,9 +79,15 @@ async def __aiter__(self) -> AsyncIterator[bytes]:
                 yield chunk
                 chunk = await self._stream.aread(self.CHUNK_SIZE)
         else:
-            # Otherwise iterate.
-            async for part in self._stream:
-                yield part
+            # Otherwise iterate, making sure the wrapped stream is closed even if the
+            # consumer stops early (e.g. an exception is raised part-way through decoding).
+            stream = self._stream.__aiter__()
+            try:
+                async for part in stream:
+                    yield part
+            finally:
+                if isinstance(stream, AsyncGenerator):
+                    await stream.aclose()
 
 
 class UnattachedStream(AsyncByteStream, SyncByteStream):

src/httpx2/httpx2/_decoders.py

@@ -8,6 +8,7 @@
 
 import codecs
 import io
+import itertools
 import sys
 import typing
 import zlib
@@ -48,11 +49,42 @@
         _zstandard_installed = False
 
 
+MAX_DECODE_CHUNK_SIZE = 2**20  # 1 MiB
+
+
+class Decompressor(typing.Protocol):
+    @property
+    def unconsumed_tail(self) -> bytes: ...
+
+    def decompress(self, data: bytes, max_length: int) -> bytes: ...
+
+    def flush(self) -> bytes: ...
+
+
+class ZlibDecompressor:
+    """
+    Drain a `zlib`/`gzip` decompressor in bounded pieces so a small compressed
+    input cannot inflate to an unbounded buffer in a single call.
+    """
+
+    def __init__(self, decompressor: Decompressor) -> None:
+        self.decompressor = decompressor
+
+    def decompress(self, data: bytes) -> typing.Iterator[bytes]:
+        decompressed = self.decompressor.decompress(data, MAX_DECODE_CHUNK_SIZE)
+        while decompressed:
+            yield decompressed
+            decompressed = self.decompressor.decompress(self.decompressor.unconsumed_tail, MAX_DECODE_CHUNK_SIZE)
+
+    def flush(self) -> bytes:
+        return self.decompressor.flush()
+
+
 class ContentDecoder:
-    def decode(self, data: bytes) -> bytes:
+    def decode(self, data: bytes) -> typing.Iterator[bytes]:
         raise NotImplementedError()  # pragma: no cover
 
-    def flush(self) -> bytes:
+    def flush(self) -> typing.Iterator[bytes]:
         raise NotImplementedError()  # pragma: no cover
 
 
@@ -61,11 +93,11 @@ class IdentityDecoder(ContentDecoder):
     Handle unencoded data.
     """
 
-    def decode(self, data: bytes) -> bytes:
-        return data
+    def decode(self, data: bytes) -> typing.Iterator[bytes]:
+        yield data
 
-    def flush(self) -> bytes:
-        return b""
+    def flush(self) -> typing.Iterator[bytes]:
+        yield from ()
 
 
 class DeflateDecoder(ContentDecoder):
@@ -77,22 +109,23 @@ class DeflateDecoder(ContentDecoder):
 
     def __init__(self) -> None:
         self.first_attempt = True
-        self.decompressor = zlib.decompressobj()
+        self.decompressor = ZlibDecompressor(zlib.decompressobj())
 
-    def decode(self, data: bytes) -> bytes:
+    def decode(self, data: bytes) -> typing.Iterator[bytes]:
         was_first_attempt = self.first_attempt
         self.first_attempt = False
         try:
-            return self.decompressor.decompress(data)
+            yield from self.decompressor.decompress(data)
         except zlib.error as exc:
             if was_first_attempt:
-                self.decompressor = zlib.decompressobj(-zlib.MAX_WBITS)
-                return self.decode(data)
-            raise DecodingError(str(exc)) from exc
+                self.decompressor = ZlibDecompressor(zlib.decompressobj(-zlib.MAX_WBITS))
+                yield from self.decode(data)
+            else:
+                raise DecodingError(str(exc)) from exc
 
-    def flush(self) -> bytes:
+    def flush(self) -> typing.Iterator[bytes]:
         try:
-            return self.decompressor.flush()
+            yield self.decompressor.flush()
         except zlib.error as exc:  # pragma: no cover
             raise DecodingError(str(exc)) from exc
 
@@ -105,17 +138,17 @@ class GZipDecoder(ContentDecoder):
     """
 
     def __init__(self) -> None:
-        self.decompressor = zlib.decompressobj(zlib.MAX_WBITS | 16)
+        self.decompressor = ZlibDecompressor(zlib.decompressobj(zlib.MAX_WBITS | 16))
 
-    def decode(self, data: bytes) -> bytes:
+    def decode(self, data: bytes) -> typing.Iterator[bytes]:
         try:
-            return self.decompressor.decompress(data)
+            yield from self.decompressor.decompress(data)
         except zlib.error as exc:
             raise DecodingError(str(exc)) from exc
 
-    def flush(self) -> bytes:
+    def flush(self) -> typing.Iterator[bytes]:
         try:
-            return self.decompressor.flush()
+            yield self.decompressor.flush()
         except zlib.error as exc:  # pragma: no cover
             raise DecodingError(str(exc)) from exc
 
@@ -140,26 +173,31 @@ def __init__(self) -> None:
 
         self.decompressor = brotli.Decompressor()
         self.seen_data = False
-        self._decompress: typing.Callable[[bytes], bytes]
+        self._decompress: typing.Callable[..., bytes]
         if hasattr(self.decompressor, "decompress"):
             # The 'brotlicffi' package.
             self._decompress = self.decompressor.decompress  # pragma: no cover
         else:
             # The 'brotli' package.
-            self._decompress = self.decompressor.process  # pragma: no cover
+            self._decompress = self.decompressor.process
 
-    def decode(self, data: bytes) -> bytes:
+    def decode(self, data: bytes) -> typing.Iterator[bytes]:
         if not data:
-            return b""
+            return
         self.seen_data = True
         try:
-            return self._decompress(data)
+            # The C backend may allocate nearly twice the requested threshold.
+            output_buffer_limit = MAX_DECODE_CHUNK_SIZE // 2
+            decompressed = self._decompress(data, output_buffer_limit=output_buffer_limit)
+            while decompressed:
+                yield decompressed
+                decompressed = self._decompress(b"", output_buffer_limit=output_buffer_limit)
         except brotli.error as exc:
             raise DecodingError(str(exc)) from exc
 
-    def flush(self) -> bytes:
+    def flush(self) -> typing.Iterator[bytes]:
         if not self.seen_data:
-            return b""
+            return
         try:
             if hasattr(self.decompressor, "finish"):
                 # Only available in the 'brotlicffi' package.
@@ -168,9 +206,9 @@ def flush(self) -> bytes:
                 # will never actually emit any data. However, it will potentially throw
                 # errors if a truncated or damaged data stream has been used.
                 self.decompressor.finish()  # pragma: no cover
-            return b""
         except brotli.error as exc:  # pragma: no cover
             raise DecodingError(str(exc)) from exc
+        yield from ()
 
 
 class ZStandardDecoder(ContentDecoder):
@@ -189,30 +227,34 @@ def __init__(self) -> None:
         self.decompressor = ZstdDecompressor()
         self.seen_data = False
 
-    def decode(self, data: bytes) -> bytes:
+    def decode(self, data: bytes) -> typing.Iterator[bytes]:
         if not data:
-            return b""
+            return
         self.seen_data = True
-        output = io.BytesIO()
         try:
             if self.decompressor.eof:
                 data = self.decompressor.unused_data + data
                 self.decompressor = ZstdDecompressor()
-            output.write(self.decompressor.decompress(data))
-            while self.decompressor.eof and self.decompressor.unused_data:
-                unused_data = self.decompressor.unused_data
+            while True:
+                decompressed = self.decompressor.decompress(data, MAX_DECODE_CHUNK_SIZE)
+                while decompressed:
+                    yield decompressed
+                    if self.decompressor.needs_input or self.decompressor.eof:
+                        break
+                    decompressed = self.decompressor.decompress(b"", MAX_DECODE_CHUNK_SIZE)
+                if not (self.decompressor.eof and self.decompressor.unused_data):
+                    break
+                data = self.decompressor.unused_data
                 self.decompressor = ZstdDecompressor()
-                output.write(self.decompressor.decompress(unused_data))
         except ZstdError as exc:
             raise DecodingError(str(exc)) from exc
-        return output.getvalue()
 
-    def flush(self) -> bytes:
+    def flush(self) -> typing.Iterator[bytes]:
         if not self.seen_data:
-            return b""
+            return
         if not self.decompressor.eof:
             raise DecodingError("Zstandard data is incomplete")  # pragma: no cover
-        return b""
+        yield from ()
 
 
 class MultiDecoder(ContentDecoder):
@@ -233,16 +275,25 @@ def __init__(self, encodings: typing.Sequence[str]) -> None:
         # Note that we reverse the order for decoding.
         self.children: list[ContentDecoder] = [SUPPORTED_DECODERS[coding]() for coding in reversed(codings)]
 
-    def decode(self, data: bytes) -> bytes:
+    def decode(self, data: bytes) -> typing.Iterator[bytes]:
+        streams: typing.Iterator[bytes] = iter((data,))
         for child in self.children:
-            data = child.decode(data)
-        return data
+            streams = self._pipe(child.decode, streams)
+        yield from streams
 
-    def flush(self) -> bytes:
-        data = b""
+    def flush(self) -> typing.Iterator[bytes]:
+        streams: typing.Iterator[bytes] = iter(())
         for child in self.children:
-            data = child.decode(data) + child.flush()
-        return data
+            streams = itertools.chain(self._pipe(child.decode, streams), child.flush())
+        yield from streams
+
+    @staticmethod
+    def _pipe(
+        decode: typing.Callable[[bytes], typing.Iterator[bytes]],
+        upstream: typing.Iterator[bytes],
+    ) -> typing.Iterator[bytes]:
+        for chunk in upstream:
+            yield from decode(chunk)
 
 
 class ByteChunker: