Repository: aio-libs/aiohttp

Issue:
 ZLibCompressor issue leading to Websocket anomalies
### Describe the bug

After upgrading aiohttp to version 3.9.0, I encountered an intermittent error where the browser reports an 'Invalid frame header' during websocket access. 
On delving deeper, I found that the websocket module in aiohttp, when compressing the data, replaced zlib.compressobj with ZLibCompressor to support max_sync_chunk_size. The specific change can be traced to this commit: https://github.com/aio-libs/aiohttp/commit/586778f19b67a26a7d7ed9cd20add9b5b4dea96c. However, ZLibCompressor was only initialized once in the WebSocketWriter and shared across multiple coroutines. However, its logic does not support usage in multiple coroutines. When different coroutines simultaneously call the websocket to send messages, the compressed results get mixed up, leading to frame inconsistencies.
After replacing ZLibCompressor with an older version of zlib.compressobj(mostly remove await on compressobj.compress to avoid coroutine switching), everything returned to normal.

### To Reproduce

```
from aiohttp.compression_utils import ZLibCompressor
import zlib
import asyncio
import random
import string

compressobj = ZLibCompressor(
    level=zlib.Z_BEST_SPEED,
    wbits=-15,
    max_sync_chunk_size=5 * 1024,
)


def generate_random_string(length):
    letters = string.ascii_lowercase
    return ''.join(random.choice(letters) for i in range(length))


async def test(compressobj):
    # test code from aiohttp.http_websocket.py: _send_frame
    # https://github.com/aio-libs/aiohttp/commit/586778f19b67a26a7d7ed9cd20add9b5b4dea96c#diff-29db742fa68fd59f39e1ca70379c4d478101a8ed9af8895c8fc7604bbf1b4b11R625
    message = generate_random_string(1232323).encode()
    message = await compressobj.compress(message)  # await cause coroutine switching
    message += compressobj.flush(zlib.Z_SYNC_FLUSH)  # data maybe from other coroutine
    return message


loop = asyncio.get_event_loop()
# parallel running 10 coroutines
tasks = [test(compressobj) for i in range(10)]
messages = loop.run_until_complete(asyncio.gather(*tasks))

for message in messages:
    decompressor = zlib.decompressobj(wbits=-15)
    message = decompressor.decompress(message) # E   zlib.error: Error -3 while decompressing data: invalid stored block lengths
```


### Expected behavior

correct compression

### Logs/tracebacks

```python-traceback
Invalid frame header
```


### Python Version

```console
Python 3.11.1
```


### aiohttp Version

```console
Name: aiohttp
Version: 3.9.0
Summary: Async http client/server framework (asyncio)
Home-page: https://github.com/aio-libs/aiohttp
Author: 
Author-email: 
License: Apache 2
Location: /Users/xiaohuanshu/PycharmProjects/bazhen/venv/lib/python3.11/site-packages
Requires: aiosignal, attrs, frozenlist, multidict, yarl
Required-by: aiohttp-debugtoolbar, aiohttp-devtools, aiohttp-jinja2, aiohttp-pydantic, pytest-aiohttp
```


### multidict Version

```console
Name: multidict
Version: 6.0.4
Summary: multidict implementation
Home-page: https://github.com/aio-libs/multidict
Author: Andrew Svetlov
Author-email: andrew.svetlov@gmail.com
License: Apache 2
Location: /Users/xiaohuanshu/PycharmProjects/bazhen/venv/lib/python3.11/site-packages
Requires: 
Required-by: aiohttp, yarl
```


### yarl Version

```console
Name: yarl
Version: 1.9.2
Summary: Yet another URL library
Home-page: https://github.com/aio-libs/yarl/
Author: Andrew Svetlov
Author-email: andrew.svetlov@gmail.com
License: Apache-2.0
Location: /Users/xiaohuanshu/PycharmProjects/bazhen/venv/lib/python3.11/site-packages
Requires: idna, multidict
Required-by: aiohttp
```


### OS

macOS

### Related component

Server

### Additional context

_No response_

### Code of Conduct

- [X] I agree to follow the aio-libs Code of Conduct


Discussion:

Resolving patch:
diff --git a/aiohttp/compression_utils.py b/aiohttp/compression_utils.py
--- a/aiohttp/compression_utils.py
+++ b/aiohttp/compression_utils.py
@@ -62,19 +62,25 @@ def __init__(
             self._compressor = zlib.compressobj(
                 wbits=self._mode, strategy=strategy, level=level
             )
+        self._compress_lock = asyncio.Lock()
 
     def compress_sync(self, data: bytes) -> bytes:
         return self._compressor.compress(data)
 
     async def compress(self, data: bytes) -> bytes:
-        if (
-            self._max_sync_chunk_size is not None
-            and len(data) > self._max_sync_chunk_size
-        ):
-            return await asyncio.get_event_loop().run_in_executor(
-                self._executor, self.compress_sync, data
-            )
-        return self.compress_sync(data)
+        async with self._compress_lock:
+            # To ensure the stream is consistent in the event
+            # there are multiple writers, we need to lock
+            # the compressor so that only one writer can
+            # compress at a time.
+            if (
+                self._max_sync_chunk_size is not None
+                and len(data) > self._max_sync_chunk_size
+            ):
+                return await asyncio.get_event_loop().run_in_executor(
+                    self._executor, self.compress_sync, data
+                )
+            return self.compress_sync(data)
 
     def flush(self, mode: int = zlib.Z_FINISH) -> bytes:
         return self._compressor.flush(mode)
diff --git a/aiohttp/http_websocket.py b/aiohttp/http_websocket.py
--- a/aiohttp/http_websocket.py
+++ b/aiohttp/http_websocket.py
@@ -626,21 +626,17 @@ async def _send_frame(
         if (compress or self.compress) and opcode < 8:
             if compress:
                 # Do not set self._compress if compressing is for this frame
-                compressobj = ZLibCompressor(
-                    level=zlib.Z_BEST_SPEED,
-                    wbits=-compress,
-                    max_sync_chunk_size=WEBSOCKET_MAX_SYNC_CHUNK_SIZE,
-                )
+                compressobj = self._make_compress_obj(compress)
             else:  # self.compress
                 if not self._compressobj:
-                    self._compressobj = ZLibCompressor(
-                        level=zlib.Z_BEST_SPEED,
-                        wbits=-self.compress,
-                        max_sync_chunk_size=WEBSOCKET_MAX_SYNC_CHUNK_SIZE,
-                    )
+                    self._compressobj = self._make_compress_obj(self.compress)
                 compressobj = self._compressobj
 
             message = await compressobj.compress(message)
+            # Its critical that we do not return control to the event
+            # loop until we have finished sending all the compressed
+            # data. Otherwise we could end up mixing compressed frames
+            # if there are multiple coroutines compressing data.
             message += compressobj.flush(
                 zlib.Z_FULL_FLUSH if self.notakeover else zlib.Z_SYNC_FLUSH
             )
@@ -678,10 +674,20 @@ async def _send_frame(
 
             self._output_size += len(header) + len(message)
 
+        # It is safe to return control to the event loop when using compression
+        # after this point as we have already sent or buffered all the data.
+
         if self._output_size > self._limit:
             self._output_size = 0
             await self.protocol._drain_helper()
 
+    def _make_compress_obj(self, compress: int) -> ZLibCompressor:
+        return ZLibCompressor(
+            level=zlib.Z_BEST_SPEED,
+            wbits=-compress,
+            max_sync_chunk_size=WEBSOCKET_MAX_SYNC_CHUNK_SIZE,
+        )
+
     def _write(self, data: bytes) -> None:
         if self.transport.is_closing():
             raise ConnectionResetError("Cannot write to closing transport")


Test patch:
diff --git a/tests/test_websocket_writer.py b/tests/test_websocket_writer.py
--- a/tests/test_websocket_writer.py
+++ b/tests/test_websocket_writer.py
@@ -1,11 +1,13 @@
 # type: ignore
+import asyncio
 import random
-from typing import Any
+from typing import Any, Callable
 from unittest import mock
 
 import pytest
 
-from aiohttp.http import WebSocketWriter
+from aiohttp import DataQueue, WSMessage
+from aiohttp.http import WebSocketReader, WebSocketWriter
 from aiohttp.test_utils import make_mocked_coro
 
 
@@ -106,3 +108,65 @@ async def test_send_compress_text_per_message(protocol: Any, transport: Any) ->
     writer.transport.write.assert_called_with(b"\x81\x04text")
     await writer.send(b"text", compress=15)
     writer.transport.write.assert_called_with(b"\xc1\x06*I\xad(\x01\x00")
+
+
+@pytest.mark.parametrize(
+    ("max_sync_chunk_size", "payload_point_generator"),
+    (
+        (16, lambda count: count),
+        (4096, lambda count: count),
+        (32, lambda count: 64 + count if count % 2 else count),
+    ),
+)
+async def test_concurrent_messages(
+    protocol: Any,
+    transport: Any,
+    max_sync_chunk_size: int,
+    payload_point_generator: Callable[[int], int],
+) -> None:
+    """Ensure messages are compressed correctly when there are multiple concurrent writers.
+
+    This test generates is parametrized to
+
+    - Generate messages that are larger than patch
+      WEBSOCKET_MAX_SYNC_CHUNK_SIZE of 16
+      where compression will run in the executor
+
+    - Generate messages that are smaller than patch
+      WEBSOCKET_MAX_SYNC_CHUNK_SIZE of 4096
+      where compression will run in the event loop
+
+    - Interleave generated messages with a
+      WEBSOCKET_MAX_SYNC_CHUNK_SIZE of 32
+      where compression will run in the event loop
+      and in the executor
+    """
+    with mock.patch(
+        "aiohttp.http_websocket.WEBSOCKET_MAX_SYNC_CHUNK_SIZE", max_sync_chunk_size
+    ):
+        writer = WebSocketWriter(protocol, transport, compress=15)
+        queue: DataQueue[WSMessage] = DataQueue(asyncio.get_running_loop())
+        reader = WebSocketReader(queue, 50000)
+        writers = []
+        payloads = []
+        for count in range(1, 64 + 1):
+            point = payload_point_generator(count)
+            payload = bytes((point,)) * point
+            payloads.append(payload)
+            writers.append(writer.send(payload, binary=True))
+        await asyncio.gather(*writers)
+
+    for call in writer.transport.write.call_args_list:
+        call_bytes = call[0][0]
+        result, _ = reader.feed_data(call_bytes)
+        assert result is False
+        msg = await queue.read()
+        bytes_data: bytes = msg.data
+        first_char = bytes_data[0:1]
+        char_val = ord(first_char)
+        assert len(bytes_data) == char_val
+        # If we have a concurrency problem, the data
+        # tends to get mixed up between messages so
+        # we want to validate that all the bytes are
+        # the same value
+        assert bytes_data == bytes_data[0:1] * char_val