[PR #9771/e80d5854 backport][3.10] Small cleanups to enabling compression in web_response (#9776)

Author: bdraco

aiohttp/web_response.py

@@ -42,6 +42,7 @@
 from .typedefs import JSONEncoder, LooseHeaders
 
 REASON_PHRASES = {http_status.value: http_status.phrase for http_status in HTTPStatus}
+LARGE_BODY_SIZE = 1024**2
 
 __all__ = ("ContentCoding", "StreamResponse", "Response", "json_response")
 
@@ -397,25 +398,26 @@ def _generate_content_type_header(
         self._headers[CONTENT_TYPE] = ctype
 
     async def _do_start_compression(self, coding: ContentCoding) -> None:
-        if coding != ContentCoding.identity:
-            assert self._payload_writer is not None
-            self._headers[hdrs.CONTENT_ENCODING] = coding.value
-            self._payload_writer.enable_compression(coding.value)
-            # Compressed payload may have different content length,
-            # remove the header
-            self._headers.popall(hdrs.CONTENT_LENGTH, None)
+        if coding is ContentCoding.identity:
+            return
+        assert self._payload_writer is not None
+        self._headers[hdrs.CONTENT_ENCODING] = coding.value
+        self._payload_writer.enable_compression(coding.value)
+        # Compressed payload may have different content length,
+        # remove the header
+        self._headers.popall(hdrs.CONTENT_LENGTH, None)
 
     async def _start_compression(self, request: "BaseRequest") -> None:
         if self._compression_force:
             await self._do_start_compression(self._compression_force)
-        else:
-            # Encoding comparisons should be case-insensitive
-            # https://www.rfc-editor.org/rfc/rfc9110#section-8.4.1
-            accept_encoding = request.headers.get(hdrs.ACCEPT_ENCODING, "").lower()
-            for value, coding in CONTENT_CODINGS.items():
-                if value in accept_encoding:
-                    await self._do_start_compression(coding)
-                    return
+            return
+        # Encoding comparisons should be case-insensitive
+        # https://www.rfc-editor.org/rfc/rfc9110#section-8.4.1
+        accept_encoding = request.headers.get(hdrs.ACCEPT_ENCODING, "").lower()
+        for value, coding in CONTENT_CODINGS.items():
+            if value in accept_encoding:
+                await self._do_start_compression(coding)
+                return
 
     async def prepare(self, request: "BaseRequest") -> Optional[AbstractStreamWriter]:
         if self._eof_sent:
@@ -765,30 +767,28 @@ async def _start(self, request: "BaseRequest") -> AbstractStreamWriter:
     async def _do_start_compression(self, coding: ContentCoding) -> None:
         if self._chunked or isinstance(self._body, Payload):
             return await super()._do_start_compression(coding)
-
-        if coding != ContentCoding.identity:
-            # Instead of using _payload_writer.enable_compression,
-            # compress the whole body
-            compressor = ZLibCompressor(
-                encoding=str(coding.value),
-                max_sync_chunk_size=self._zlib_executor_size,
-                executor=self._zlib_executor,
-            )
-            assert self._body is not None
-            if self._zlib_executor_size is None and len(self._body) > 1024 * 1024:
-                warnings.warn(
-                    "Synchronous compression of large response bodies "
-                    f"({len(self._body)} bytes) might block the async event loop. "
-                    "Consider providing a custom value to zlib_executor_size/"
-                    "zlib_executor response properties or disabling compression on it."
-                )
-            self._compressed_body = (
-                await compressor.compress(self._body) + compressor.flush()
+        if coding is ContentCoding.identity:
+            return
+        # Instead of using _payload_writer.enable_compression,
+        # compress the whole body
+        compressor = ZLibCompressor(
+            encoding=coding.value,
+            max_sync_chunk_size=self._zlib_executor_size,
+            executor=self._zlib_executor,
+        )
+        assert self._body is not None
+        if self._zlib_executor_size is None and len(self._body) > LARGE_BODY_SIZE:
+            warnings.warn(
+                "Synchronous compression of large response bodies "
+                f"({len(self._body)} bytes) might block the async event loop. "
+                "Consider providing a custom value to zlib_executor_size/"
+                "zlib_executor response properties or disabling compression on it."
             )
-            assert self._compressed_body is not None
-
-            self._headers[hdrs.CONTENT_ENCODING] = coding.value
-            self._headers[hdrs.CONTENT_LENGTH] = str(len(self._compressed_body))
+        self._compressed_body = (
+            await compressor.compress(self._body) + compressor.flush()
+        )
+        self._headers[hdrs.CONTENT_ENCODING] = coding.value
+        self._headers[hdrs.CONTENT_LENGTH] = str(len(self._compressed_body))
 
 
 def json_response(

tests/test_web_response.py

@@ -480,6 +480,24 @@ async def test_force_compression_deflate() -> None:
     assert "deflate" == resp.headers.get(hdrs.CONTENT_ENCODING)
 
 
+async def test_force_compression_deflate_large_payload() -> None:
+    """Make sure a warning is thrown for large payloads compressed in the event loop."""
+    req = make_request(
+        "GET", "/", headers=CIMultiDict({hdrs.ACCEPT_ENCODING: "gzip, deflate"})
+    )
+    resp = Response(body=b"large")
+
+    resp.enable_compression(ContentCoding.deflate)
+    assert resp.compression
+
+    with pytest.warns(
+        Warning, match="Synchronous compression of large response bodies"
+    ), mock.patch("aiohttp.web_response.LARGE_BODY_SIZE", 2):
+        msg = await resp.prepare(req)
+        assert msg is not None
+    assert "deflate" == resp.headers.get(hdrs.CONTENT_ENCODING)
+
+
 async def test_force_compression_no_accept_deflate() -> None:
     req = make_request("GET", "/")
     resp = StreamResponse()