From 489f411f6baad63e3989bf60266f4d5eecb13c69 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Thu, 10 Sep 2026 21:45:31 -0400 Subject: [PATCH] perf(video): buffer upload metadata in a bounded bytearray, not a list[bytes] The multipart parser hands the `metadata` field over in whatever pieces the client sent it, and `_VideoUploadStreamParser` kept one `bytes` object per piece. Per-object overhead then dominated: a client that dribbled the field in 2-byte chunks retained ~22x the payload, so the 1 MiB `MAX_UPLOAD_METADATA_SIZE` cap bounded payload but not memory (~22 MiB per upload, x2 concurrent slots). Accumulate into a `bytearray` instead, the way the header buffers already do. The buffer's own length is now the size the cap checks, so the separate `_metadata_size` counter goes away, and `_on_part_end` decodes the buffer directly. Tests: feed a 64 KiB field 1, 2 and 3 bytes at a time and assert retained allocation stays under 2x the payload (the old code retained 8.6x / 21.7x / 15.0x); and check the cap still counts the whole field across chunks. Closes #9563 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01L5Y1XMvgKRCr2pu6wm8rs1 --- invokeai/app/api/routers/videos.py | 21 +++---- tests/app/api/test_video_upload_limits.py | 72 +++++++++++++++++++++++ 2 files changed, 83 insertions(+), 10 deletions(-) diff --git a/invokeai/app/api/routers/videos.py b/invokeai/app/api/routers/videos.py index 9e9a9e890d2..2d1dd69054d 100644 --- a/invokeai/app/api/routers/videos.py +++ b/invokeai/app/api/routers/videos.py @@ -215,8 +215,12 @@ def __init__(self, destination: BinaryIO) -> None: self._header_value = bytearray() self._headers: dict[bytes, bytes] = {} self._part_name: Optional[bytes] = None - self._metadata_chunks: list[bytes] = [] - self._metadata_size = 0 + # A bytearray, not a list of per-callback chunks: the parser hands the metadata + # field over in whatever pieces the client sent it, and one `bytes` object per + # piece would let a client that dribbles the field in 2-byte pieces hold ~22x the + # size cap in memory. `len(self._metadata)` is the retained size, so the cap + # below bounds memory, not just payload. + self._metadata = bytearray() self.filename: Optional[str] = None self.content_type: Optional[str] = None self.metadata: Optional[str] = None @@ -247,8 +251,7 @@ def _on_part_begin(self) -> None: self._header_field = bytearray() self._header_value = bytearray() self._part_name = None - self._metadata_chunks = [] - self._metadata_size = 0 + self._metadata = bytearray() def _on_header_field(self, data: bytes, start: int, end: int) -> None: self._header_field.extend(data[start:end]) @@ -288,25 +291,23 @@ def _on_part_data(self, data: bytes, start: int, end: int) -> None: ) self._destination.write(chunk) elif self._part_name == b"metadata": - self._metadata_size += len(chunk) - if self._metadata_size > MAX_UPLOAD_METADATA_SIZE: + if len(self._metadata) + len(chunk) > MAX_UPLOAD_METADATA_SIZE: raise HTTPException( status_code=413, detail=f"Video metadata exceeds maximum size ({MAX_UPLOAD_METADATA_SIZE} bytes)", ) - self._metadata_chunks.append(chunk) + self._metadata.extend(chunk) # Any other field is dropped rather than buffered: an unknown part must not be a # way to make the server hold arbitrary bytes in memory. def _on_part_end(self) -> None: if self._part_name == b"metadata": try: - self.metadata = b"".join(self._metadata_chunks).decode("utf-8") + self.metadata = self._metadata.decode("utf-8") except UnicodeDecodeError as error: raise HTTPException(status_code=422, detail="Metadata must be UTF-8 encoded") from error self._part_name = None - self._metadata_chunks = [] - self._metadata_size = 0 + self._metadata = bytearray() def _on_end(self) -> None: self.saw_end = True diff --git a/tests/app/api/test_video_upload_limits.py b/tests/app/api/test_video_upload_limits.py index 726ecca835a..ca5941822bb 100644 --- a/tests/app/api/test_video_upload_limits.py +++ b/tests/app/api/test_video_upload_limits.py @@ -8,8 +8,10 @@ """ import asyncio +import gc import tempfile import time +import tracemalloc from pathlib import Path from types import SimpleNamespace from typing import Any @@ -1110,3 +1112,73 @@ async def send(message): # Only the first chunk (under the cap) reached the app as a body message. body_bytes = sum(len(m.get("body", b"")) for m in seen if m["type"] == "http.request") assert body_bytes == 600 + + +def _metadata_part_head() -> bytes: + return f'--{BOUNDARY}\r\nContent-Disposition: form-data; name="metadata"\r\n\r\n'.encode() + + +def _closing_boundary() -> bytes: + return f"\r\n--{BOUNDARY}--\r\n".encode() + + +@pytest.mark.parametrize("chunk_size", [1, 2, 3]) +def test_upload_metadata_buffer_is_bounded_by_its_size_not_its_chunk_count(chunk_size: int): + """A client that dribbles the metadata field in tiny pieces must not amplify what the + server retains. + + The parser hands the field over in whatever pieces the client sent it. Buffering one + `bytes` object per piece retained 8-22x the payload (per-object overhead dominates; + single-byte `bytes` are interned, so 2-byte pieces are the worst case), which turned + the 1 MiB metadata cap into a ~22 MiB one per upload. The buffer must be flat so the + cap bounds memory, not just payload. + """ + payload = b"x" * (64 * 1024) + with tempfile.TemporaryFile() as destination: + callbacks = videos._VideoUploadStreamParser(destination) + parser = MultipartParser(BOUNDARY.encode(), callbacks.callbacks) + parser.write(_metadata_part_head()) + + was_tracing = tracemalloc.is_tracing() + if not was_tracing: + tracemalloc.start() + try: + gc.collect() # so an unrelated collection mid-loop cannot skew the delta + before, _ = tracemalloc.get_traced_memory() + for start in range(0, len(payload), chunk_size): + parser.write(payload[start : start + chunk_size]) + after, _ = tracemalloc.get_traced_memory() + finally: + if not was_tracing: + tracemalloc.stop() + + retained = after - before + # bytearray over-allocates by at most ~12.5%; anything near 2x means per-chunk objects. + assert retained < 2 * len(payload), f"retained {retained / len(payload):.1f}x the payload" + + parser.write(_closing_boundary()) + parser.finalize() + assert callbacks.metadata == payload.decode() + + +def test_upload_metadata_cap_counts_the_whole_field_across_chunks(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(videos, "MAX_UPLOAD_METADATA_SIZE", 32) + + def parse(payload: bytes, chunk_size: int) -> videos._VideoUploadStreamParser: + with tempfile.TemporaryFile() as destination: + callbacks = videos._VideoUploadStreamParser(destination) + parser = MultipartParser(BOUNDARY.encode(), callbacks.callbacks) + parser.write(_metadata_part_head()) + for start in range(0, len(payload), chunk_size): + parser.write(payload[start : start + chunk_size]) + parser.write(_closing_boundary()) + parser.finalize() + return callbacks + + # Exactly at the cap, split unevenly across chunks: accepted and reassembled intact. + assert parse(b"a" * 32, chunk_size=5).metadata == "a" * 32 + + # One byte over, where no single chunk is anywhere near the cap: still rejected. + with pytest.raises(HTTPException) as error: + parse(b"a" * 33, chunk_size=5) + assert error.value.status_code == 413