Add byte iter memory optimization to async media uploads

Author: tulir

mautrix/api.py

@@ -22,6 +22,7 @@
 
 from mautrix import __optional_imports__, __version__ as mautrix_version
 from mautrix.errors import MatrixConnectionError, MatrixRequestError, make_request_error
+from mautrix.util.async_iter_bytes import AsyncBody, async_iter_bytes
 from mautrix.util.logging import TraceLogger
 from mautrix.util.opt_prometheus import Counter
 
@@ -155,7 +156,6 @@ def replace(self, find: str, replace: str) -> PathBuilder:
 """
 
 _req_id = 0
-AsyncBody = AsyncGenerator[Union[bytes, bytearray, memoryview], None]
 
 
 def _next_global_req_id() -> int:
@@ -164,12 +164,6 @@ def _next_global_req_id() -> int:
     return _req_id
 
 
-async def _async_iter_bytes(data: bytearray | bytes, chunk_size: int = 1024**2) -> AsyncBody:
-    with memoryview(data) as mv:
-        for i in range(0, len(data), chunk_size):
-            yield mv[i : i + chunk_size]
-
-
 class HTTPAPI:
     """HTTPAPI is a simple asyncio Matrix API request sender."""
 
@@ -395,7 +389,7 @@ async def request(
                 method, log_url, content, orig_content, query_params, headers, req_id, sensitive
             )
             API_CALLS.labels(method=metrics_method).inc()
-            req_content = _async_iter_bytes(content) if do_fake_iter else content
+            req_content = async_iter_bytes(content) if do_fake_iter else content
             start = time.monotonic()
             try:
                 resp_data, resp = await self._send(

mautrix/client/api/modules/media_repository.py

@@ -20,6 +20,7 @@
     MXOpenGraph,
     SerializerError,
 )
+from mautrix.util.async_iter_bytes import async_iter_bytes
 from mautrix.util.opt_prometheus import Histogram
 
 from ..base import BaseClientAPI
@@ -286,15 +287,20 @@ async def _upload_to_url(
         headers: dict[str, str],
         data: bytes | bytearray | AsyncIterable[bytes],
         post_upload_query: dict[str, str],
+        min_iter_size: int = 25 * 1024 * 1024,
     ) -> None:
         retry_count = self.api.default_retry_count
         backoff = 4
+        do_fake_iter = data and hasattr(data, "__len__") and len(data) > min_iter_size
+        if do_fake_iter:
+            headers["Content-Length"] = str(len(data))
         while True:
             self.log.debug("Uploading media to external URL %s", upload_url)
             upload_response = None
             try:
+                req_data = async_iter_bytes(data) if do_fake_iter else data
                 upload_response = await self.api.session.put(
-                    upload_url, data=data, headers=headers
+                    upload_url, data=req_data, headers=headers
                 )
                 upload_response.raise_for_status()
             except Exception as e:

mautrix/util/__init__.py

@@ -1,22 +1,25 @@
 __all__ = [
+    "async_db",
+    "config",
+    "db",
     "formatter",
     "logging",
-    "config",
-    "signed_token",
-    "simple_template",
-    "manhole",
-    "markdown",
-    "simple_lock",
+    "async_getter_lock",
+    "async_iter_bytes",
+    "bridge_state",
+    "color_log",
+    "ffmpeg",
     "file_store",
-    "program",
-    "async_db",
-    "db",
-    "opt_prometheus",
+    "format_duration",
     "magic",
-    "bridge_state",
+    "manhole",
+    "markdown",
     "message_send_checkpoint",
-    "variation_selector",
-    "format_duration",
-    "ffmpeg",
+    "opt_prometheus",
+    "program",
+    "signed_token",
+    "simple_lock",
+    "simple_template",
     "utf16_surrogate",
+    "variation_selector",
 ]

mautrix/util/async_iter_bytes.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2023 Tulir Asokan
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+from typing import AsyncGenerator, Union
+
+AsyncBody = AsyncGenerator[Union[bytes, bytearray, memoryview], None]
+
+
+async def async_iter_bytes(data: bytearray | bytes, chunk_size: int = 1024**2) -> AsyncBody:
+    """
+    Return memory views into a byte array in chunks. This is used to prevent aiohttp from copying
+    the entire request body.
+
+    Args:
+        data: The underlying data to iterate through.
+        chunk_size: How big each returned chunk should be.
+
+    Returns:
+        An async generator that yields the given data in chunks.
+    """
+    with memoryview(data) as mv:
+        for i in range(0, len(data), chunk_size):
+            yield mv[i : i + chunk_size]
+
+
+__all__ = ["AsyncBody", "async_iter_bytes"]