Implement zero copy writes for the WebSocket writer

CHANGES/9634.feature.rst

@@ -0,0 +1 @@
+Implemented zero copy writes for WebSockets when using Python 3.12+ -- by :user:`bdraco`.

aiohttp/_websocket/writer.py

@@ -135,12 +135,16 @@ async def send_frame(
             mask = PACK_RANDBITS(self.get_random_bits())
             message = bytearray(message)
             websocket_mask(mask, message)
-            self.transport.write(header + mask + message)
+            self.transport.writelines((header, mask, message))
             self._output_size += MASK_LEN
         elif msg_length > MSG_SIZE:
-            self.transport.write(header)
-            self.transport.write(message)
+            # For large messages, we use writelines to avoid copying the
+            # entire message into a new buffer. This is a performance
+            # optimization to avoid unnecessary memory allocations.
+            self.transport.writelines((header, message))
         else:
+            # If the message is small, its faster to copy it into a new
+            # buffer and send it all at once.
             self.transport.write(header + message)
 
         self._output_size += header_len + msg_length

tests/test_benchmarks_http_websocket.py

@@ -1,10 +1,12 @@
 """codspeed benchmarks for http websocket."""
 
 import asyncio
+from typing import Iterable, Union
 
 from pytest_codspeed import BenchmarkFixture
 
 from aiohttp import DataQueue
+from aiohttp._websocket.helpers import MSG_SIZE
 from aiohttp.base_protocol import BaseProtocol
 from aiohttp.http_websocket import (
     WebSocketReader,
@@ -41,7 +43,10 @@ def is_closing(self) -> bool:
         """Swallow is_closing."""
         return False
 
-    def write(self, data: bytes) -> None:
+    def write(self, data: Iterable[Union[bytes, bytearray, memoryview]]) -> None:
+        """Swallow writes."""
+
+    def writelines(self, data: Iterable[Union[bytes, bytearray, memoryview]]) -> None:
         """Swallow writes."""
 
 
@@ -67,6 +72,22 @@ def _run() -> None:
         loop.run_until_complete(_send_one_hundred_websocket_text_messages())
 
 
+def test_send_one_hundred_large_websocket_text_messages(
+    loop: asyncio.AbstractEventLoop, benchmark: BenchmarkFixture
+) -> None:
+    """Benchmark sending 100 WebSocket text messages."""
+    writer = WebSocketWriter(MockProtocol(loop=loop), MockTransport())
+    raw_message = b"x" * MSG_SIZE * 2
+
+    async def _send_one_hundred_websocket_text_messages() -> None:
+        for _ in range(100):
+            await writer.send_frame(raw_message, WSMsgType.TEXT)
+
+    @benchmark
+    def _run() -> None:
+        loop.run_until_complete(_send_one_hundred_websocket_text_messages())
+
+
 def test_send_one_hundred_websocket_text_messages_with_mask(
     loop: asyncio.AbstractEventLoop, benchmark: BenchmarkFixture
 ) -> None: