From d417cb45985df43c5af9ba88a3c3be0776ed6e76 Mon Sep 17 00:00:00 2001 From: serversidehannes Date: Fri, 10 Jul 2026 15:55:45 +0200 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20stream=20sealed=20frames=20as=20upl?= =?UTF-8?q?oad=20bodies=20=E2=80=94=20O(part)=20->=20O(frame)=20memory?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Upload paths buffered a full internal part of ciphertext before handing it to the backend (up to ~240MB for multi-GB copies; the reason pods need 1Gi). FramedStreamBody now yields sealed 8MiB AES-GCM frames (in 1MB slices) as encryption produces them, with exact Content-Length via framed_ciphertext_size and a running ciphertext MD5 for backend ETag verification. - UploadPart framed path: multi-frame internal parts stream; single-frame parts stay buffered bytes so botocore can replay them. A mid-part backend failure surfaces as a failed client part (rclone/barman re-send it). - Copy pump: internal parts stream with pump-owned retry — on failure the plaintext source is reopened at the part's offset (ranged re-read) and the client-part MD5 is committed per attempt from a copy() snapshot, so retried bytes are never double-hashed. Short sources now fail loudly instead of writing wrong sizes. - Governor formulas recalibrated (tracemalloc): streaming_upload_peak flat 4*FRAME (32MB, real ~24MB) for multi-frame parts; copy_chunk_peak flat 4*FRAME + 2*MAX_BUFFER (48MB, real ~32MB). Clamp machinery is now vestigial. Closes #128 --- s3proxy/crypto.py | 59 +++-- s3proxy/handlers/multipart/copy.py | 172 +++++++++----- s3proxy/handlers/multipart/upload_part.py | 83 ++++--- s3proxy/streaming/__init__.py | 3 + s3proxy/streaming/framed_body.py | 114 +++++++++ tests/conftest.py | 19 ++ tests/unit/test_copy_internal_retry.py | 223 ++++++++++++++++++ tests/unit/test_copy_per_part_memory.py | 3 +- tests/unit/test_copy_reservation_vs_real.py | 27 ++- tests/unit/test_framed_stream_body.py | 107 +++++++++ tests/unit/test_governor_peak_and_clamp.py | 26 +- tests/unit/test_memory_concurrency.py | 9 +- tests/unit/test_streaming_framed_upload.py | 78 +++++- tests/unit/test_upload_reservation_vs_real.py | 38 +-- 14 files changed, 786 insertions(+), 175 deletions(-) create mode 100644 s3proxy/streaming/framed_body.py create mode 100644 tests/unit/test_copy_internal_retry.py create mode 100644 tests/unit/test_framed_stream_body.py diff --git a/s3proxy/crypto.py b/s3proxy/crypto.py index b58ca10..c402e99 100644 --- a/s3proxy/crypto.py +++ b/s3proxy/crypto.py @@ -231,23 +231,24 @@ def memory_bounded_part_size( def streaming_upload_peak(content_length: int) -> int: """Peak memory the framed UploadPart path holds for one in-flight request. - The path encrypts one internal part at a time. At peak it stacks the - accumulated ciphertext (~part), one plaintext frame being sealed, and (for - signed uploads) aiobotocore's copy of the part body for the HTTP request. - - Small single-frame parts (part <= FRAME_PLAINTEXT_SIZE): encrypt transient - dominates, bound by ``4*part`` (measured: 100KB -> 400KB). - - Multi-frame parts: ``2*part + FRAME_PLAINTEXT_SIZE`` (measured tracemalloc - with transport body copy: 16MB -> 24.5MB, 512MB/25.6MB internal -> 56.1MB, - 4GB/204MB internal -> 416.7MB). The old ``2*part + 2*frame`` formula - double-counted the frame term and scaled with multi-GB Content-Length even - though only one frame is in flight — starving concurrent ~50MB Scylla parts. + Multi-frame internal parts stream sealed frames to the backend + (FramedStreamBody): at peak the reader holds ~one frame of buffered source, + one plaintext frame is being sealed, and the sealed frame sits in the + transport write buffer — O(FRAME_PLAINTEXT_SIZE), independent of part size + and Content-Length. The backend hop is UNSIGNED-PAYLOAD, so botocore never + re-reads or copies the body. + + Small single-frame parts (part <= FRAME_PLAINTEXT_SIZE) are buffered whole + so botocore can replay them: encrypt transient dominates, bound by + ``4*part`` (measured: 100KB -> 400KB). + + The rare aws-chunked path still buffers whole internal parts under this + same reservation; the pod memory limit is the backstop there. """ part = memory_bounded_part_size(content_length) if part <= FRAME_PLAINTEXT_SIZE: return 4 * part - return 2 * part + FRAME_PLAINTEXT_SIZE + return 4 * FRAME_PLAINTEXT_SIZE def streaming_governor_clamped_reserve(honest_peak: int, budget_bytes: int) -> int: @@ -264,10 +265,9 @@ def streaming_governor_clamped_reserve(honest_peak: int, budget_bytes: int) -> i def copy_governor_clamped_reserve(honest_peak: int, budget_bytes: int) -> int: """Reservation when a copy's honest peak exceeds the governor budget. - Unlike uploads, a multi-GB copy's honest peak reflects real per-chunk work - (238MB internal parts for a 4.7GB Scylla manifest). Clamping to the routine - upload peak (~59MB) under-reserves and admits several concurrent copies that - each need ~500MB+ RSS. Monopolize the budget slot instead. + Nearly vestigial now that copy chunks stream (copy_chunk_peak is O(frame), + far below any sane budget); kept as the safety valve for tiny budgets. + Monopolize the budget slot rather than under-reserving. """ return min(honest_peak, budget_bytes) @@ -275,12 +275,9 @@ def copy_governor_clamped_reserve(honest_peak: int, budget_bytes: int) -> int: def governor_memory_footprint(content_length: int) -> int: """Memory to reserve for a framed upload at the request gate. - Uses the honest per-part peak for routine client part sizes (≤512MB) but - caps at the 512MB-workload peak for larger Content-Length values. Multi-GB - single PutObjects are rare in backup traffic (Scylla uses ~50MB multipart - parts); capping prevents a 6GB Content-Length from reserving 500MB+ and - starving concurrent parts. A lone multi-GB PutObject may exceed its - reservation — the pod memory limit is the backstop. + streaming_upload_peak is O(frame) for multi-frame parts, so the routine cap + only still matters for the buffered aws-chunked path; the pod memory limit + remains the backstop there. """ honest = streaming_upload_peak(content_length) routine_cap = streaming_upload_peak(STREAMING_GOVERNOR_CLIENT_PART_BYTES) @@ -302,21 +299,19 @@ def copy_internal_part_size(plaintext_size: int) -> int: def copy_chunk_peak(chunk_plaintext_bytes: int) -> int: - """Peak memory while encrypting one internal copy chunk (streaming path). + """Peak memory while streaming one internal copy chunk. - The streaming copy path acquires/releases this amount per internal part. - Reservation covers read-buffer slack plus framed encrypt peak and the - ciphertext buffer through S3 upload (released after del ciphertext). + The pump uploads each chunk as a FramedStreamBody, so the ciphertext is + never accumulated: the reservation covers the reader's buffered source + frame, the frame being sealed, the sealed frame in the transport buffer, + and read-buffer slack — O(frame), independent of chunk size. """ framed = ( 4 * chunk_plaintext_bytes if chunk_plaintext_bytes <= FRAME_PLAINTEXT_SIZE - else 2 * chunk_plaintext_bytes + FRAME_PLAINTEXT_SIZE + else 4 * FRAME_PLAINTEXT_SIZE ) - peak = framed + 2 * MAX_BUFFER_SIZE - if chunk_plaintext_bytes > 32 * 1024 * 1024: - peak += chunk_plaintext_bytes // 7 - return peak + return framed + 2 * MAX_BUFFER_SIZE # UploadPartCopy passthrough moves bytes server-side; in-process peak is tiny. diff --git a/s3proxy/handlers/multipart/copy.py b/s3proxy/handlers/multipart/copy.py index c171e69..f270f9f 100644 --- a/s3proxy/handlers/multipart/copy.py +++ b/s3proxy/handlers/multipart/copy.py @@ -7,7 +7,7 @@ import math import os import time -from collections.abc import AsyncIterator, Coroutine +from collections.abc import AsyncIterator, Callable, Coroutine from dataclasses import dataclass from datetime import UTC, datetime from urllib.parse import quote @@ -30,6 +30,7 @@ load_upload_state, persist_upload_state, ) +from ...streaming import FramedStreamBody from ...utils import format_iso8601 from ..base import BaseHandler from .upload_part import _PlaintextReader @@ -59,6 +60,11 @@ # part 1 is large enough that a client part 2 will follow (avoids EntityTooSmall). HYBRID_TAIL_DEFER_MIN_CLIENT_PART = 1024 * 1024 * 1024 # 1 GiB +# Streamed internal-part bodies cannot be replayed by botocore's retry layer, +# so the copy pump retries them itself by reopening the source at the failed +# part's offset. +COPY_INTERNAL_PART_ATTEMPTS = 3 + def reset_copy_pipeline_semaphore(limit: int | None = None) -> None: """Reset the global copy pipeline semaphore (testing only).""" @@ -1215,16 +1221,18 @@ async def _streaming_copy_part_inner( internal_part_start=internal_part_start, ) - src_iter = self._iter_copy_source( - client, - src_bucket, - src_key, - copy_source_range, - src_wrapped_dek, - src_multipart_meta, - head_resp, - src_metadata, - ) + def open_source(skip_bytes: int) -> AsyncIterator[bytes]: + return self._iter_copy_source( + client, + src_bucket, + src_key, + copy_source_range, + src_wrapped_dek, + src_multipart_meta, + head_resp, + src_metadata, + skip_bytes=skip_bytes, + ) internal_parts, total_plaintext, total_ciphertext, md5 = await self._pump_copy_chunks( client, @@ -1233,10 +1241,12 @@ async def _streaming_copy_part_inner( upload_id, part_num, state, - src_iter, + open_source(0), chunk_size, internal_part_start, leading_plaintext=deferred_tail or b"", + expected_plaintext=plaintext_size + len(deferred_tail), + reopen_source=open_source, ) etag = md5.hexdigest() @@ -1282,14 +1292,26 @@ async def _iter_copy_source( src_multipart_meta, head_resp: dict, src_metadata: dict, + skip_bytes: int = 0, ) -> AsyncIterator[bytes]: - """Yield raw plaintext bytes from the copy source, one part/chunk at a time.""" + """Yield raw plaintext bytes from the copy source, one part/chunk at a time. + + skip_bytes skips that many plaintext bytes from the start of the + (possibly ranged) source — used to rebuild the stream at an internal + part boundary when the copy pump retries a failed part. + """ if src_multipart_meta: total = src_multipart_meta.total_plaintext_size if copy_source_range: range_start, range_end = self._parse_copy_source_range(copy_source_range, total) + elif skip_bytes: + range_start, range_end = 0, total - 1 else: range_start, range_end = None, None + if skip_bytes: + range_start += skip_bytes + if range_start > range_end: + return dek = crypto.unwrap_key( src_multipart_meta.wrapped_dek, self.keyring.key_by_id(src_multipart_meta.kid), @@ -1306,9 +1328,20 @@ async def _iter_copy_source( if copy_source_range: start, end = self._parse_copy_source_range(copy_source_range, len(plaintext)) plaintext = plaintext[start : end + 1] - yield plaintext + if skip_bytes: + plaintext = plaintext[skip_bytes:] + if plaintext: + yield plaintext + return else: - resp = await client.get_object(src_bucket, src_key, range_header=copy_source_range) + range_header = copy_source_range + if skip_bytes: + if copy_source_range: + raw_start, raw_end = self._parse_raw_copy_source_range(copy_source_range) + range_header = f"bytes={raw_start + skip_bytes}-{raw_end}" + else: + range_header = f"bytes={skip_bytes}-" + resp = await client.get_object(src_bucket, src_key, range_header=range_header) async with resp["Body"] as body: # resp["Body"] enters as an aiohttp ClientResponse, whose read() # takes no size arg; stream via its StreamReader in bounded chunks @@ -1332,65 +1365,76 @@ async def _pump_copy_chunks( internal_part_start: int, *, leading_plaintext: bytes = b"", + expected_plaintext: int, + reopen_source: Callable[[int], AsyncIterator[bytes]] | None = None, ) -> tuple[list[InternalPartMetadata], int, int, object]: """Frame-encrypt the copy source into internal S3 parts, one at a time. - Mirrors UploadPartMixin._stream_and_upload_framed: reads FRAME_PLAINTEXT_SIZE - plaintext frames from the source, encrypts each with encrypt_frame, - accumulates one internal part's ciphertext, then uploads it before starting - the next. Memory governor reservation is per internal part (acquire before - encrypt, hold through upload, release after ciphertext is freed) so RSS - matches what the limiter tracks and a multi-GB copy does not hold one - reservation for its entire duration. + Each internal part is uploaded as a FramedStreamBody that yields sealed + frames while the source is being read, so peak memory is O(frame) + instead of O(chunk_size). The part sizes are derived up front from + expected_plaintext (source metadata is authoritative); a source that + ends early fails loudly instead of writing an object with wrong sizes. + + A streamed body cannot be replayed by botocore, so failed internal + parts are retried here: reopen_source(offset) rebuilds the plaintext + stream at the failed part's offset (relative to the source, after + leading_plaintext). The client-part MD5 is committed per attempt from a + copy() snapshot so retried bytes are never hashed twice. """ reader = _PlaintextReader(src_iter, prefix=leading_plaintext) md5 = hashlib.md5(usedforsecurity=False) internal_parts: list[InternalPartMetadata] = [] - total_plaintext = 0 total_ciphertext = 0 - internal_part_num = internal_part_start + offset = 0 + index = 0 - prefetch: bytes | None = None - while True: - if prefetch is None: - prefetch = await reader.read(min(crypto.FRAME_PLAINTEXT_SIZE, chunk_size)) - if not prefetch: + while offset < expected_plaintext: + part_plaintext = min(chunk_size, expected_plaintext - offset) + ipn = internal_part_start + index + ct_size = crypto.framed_ciphertext_size(part_plaintext) + part_reserve = crypto.copy_chunk_peak(part_plaintext) + + for attempt in range(1, COPY_INTERNAL_PART_ATTEMPTS + 1): + md5_attempt = md5.copy() + body = FramedStreamBody( + reader, + part_plaintext, + state.dek, + upload_id, + ipn, + plaintext_hashes=(md5_attempt,), + ) + upload_start = time.monotonic() + try: + async with concurrency.reserve_copy_memory(part_reserve): + resp = await client.upload_part(bucket, key, upload_id, ipn, body) + md5 = md5_attempt break - - part_reserve = crypto.copy_chunk_peak(chunk_size) - ciphertext = bytearray() - part_plaintext = 0 - frame_idx = 0 - async with concurrency.reserve_copy_memory(part_reserve): - frame_pt = prefetch - prefetch = None - while True: - md5.update(frame_pt) - part_plaintext += len(frame_pt) - ciphertext.extend( - crypto.encrypt_frame( - frame_pt, state.dek, upload_id, internal_part_num, frame_idx - ) + except Exception as e: + if attempt == COPY_INTERNAL_PART_ATTEMPTS or reopen_source is None: + raise + logger.warning( + "COPY_INTERNAL_PART_RETRY", + bucket=bucket, + key=key, + client_part=part_num, + internal_part=ipn, + attempt=attempt, + plaintext_offset=offset, + error_type=type(e).__name__, + error=str(e), ) - frame_idx += 1 - if part_plaintext >= chunk_size: - break - frame_pt = await reader.read( - min(crypto.FRAME_PLAINTEXT_SIZE, chunk_size - part_plaintext) - ) - if not frame_pt: - break - - upload_start = time.monotonic() - resp = await client.upload_part( - bucket, key, upload_id, internal_part_num, ciphertext - ) - del ciphertext + if offset < len(leading_plaintext): + reader = _PlaintextReader( + reopen_source(0), prefix=leading_plaintext[offset:] + ) + else: + reader = _PlaintextReader(reopen_source(offset - len(leading_plaintext))) - ct_size = crypto.framed_ciphertext_size(part_plaintext) internal_parts.append( InternalPartMetadata( - internal_part_number=internal_part_num, + internal_part_number=ipn, plaintext_size=part_plaintext, ciphertext_size=ct_size, etag=resp["ETag"].strip('"'), @@ -1401,12 +1445,12 @@ async def _pump_copy_chunks( bucket=bucket, key=key, client_part=part_num, - internal_part=internal_part_num, + internal_part=ipn, plaintext_mb=f"{part_plaintext / 1024 / 1024:.2f}MB", elapsed_sec=f"{time.monotonic() - upload_start:.2f}s", ) - total_plaintext += part_plaintext + offset += part_plaintext total_ciphertext += ct_size - internal_part_num += 1 + index += 1 - return internal_parts, total_plaintext, total_ciphertext, md5 + return internal_parts, offset, total_ciphertext, md5 diff --git a/s3proxy/handlers/multipart/upload_part.py b/s3proxy/handlers/multipart/upload_part.py index cfc3048..6447528 100644 --- a/s3proxy/handlers/multipart/upload_part.py +++ b/s3proxy/handlers/multipart/upload_part.py @@ -24,7 +24,7 @@ PartMetadata, StateMissingError, ) -from ...streaming import decode_aws_chunked_stream +from ...streaming import FramedStreamBody, decode_aws_chunked_stream from ..base import BaseHandler logger: BoundLogger = structlog.get_logger(__name__) @@ -411,10 +411,13 @@ async def _stream_and_upload_framed( Used for unsigned and large signed uploads (e.g. barman backups) where Content-Length is known and the body is read directly (not aws-chunked). Reads the body once, splitting it into internal S3 parts of - optimal_part_size. Each internal part is encrypted frame-by-frame - (see crypto.encrypt_frame) into a ciphertext buffer, then uploaded as - bytes. Peak memory is O(part ciphertext + FRAME_PLAINTEXT_SIZE) — one - frame of plaintext at a time, not plaintext + ciphertext together. + optimal_part_size. Multi-frame internal parts stream sealed frames to + the backend as encryption produces them (FramedStreamBody), so peak + memory is O(frame) regardless of part size. Single-frame parts are + buffered as bytes so botocore can replay them on its internal retries. + + A streamed body cannot be replayed: a mid-part backend failure fails + the whole client part and the client (rclone/barman) re-sends it. """ md5_hash = hashlib.md5(usedforsecurity=False) sha256_hash = hashlib.sha256() @@ -424,46 +427,52 @@ async def _stream_and_upload_framed( internal_part_num = internal_part_start remaining_total = content_length + def raise_stream_short(ipn: int, expected: int, received: int) -> NoReturn: + logger.error( + "UPLOAD_PART_STREAM_SHORT", + bucket=bucket, + key=key, + client_part=part_num, + internal_part=ipn, + expected_bytes=expected, + received_bytes=received, + content_length=content_length, + ) + raise S3Error.bad_request( + f"Upload body ended early: got {received}B of {expected}B " + f"for client part {part_num} internal part {ipn}" + ) + while remaining_total > 0: part_pt_size = min(optimal_part_size, remaining_total) ct_size = crypto.framed_ciphertext_size(part_pt_size) ipn = internal_part_num - ciphertext = bytearray() - remaining = part_pt_size - frame_idx = 0 - while remaining > 0: - frame_pt = await reader.read(min(crypto.FRAME_PLAINTEXT_SIZE, remaining)) - if not frame_pt: - break + upload_start = time.monotonic() + if part_pt_size <= crypto.FRAME_PLAINTEXT_SIZE: + frame_pt = await reader.read(part_pt_size) + if len(frame_pt) < part_pt_size: + raise_stream_short(ipn, part_pt_size, len(frame_pt)) md5_hash.update(frame_pt) sha256_hash.update(frame_pt) - remaining -= len(frame_pt) - ciphertext.extend( - crypto.encrypt_frame(frame_pt, state.dek, upload_id, ipn, frame_idx) - ) - frame_idx += 1 - - if remaining > 0: - bytes_read = part_pt_size - remaining - logger.error( - "UPLOAD_PART_STREAM_SHORT", - bucket=bucket, - key=key, - client_part=part_num, - internal_part=ipn, - expected_bytes=part_pt_size, - received_bytes=bytes_read, - content_length=content_length, - ) - raise S3Error.bad_request( - f"Upload body ended early: got {bytes_read}B of {part_pt_size}B " - f"for client part {part_num} internal part {ipn}" + ciphertext = crypto.encrypt_frame(frame_pt, state.dek, upload_id, ipn, 0) + resp = await client.upload_part(bucket, key, upload_id, ipn, ciphertext) + del ciphertext + else: + body = FramedStreamBody( + reader, + part_pt_size, + state.dek, + upload_id, + ipn, + plaintext_hashes=(md5_hash, sha256_hash), ) - - upload_start = time.monotonic() - resp = await client.upload_part(bucket, key, upload_id, ipn, ciphertext) - del ciphertext + try: + resp = await client.upload_part(bucket, key, upload_id, ipn, body) + except Exception: + if body.short_read: + raise_stream_short(ipn, part_pt_size, body.bytes_read) + raise etag = resp["ETag"].strip('"') logger.info( "INTERNAL_PART_UPLOADED", diff --git a/s3proxy/streaming/__init__.py b/s3proxy/streaming/__init__.py index ae9ebe5..3785b0a 100644 --- a/s3proxy/streaming/__init__.py +++ b/s3proxy/streaming/__init__.py @@ -9,9 +9,12 @@ decode_aws_chunked, decode_aws_chunked_stream, ) +from .framed_body import FramedStreamBody, SourceExhaustedError __all__ = [ "STREAM_CHUNK_SIZE", + "FramedStreamBody", + "SourceExhaustedError", "chunked", "decode_aws_chunked", "decode_aws_chunked_stream", diff --git a/s3proxy/streaming/framed_body.py b/s3proxy/streaming/framed_body.py new file mode 100644 index 0000000..ea46423 --- /dev/null +++ b/s3proxy/streaming/framed_body.py @@ -0,0 +1,114 @@ +"""Streaming upload body that seals AES-GCM frames on the fly. + +Upload paths used to accumulate a whole internal part of ciphertext before +handing it to the backend client, so peak memory tracked the internal part +size (hundreds of MB for multi-GB client parts). FramedStreamBody instead +yields each sealed frame as encryption produces it: peak memory is O(frame), +independent of part size. + +The body reports the exact framed ciphertext size via __len__, so botocore +sends a normal Content-Length request (aiohttp only falls back to chunked +transfer encoding when no length is known). It tracks a running MD5 of the +ciphertext for backend ETag verification (see client.s3.verify_backend_etag). + +Single-shot: botocore cannot replay it on its internal retries (no seek). +Callers own retry semantics — the copy pump rebuilds the source stream and +retries the internal part itself; the upload path surfaces the failure so the +client re-sends its part. +""" + +from __future__ import annotations + +import hashlib +from typing import TYPE_CHECKING, Any + +from .. import crypto + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + +class SourceExhaustedError(Exception): + """The plaintext source ended before yielding the promised bytes.""" + + +# Sealed frames are handed to the transport in slices this big, so the consumer +# never pins a whole 8MB frame while the next one is being produced (that held +# frame was a full quarter of the measured peak). +FRAME_YIELD_SLICE = 1024 * 1024 + + +class FramedStreamBody: + """Async-iterable request body of sealed AES-GCM frames for one internal part.""" + + def __init__( + self, + reader: Any, + plaintext_size: int, + dek: bytes, + upload_id: str, + part_number: int, + *, + plaintext_hashes: tuple[Any, ...] = (), + ) -> None: + if plaintext_size <= 0: + raise ValueError("plaintext_size must be positive") + self._reader = reader + self._plaintext_size = plaintext_size + self._dek = dek + self._upload_id = upload_id + self._part_number = part_number + self._plaintext_hashes = plaintext_hashes + self._ciphertext_md5 = hashlib.md5(usedforsecurity=False) + self._consumed = False + self.bytes_read = 0 + self.short_read = False + + def __len__(self) -> int: + return crypto.framed_ciphertext_size(self._plaintext_size) + + def read(self, *args: Any, **kwargs: Any) -> bytes: + # Present only so botocore's blob param validation accepts the body + # (bytes-like or has `read`). Nothing may consume the body this way: + # payload signing and flexible checksums are disabled on the backend + # client, and aiohttp streams via __aiter__. + raise NotImplementedError("FramedStreamBody is consumed via async iteration") + + def ciphertext_md5_hexdigest(self) -> str: + return self._ciphertext_md5.hexdigest() + + def __aiter__(self) -> AsyncIterator[bytes]: + if self._consumed: + raise RuntimeError("FramedStreamBody is single-shot and was already consumed") + self._consumed = True + return self._frames() + + async def _frames(self) -> AsyncIterator[bytes]: + remaining = self._plaintext_size + frame_index = 0 + while remaining > 0: + want = min(crypto.FRAME_PLAINTEXT_SIZE, remaining) + frame_pt = await self._reader.read(want) + self.bytes_read += len(frame_pt) + if len(frame_pt) < want: + self.short_read = True + raise SourceExhaustedError( + f"source ended after {self.bytes_read} of {self._plaintext_size} " + f"plaintext bytes for internal part {self._part_number}" + ) + for h in self._plaintext_hashes: + h.update(frame_pt) + remaining -= want + frame = crypto.encrypt_frame( + frame_pt, self._dek, self._upload_id, self._part_number, frame_index + ) + del frame_pt + self._ciphertext_md5.update(frame) + frame_index += 1 + view = memoryview(frame) + for offset in range(0, len(frame), FRAME_YIELD_SLICE): + yield bytes(view[offset : offset + FRAME_YIELD_SLICE]) + # Drop the frame before the next read so the generator never holds + # two frames at once (keeps the peak at ~2 frames, not 4). + view.release() + del frame diff --git a/tests/conftest.py b/tests/conftest.py index 759248e..cf8ce7e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -104,6 +104,23 @@ def credentials_store(): # ============================================================================ +async def materialize_body(body) -> bytes: + """Collect an upload body into bytes, like the real backend transport. + + Streaming bodies (FramedStreamBody) declare their exact size via len(); + enforce that the yielded bytes match it, since the real request is sent + with that Content-Length. + """ + if isinstance(body, (bytes, bytearray)): + return bytes(body) + declared = len(body) + data = b"".join([chunk async for chunk in body]) + assert len(data) == declared, ( + f"streaming body yielded {len(data)} bytes but declared Content-Length {declared}" + ) + return data + + class MockS3Response: """Mock S3 response object.""" @@ -173,6 +190,7 @@ async def put_object( ) -> dict: """Store an object.""" self.call_history.append(("put_object", {"bucket": bucket, "key": key})) + body = await materialize_body(body) self.objects[self._key(bucket, key)] = { "Body": body, "Metadata": metadata or {}, @@ -411,6 +429,7 @@ async def upload_part( if upload_id not in self.multipart_uploads: raise self._not_found_error(f"upload {upload_id}") + body = await materialize_body(body) etag = hashlib.md5(body).hexdigest() self.multipart_uploads[upload_id]["Parts"][part_number] = { "Body": body, diff --git a/tests/unit/test_copy_internal_retry.py b/tests/unit/test_copy_internal_retry.py new file mode 100644 index 0000000..44ef0cd --- /dev/null +++ b/tests/unit/test_copy_internal_retry.py @@ -0,0 +1,223 @@ +"""Copy pump retry semantics for streamed internal-part bodies. + +Streamed bodies cannot be replayed by botocore, so _pump_copy_chunks retries a +failed internal part itself: it reopens the plaintext source at the part's +offset and re-streams it. These tests inject mid-body backend failures and +assert the copy self-heals with correct data, correct client ETag (the MD5 +snapshot must not double-hash retried bytes), and correct reopen offsets. +""" + +from __future__ import annotations + +import asyncio +import hashlib + +import pytest + +from s3proxy import concurrency, crypto +from s3proxy.handlers.multipart import MultipartHandlerMixin +from s3proxy.state import MultipartUploadState + +MB = 1024 * 1024 + + +class _Body: + def __init__(self, data: bytes, chunk: int = MB): + self._data = data + self._chunk = chunk + self._pos = 0 + self.content = self + + async def read(self, n: int) -> bytes: + take = min(n, self._chunk, len(self._data) - self._pos) + out = self._data[self._pos : self._pos + take] + self._pos += take + return out + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + +class _FlakyClient: + """Unencrypted source + upload_part that drops the connection mid-body.""" + + def __init__(self, data: bytes, fail_parts: dict[int, int]): + self.data = data + self.fail = dict(fail_parts) + self.uploaded: dict[int, bytes] = {} + self.get_ranges: list[str | None] = [] + + async def get_object(self, bucket, key, range_header=None): + self.get_ranges.append(range_header) + body = self.data + if range_header: + spec = range_header.replace("bytes=", "") + start_s, end_s = spec.split("-") + start = int(start_s) + end = int(end_s) if end_s else len(body) - 1 + body = body[start : end + 1] + return {"Body": _Body(body)} + + async def upload_part(self, bucket, key, upload_id, part_number, body): + collected = bytearray() + if isinstance(body, (bytes, bytearray)): + collected += body + else: + declared = len(body) + async for chunk in body: + collected += chunk + if self.fail.get(part_number, 0) and len(collected) > declared // 2: + break + if self.fail.get(part_number, 0): + self.fail[part_number] -= 1 + raise ConnectionError("backend reset mid-part") + self.uploaded[part_number] = bytes(collected) + return {"ETag": f'"{hashlib.md5(bytes(collected), usedforsecurity=False).hexdigest()}"'} + + +class _Mgr: + def __init__(self, deferred_tail: bytes = b""): + self.deferred_tail = deferred_tail + self.part_meta = None + self.allocated = 0 + + async def allocate_internal_parts(self, bucket, key, upload_id, count, client_part_number): + start = self.allocated + 1 + self.allocated += count + return start + + async def add_part(self, bucket, key, upload_id, part_meta): + self.part_meta = part_meta + + async def take_deferred_copy_tail(self, bucket, key, upload_id): + tail, self.deferred_tail = self.deferred_tail, b"" + return tail + + +def _handler(mgr) -> MultipartHandlerMixin: + h = MultipartHandlerMixin.__new__(MultipartHandlerMixin) + h.multipart_manager = mgr + return h + + +async def _run_copy(client, mgr, plaintext_size: int) -> tuple[MultipartUploadState, object]: + state = MultipartUploadState(dek=crypto.generate_dek(), bucket="b", key="k", upload_id="u") + resp = await _handler(mgr)._streaming_copy_part_inner( + client, + "b", + "k", + "u", + 1, + state, + "b", + "src", + None, + None, + None, + {}, + {}, + plaintext_size, + ) + return state, resp + + +def _recovered_plaintext(client: _FlakyClient, mgr: _Mgr, dek: bytes) -> bytes: + parts = sorted(mgr.part_meta.internal_parts, key=lambda m: m.internal_part_number) + return b"".join( + crypto.decrypt_framed(client.uploaded[m.internal_part_number], dek, m.plaintext_size) + for m in parts + ) + + +@pytest.fixture(autouse=True) +def _fresh_governor(): + concurrency.reset_state() + concurrency.set_memory_limit(192) + yield + concurrency.reset_state() + + +@pytest.mark.asyncio +async def test_mid_part_failure_retries_and_roundtrips(): + data = bytes(80 * MB) # 32 + 32 + 16 MB internal parts + client = _FlakyClient(data, fail_parts={2: 1}) + mgr = _Mgr() + + state, _ = await _run_copy(client, mgr, len(data)) + + assert _recovered_plaintext(client, mgr, state.dek) == data + # Client ETag is the plaintext MD5; a double-hashed retry would corrupt it. + assert mgr.part_meta.etag == hashlib.md5(data, usedforsecurity=False).hexdigest() + assert mgr.part_meta.plaintext_size == len(data) + # The retry reopened the source at the failed part's offset (32MB). + assert f"bytes={32 * MB}-" in client.get_ranges + + +@pytest.mark.asyncio +async def test_retry_exhaustion_propagates_failure(): + data = bytes(40 * MB) + client = _FlakyClient(data, fail_parts={1: 10}) + mgr = _Mgr() + + with pytest.raises(ConnectionError): + await _run_copy(client, mgr, len(data)) + assert mgr.part_meta is None + assert concurrency.get_active_memory() == 0 + + +@pytest.mark.asyncio +async def test_retry_with_deferred_tail_prefix(): + """A failure in the part containing the deferred tail must rebuild the + stream as tail-remainder + reopened source.""" + tail = bytes([7]) * (2 * MB) + data = bytes([9]) * (40 * MB) + client = _FlakyClient(data, fail_parts={1: 1}) + mgr = _Mgr(deferred_tail=tail) + + state, _ = await _run_copy(client, mgr, len(data)) + + assert _recovered_plaintext(client, mgr, state.dek) == tail + data + expected_md5 = hashlib.md5(tail + data, usedforsecurity=False).hexdigest() + assert mgr.part_meta.etag == expected_md5 + assert mgr.part_meta.plaintext_size == len(tail) + len(data) + + +@pytest.mark.asyncio +async def test_short_source_fails_loudly_not_silently_truncated(): + """Source metadata promised more bytes than the stream delivers: the copy + must error (client retries) instead of storing an object with wrong sizes.""" + data = bytes(20 * MB) + client = _FlakyClient(data, fail_parts={}) + mgr = _Mgr() + + with pytest.raises(Exception, match="source ended"): + await _run_copy(client, mgr, 40 * MB) + assert mgr.part_meta is None + + +@pytest.mark.asyncio +async def test_no_failure_means_single_pass_no_reopen(): + data = bytes(64 * MB) + client = _FlakyClient(data, fail_parts={}) + mgr = _Mgr() + + state, _ = await _run_copy(client, mgr, len(data)) + assert _recovered_plaintext(client, mgr, state.dek) == data + assert client.get_ranges == [None] + + +@pytest.mark.asyncio +async def test_concurrent_retrying_copies_release_reservations(): + """Governor must return to zero even when copies fail and retry.""" + data = bytes(40 * MB) + + async def one(fail: dict[int, int]): + client = _FlakyClient(data, fail_parts=fail) + mgr = _Mgr() + await _run_copy(client, mgr, len(data)) + + await asyncio.gather(one({1: 1}), one({2: 1}), one({})) + assert concurrency.get_active_memory() == 0 diff --git a/tests/unit/test_copy_per_part_memory.py b/tests/unit/test_copy_per_part_memory.py index 358afe4..d2f4f4c 100644 --- a/tests/unit/test_copy_per_part_memory.py +++ b/tests/unit/test_copy_per_part_memory.py @@ -298,7 +298,8 @@ async def tracking_acquire(bytes_needed: int, *, copy: bool = False): def test_copy_chunk_peak_32mb_uses_framed_formula_not_small_buffer(): """32MB internal parts must not hit the small-buffer 3x formula (96MB bug).""" peak = crypto.copy_chunk_peak(32 * MB) - assert peak == 88 * MB, f"32MB chunk peak should be 88MB, got {peak / MB}MB" + expected = 4 * crypto.FRAME_PLAINTEXT_SIZE + 2 * crypto.MAX_BUFFER_SIZE + assert peak == expected, f"32MB chunk peak should be {expected / MB}MB, got {peak / MB}MB" assert peak < 96 * MB diff --git a/tests/unit/test_copy_reservation_vs_real.py b/tests/unit/test_copy_reservation_vs_real.py index d57a495..bc35c49 100644 --- a/tests/unit/test_copy_reservation_vs_real.py +++ b/tests/unit/test_copy_reservation_vs_real.py @@ -1,17 +1,16 @@ """The copy reservation must bound the streaming UploadPartCopy path's REAL peak. Analogue of test_upload_reservation_vs_real.py, for the copy path. This is the -test that would have caught the ScyllaDB backup OOM: copy_pipeline_peak() -reserved a size-independent 32MB (4*MAX_BUFFER_SIZE) on the false premise that -the copy "streams in MAX_BUFFER_SIZE chunks", while _pump_copy_chunks actually -buffers a full internal part and _upload_internal_part_with_semaphore holds -plaintext + ciphertext of that part at once, times MAX_PARALLEL_INTERNAL_UPLOADS. -The governor therefore admitted ~6x too many concurrent copies -> OOMKilled. +test that would have caught the ScyllaDB backup OOM: an earlier +copy_pipeline_peak() reserved a size-independent value on a false premise about +what the pump buffered, so the governor admitted ~6x too many concurrent +copies -> OOMKilled. The pump now streams each internal part as sealed frames +(FramedStreamBody), so the honest peak really is O(frame); this test keeps the +reservation honest against the real allocation profile. It drives the ACTUAL handler method (not a re-implementation) under tracemalloc, so it fails if the reservation drifts below reality OR the copy path starts -allocating more. The mock client copies the uploaded body the way aiobotocore -does for the signed HTTP request, which is part of the real peak. +allocating more. """ import hashlib @@ -71,10 +70,16 @@ async def get_object(self, bucket, key, range_header=None): return {"Body": _Body(self._total, crypto.MAX_BUFFER_SIZE)} async def upload_part(self, bucket, key, upload_id, part_number, body): - # aiobotocore copies the body to sign and send it; mirror that so the + # The backend hop is UNSIGNED-PAYLOAD with streaming bodies: the + # transport consumes frames without retaining them. Mirror that so the # measured peak reflects what the real transport holds. - sent = bytes(body) - return {"ETag": hashlib.md5(sent).hexdigest()} + md5 = hashlib.md5(usedforsecurity=False) + if isinstance(body, (bytes, bytearray)): + md5.update(body) + else: + async for chunk in body: + md5.update(chunk) + return {"ETag": md5.hexdigest()} def _handler(): diff --git a/tests/unit/test_framed_stream_body.py b/tests/unit/test_framed_stream_body.py new file mode 100644 index 0000000..4d13e44 --- /dev/null +++ b/tests/unit/test_framed_stream_body.py @@ -0,0 +1,107 @@ +"""FramedStreamBody: sealed frames on the fly, exact Content-Length, no replay.""" + +import hashlib +import os + +import pytest + +from s3proxy import crypto +from s3proxy.handlers.multipart.upload_part import _PlaintextReader +from s3proxy.streaming import FramedStreamBody, SourceExhaustedError + +UPLOAD_ID = "u" * 40 +MB = 1024 * 1024 + + +async def _chunks(data: bytes, chunk: int = 1 << 20): + for i in range(0, len(data), chunk): + yield data[i : i + chunk] + + +def _body(data: bytes, size: int | None = None, **kwargs) -> FramedStreamBody: + return FramedStreamBody( + _PlaintextReader(_chunks(data)), + size if size is not None else len(data), + b"k" * 32, + UPLOAD_ID, + 7, + **kwargs, + ) + + +async def _collect(body: FramedStreamBody) -> bytes: + return b"".join([frame async for frame in body]) + + +@pytest.mark.asyncio +async def test_yields_decryptable_frames_with_exact_declared_length(): + dek = crypto.generate_dek() + plaintext = os.urandom(20 * MB) # 3 frames: 8, 8, 4 + body = FramedStreamBody(_PlaintextReader(_chunks(plaintext)), len(plaintext), dek, UPLOAD_ID, 7) + + ciphertext = await _collect(body) + assert len(ciphertext) == len(body) == crypto.framed_ciphertext_size(len(plaintext)) + assert crypto.decrypt_framed(ciphertext, dek, len(plaintext)) == plaintext + + +@pytest.mark.asyncio +async def test_frame_boundaries_and_nonces_match_buffered_encryption(): + """Streamed frames must be byte-identical to the old buffered encrypt loop.""" + dek = crypto.generate_dek() + plaintext = os.urandom(10 * MB) + body = FramedStreamBody(_PlaintextReader(_chunks(plaintext)), len(plaintext), dek, UPLOAD_ID, 7) + streamed = await _collect(body) + + buffered = bytearray() + for idx in range(crypto.frame_count(len(plaintext))): + start = idx * crypto.FRAME_PLAINTEXT_SIZE + frame_pt = plaintext[start : start + crypto.FRAME_PLAINTEXT_SIZE] + buffered.extend(crypto.encrypt_frame(frame_pt, dek, UPLOAD_ID, 7, idx)) + assert streamed == bytes(buffered) + + +@pytest.mark.asyncio +async def test_tracks_plaintext_hashes_and_ciphertext_md5(): + plaintext = os.urandom(9 * MB) + md5 = hashlib.md5(usedforsecurity=False) + sha = hashlib.sha256() + body = _body(plaintext, plaintext_hashes=(md5, sha)) + + ciphertext = await _collect(body) + assert md5.hexdigest() == hashlib.md5(plaintext, usedforsecurity=False).hexdigest() + assert sha.hexdigest() == hashlib.sha256(plaintext).hexdigest() + assert body.ciphertext_md5_hexdigest() == hashlib.md5(ciphertext).hexdigest() + + +@pytest.mark.asyncio +async def test_short_source_raises_and_flags(): + plaintext = os.urandom(3 * MB) + body = _body(plaintext, size=9 * MB) + + with pytest.raises(SourceExhaustedError): + await _collect(body) + assert body.short_read + assert body.bytes_read == 3 * MB + + +@pytest.mark.asyncio +async def test_single_shot_cannot_be_replayed(): + body = _body(b"x" * MB) + await _collect(body) + with pytest.raises(RuntimeError, match="single-shot"): + await _collect(body) + + +def test_botocore_compat_surface(): + """Blob param validation needs `read`; length must be known via len().""" + body = _body(b"x" * MB) + assert hasattr(body, "read") + with pytest.raises(NotImplementedError): + body.read() + assert len(body) == crypto.framed_ciphertext_size(MB) + assert not hasattr(body, "seek") # botocore must not silently replay it + + +def test_rejects_empty_plaintext(): + with pytest.raises(ValueError): + _body(b"", size=0) diff --git a/tests/unit/test_governor_peak_and_clamp.py b/tests/unit/test_governor_peak_and_clamp.py index c650743..deb7ede 100644 --- a/tests/unit/test_governor_peak_and_clamp.py +++ b/tests/unit/test_governor_peak_and_clamp.py @@ -56,6 +56,11 @@ async def add_part(self, *a, **k): class _DiscardingClient: async def upload_part(self, bucket, key, upload_id, part_number, body): + # Consume streaming bodies chunk-by-chunk like the real transport, + # retaining nothing, so measured peaks reflect prod behavior. + if not isinstance(body, (bytes, bytearray)): + async for _chunk in body: + pass del body return {"ETag": '"0"'} @@ -108,13 +113,10 @@ async def _measure_framed_peak(content_length: int) -> int: [16, 50, 512, 1024, 4096], ) def test_streaming_upload_peak_formula(mb: int): - """Peak formula must match the two measured regimes (small 4*part, large 2*part+frame).""" + """Small parts buffer (4*part); multi-frame parts stream at a flat O(frame) peak.""" cl = mb * MB part = crypto.memory_bounded_part_size(cl) - if part <= crypto.FRAME_PLAINTEXT_SIZE: - expected = 4 * part - else: - expected = 2 * part + crypto.FRAME_PLAINTEXT_SIZE + expected = 4 * part if part <= crypto.FRAME_PLAINTEXT_SIZE else 4 * crypto.FRAME_PLAINTEXT_SIZE assert crypto.streaming_upload_peak(cl) == expected @@ -148,7 +150,12 @@ def test_multi_gb_estimate_capped_at_routine_workload_peak(): @pytest.mark.asyncio async def test_clamp_does_not_monopolize_budget(): - """Concurrent ~50MB part must not stall behind one clamped multi-GB estimate.""" + """Concurrent ~50MB part must not stall behind one multi-GB upload. + + With frame-streaming bodies the honest peak is O(frame) for any + Content-Length, so the clamp scenario can no longer arise at all: a + multi-GB estimate reserves the same as a routine part and coexists. + """ reset_state() set_memory_limit(312) try: @@ -158,13 +165,12 @@ async def test_clamp_does_not_monopolize_budget(): r_scylla = await try_acquire_memory(scylla) assert r_scylla == scylla - # Honest peak for multi-GB body exceeds 312MB budget before governor cap. huge_honest = crypto.streaming_upload_peak(5 * 1024 * MB) - assert huge_honest > 312 * MB + routine = crypto.streaming_upload_peak(crypto.STREAMING_GOVERNOR_CLIENT_PART_BYTES) + assert huge_honest == routine huge_reserved = crypto.governor_memory_footprint(5 * 1024 * MB) r_huge = await try_acquire_memory(huge_reserved) - routine = crypto.streaming_upload_peak(crypto.STREAMING_GOVERNOR_CLIENT_PART_BYTES) assert r_huge == routine assert r_huge < 312 * MB assert scylla + r_huge < 312 * MB @@ -390,4 +396,4 @@ def test_copy_peak_is_bounded_and_independent_of_object_size(): assert peak == crypto.copy_pipeline_peak(64 * MB) # size-independent assert peak < 128 * MB # far below the old object/20 peak that forced full-budget clamping - assert peak < crypto.streaming_upload_peak(huge) // 4 + assert peak < _old_streaming_upload_peak(huge) // 4 diff --git a/tests/unit/test_memory_concurrency.py b/tests/unit/test_memory_concurrency.py index 2997c72..de4c592 100644 --- a/tests/unit/test_memory_concurrency.py +++ b/tests/unit/test_memory_concurrency.py @@ -50,10 +50,9 @@ def test_small_file_reserves_streaming_peak(self): assert footprint == crypto.streaming_upload_peak(100 * 1024) == 400 * 1024 def test_large_file_reserves_framed_peak_not_part_size(self): - """Large PUTs must reserve the framed path's true peak, NOT the bare - internal-part size -- reserving the part size under-counted ~3x and let - the limiter admit too many concurrent uploads (the OOM). The peak must - strictly exceed the part size.""" + """Large PUTs must reserve the framed path's true peak. With frame + streaming that peak is O(frame) regardless of Content-Length, so every + multi-frame upload reserves the same flat 4*FRAME_PLAINTEXT_SIZE.""" import s3proxy.concurrency as concurrency_module from s3proxy import crypto @@ -63,7 +62,7 @@ def test_large_file_reserves_framed_peak_not_part_size(self): assert footprint == crypto.governor_memory_footprint(cl) if mb <= 512: assert footprint == crypto.streaming_upload_peak(cl) - assert footprint >= crypto.memory_bounded_part_size(cl) + assert footprint == 4 * crypto.FRAME_PLAINTEXT_SIZE def test_minimum_reservation_enforced(self): """0-byte file should still reserve MIN_RESERVATION (64KB).""" diff --git a/tests/unit/test_streaming_framed_upload.py b/tests/unit/test_streaming_framed_upload.py index b76846e..1142c8c 100644 --- a/tests/unit/test_streaming_framed_upload.py +++ b/tests/unit/test_streaming_framed_upload.py @@ -48,7 +48,9 @@ async def test_framed_upload_roundtrips_and_sets_metadata(): class _Client: async def upload_part(self, bucket, key, upload_id, part_number, body): - captured[part_number] = body + if not isinstance(body, (bytes, bytearray)): + body = b"".join([chunk async for chunk in body]) + captured[part_number] = bytes(body) return {"ETag": f'"{part_number:032x}"'} mgr = _Manager() @@ -85,6 +87,11 @@ async def upload_part(self, bucket, key, upload_id, part_number, body): class _DiscardingClient: async def upload_part(self, bucket, key, upload_id, part_number, body): + # Consume streaming bodies chunk-by-chunk, retaining nothing, like the + # real (unsigned, streaming) transport. + if not isinstance(body, (bytes, bytearray)): + async for _chunk in body: + pass del body return {"ETag": f'"{part_number:032x}"'} @@ -128,6 +135,75 @@ async def _measure_framed_peak(client_part_size: int) -> int: return peak +@pytest.mark.asyncio +async def test_mid_part_backend_failure_surfaces_to_client(): + """A streamed body can't be replayed by botocore: a mid-part backend + failure must propagate (the client re-sends the part) and must not record + the part in upload state.""" + + class _FailingClient: + async def upload_part(self, bucket, key, upload_id, part_number, body): + declared = len(body) + consumed = 0 + async for chunk in body: + consumed += len(chunk) + if part_number >= 2 and consumed > declared // 2: + raise ConnectionError("backend reset mid-part") + return {"ETag": f'"{part_number:032x}"'} + + mgr = _Manager() + plaintext = os.urandom(30 * 1024 * 1024) + optimal = 12 * 1024 * 1024 + with pytest.raises(ConnectionError): + await _handler(mgr)._stream_and_upload_framed( + _Request(plaintext), + _FailingClient(), + "b", + "k", + UPLOAD_ID, + 1, + types.SimpleNamespace(dek=crypto.generate_dek()), + len(plaintext), + optimal, + 1, + 3, + ) + assert mgr.part_meta is None + + +@pytest.mark.asyncio +async def test_body_shorter_than_content_length_is_bad_request(): + from s3proxy.errors import S3Error + + class _ConsumingClient: + async def upload_part(self, bucket, key, upload_id, part_number, body): + if not isinstance(body, (bytes, bytearray)): + async for _ in body: + pass + return {"ETag": f'"{part_number:032x}"'} + + mgr = _Manager() + plaintext = os.urandom(10 * 1024 * 1024) + declared = 30 * 1024 * 1024 # stream ends 20MB early + with pytest.raises(S3Error) as exc_info: + await _handler(mgr)._stream_and_upload_framed( + _Request(plaintext), + _ConsumingClient(), + "b", + "k", + UPLOAD_ID, + 1, + types.SimpleNamespace(dek=crypto.generate_dek()), + declared, + 12 * 1024 * 1024, + 1, + 3, + ) + assert exc_info.value.status_code == 400 + assert "ended early" in exc_info.value.message + assert mgr.part_meta is None + + @pytest.mark.asyncio async def test_framed_upload_memory_is_independent_of_part_size(): """Peak memory tracks one internal part, not the client part size. With the diff --git a/tests/unit/test_upload_reservation_vs_real.py b/tests/unit/test_upload_reservation_vs_real.py index 22c9fc7..e34c1ea 100644 --- a/tests/unit/test_upload_reservation_vs_real.py +++ b/tests/unit/test_upload_reservation_vs_real.py @@ -33,10 +33,16 @@ async def add_part(self, *a, **k): class _Client: async def upload_part(self, bucket, key, upload_id, part_number, body): - # aiobotocore copies the body to sign and send it; mirror that so the + # The backend hop is UNSIGNED-PAYLOAD with streaming bodies: the + # transport consumes frames without retaining them. Mirror that so the # measured peak reflects what the real transport holds. - sent = bytes(body) - return {"ETag": hashlib.md5(sent).hexdigest()} + md5 = hashlib.md5(usedforsecurity=False) + if isinstance(body, (bytes, bytearray)): + md5.update(body) + else: + async for chunk in body: + md5.update(chunk) + return {"ETag": md5.hexdigest()} class _Request: @@ -87,14 +93,17 @@ async def test_reservation_bounds_real_framed_peak(mb): real_peak = await _measure_peak(content_length) reserved = estimate_memory_footprint("PUT", content_length) - # The whole point: the reservation must cover the REAL peak, not the bare - # part size. A bare-part-size reservation (the old behaviour) fails here. + # The whole point: the reservation must cover the REAL peak. assert reserved >= real_peak, ( f"{mb}MB part: reserved {reserved / MB:.1f}MB < real peak " f"{real_peak / MB:.1f}MB -- governor under-counts, will OOM" ) - # And the old under-count must be demonstrably insufficient. - assert crypto.memory_bounded_part_size(content_length) < real_peak + # Frame streaming: the real peak must stay O(frame), never O(part). + if crypto.memory_bounded_part_size(content_length) > crypto.FRAME_PLAINTEXT_SIZE: + assert real_peak <= 4 * crypto.FRAME_PLAINTEXT_SIZE, ( + f"{mb}MB part: real peak {real_peak / MB:.1f}MB scales with part size " + f"-- frame streaming regressed to whole-part buffering" + ) # Deployed governor budget (chart/values.yaml performance.memoryLimitMb). @@ -119,21 +128,22 @@ def test_es_part_fits_deploy_budget_with_concurrency(): @pytest.mark.asyncio -async def test_scylla_50mb_signed_part_reservation_covers_transport_copy(): - """Prod Scylla SST parts are ~50MB signed uploads with aiobotocore body copy.""" +async def test_scylla_50mb_part_peak_is_o_frame_and_covered(): + """Prod Scylla SST parts are ~50MB uploads; frame streaming keeps the real + peak at a few frames and the reservation must still cover it.""" content_length = 50 * MB real_peak = await _measure_peak(content_length) reserved = estimate_memory_footprint("PUT", content_length) - assert 24 * MB <= real_peak <= 40 * MB + assert real_peak <= 36 * MB assert reserved >= real_peak assert reserved < 40 * MB -def test_multi_gb_content_length_capped_not_honest_peak(): - """6.6GB Content-Length must not reserve the honest ~600MB+ signed peak.""" +def test_multi_gb_content_length_reserves_o_frame_not_o_part(): + """6.6GB Content-Length reserves the same flat O(frame) peak as a 512MB part.""" cl = 6_597_946_546 honest = crypto.streaming_upload_peak(cl) reserved = estimate_memory_footprint("PUT", cl) - assert honest > 500 * MB + assert honest == crypto.streaming_upload_peak(512 * MB) + assert reserved == honest assert reserved < 80 * MB - assert reserved < honest From 13bdf5950c7eb1de1fa82bc9ebceba8eddc23b05 Mon Sep 17 00:00:00 2001 From: serversidehannes Date: Fri, 10 Jul 2026 16:02:21 +0200 Subject: [PATCH 2/2] fix: clamp copy expected bytes to real segment extent; recalibrate remaining governor tests Scylla manifest sidecars overstate total_plaintext_size relative to the stored segments; the pump's strict Content-Length expectation must come from the segments' real extent or every inflated-metadata copy aborts as a short source. Also update the four remaining tests pinned to the old peak formulas. --- s3proxy/handlers/multipart/copy.py | 19 ++++++++++++++++++- tests/unit/test_copy_memory_governing.py | 3 +-- tests/unit/test_copy_pipeline_cap.py | 4 ++-- tests/unit/test_governor_peak_and_clamp.py | 13 +++++++++---- 4 files changed, 30 insertions(+), 9 deletions(-) diff --git a/s3proxy/handlers/multipart/copy.py b/s3proxy/handlers/multipart/copy.py index f270f9f..d1fc6bd 100644 --- a/s3proxy/handlers/multipart/copy.py +++ b/s3proxy/handlers/multipart/copy.py @@ -1234,6 +1234,23 @@ def open_source(skip_bytes: int) -> AsyncIterator[bytes]: skip_bytes=skip_bytes, ) + # Streamed bodies declare Content-Length up front, so the pump needs the + # number of bytes the source will actually yield. Scylla manifest sidecar + # metadata overstates total_plaintext_size relative to the stored + # segments (prod shape), so clamp the requested range to the segments' + # real extent — the claimed total would over-promise and abort the copy. + source_bytes = plaintext_size + if src_multipart_meta: + real_total = sum(p.plaintext_size for p in src_multipart_meta.parts) + if copy_source_range: + range_start, range_end = self._parse_copy_source_range( + copy_source_range, src_multipart_meta.total_plaintext_size + ) + else: + range_start, range_end = 0, src_multipart_meta.total_plaintext_size - 1 + range_end = min(range_end, real_total - 1) + source_bytes = max(0, range_end - range_start + 1) + internal_parts, total_plaintext, total_ciphertext, md5 = await self._pump_copy_chunks( client, bucket, @@ -1245,7 +1262,7 @@ def open_source(skip_bytes: int) -> AsyncIterator[bytes]: chunk_size, internal_part_start, leading_plaintext=deferred_tail or b"", - expected_plaintext=plaintext_size + len(deferred_tail), + expected_plaintext=source_bytes + len(deferred_tail), reopen_source=open_source, ) diff --git a/tests/unit/test_copy_memory_governing.py b/tests/unit/test_copy_memory_governing.py index ea65179..e5e2f0b 100644 --- a/tests/unit/test_copy_memory_governing.py +++ b/tests/unit/test_copy_memory_governing.py @@ -29,8 +29,7 @@ def test_copy_pipeline_peak_is_bounded_regardless_of_object_size(): # ~90MB, identical to a 64MB copy. peaks = [crypto.copy_pipeline_peak(s) for s in (64 * MB, 512 * MB, 4767 * MB, 5 * 1024 * MB)] assert all(p == peaks[0] for p in peaks), peaks - part = crypto.COPY_INTERNAL_PART_SIZE - assert peaks[0] == 2 * part + crypto.FRAME_PLAINTEXT_SIZE + 2 * crypto.MAX_BUFFER_SIZE + assert peaks[0] == 4 * crypto.FRAME_PLAINTEXT_SIZE + 2 * crypto.MAX_BUFFER_SIZE assert peaks[0] < 128 * MB diff --git a/tests/unit/test_copy_pipeline_cap.py b/tests/unit/test_copy_pipeline_cap.py index 9c20cd1..ab6b03f 100644 --- a/tests/unit/test_copy_pipeline_cap.py +++ b/tests/unit/test_copy_pipeline_cap.py @@ -114,8 +114,8 @@ async def test_backpressure_logged_once_per_acquire_wait(): original_timeout = concurrency.BACKPRESSURE_TIMEOUT concurrency.BACKPRESSURE_TIMEOUT = 5 concurrency.reset_state() - # 96MB budget: one 88MB chunk leaves no room for a second → must wait. - concurrency.set_memory_limit(96) + # Budget fits one ~48MB chunk but not two → the second acquire must wait. + concurrency.set_memory_limit(64) reset_copy_pipeline_semaphore(1) held = await concurrency.try_acquire_copy_memory( diff --git a/tests/unit/test_governor_peak_and_clamp.py b/tests/unit/test_governor_peak_and_clamp.py index deb7ede..d5d9ae9 100644 --- a/tests/unit/test_governor_peak_and_clamp.py +++ b/tests/unit/test_governor_peak_and_clamp.py @@ -205,7 +205,7 @@ def test_streaming_governor_clamped_reserve(): routine = crypto.streaming_upload_peak(crypto.STREAMING_GOVERNOR_CLIENT_PART_BYTES) assert crypto.streaming_governor_clamped_reserve(honest, budget) == routine - small = 40 * MB + small = 20 * MB assert crypto.streaming_governor_clamped_reserve(small, budget) == small @@ -369,17 +369,22 @@ async def test_request_handler_put_reserves_governor_footprint(): @pytest.mark.asyncio async def test_many_concurrent_scylla_acquires_via_gather(): - """Flood of concurrent ~50MB part admissions under prod budget.""" + """Flood of concurrent ~50MB part admissions under prod budget. + + At the flat 32MB streaming reservation a 312MB budget admits 9 at once. + """ reset_state() set_memory_limit(312) try: per_part = estimate_memory_footprint("PUT", SCYLLA_PART_BYTES) + count = (312 * MB) // per_part + assert count >= 8 async def one(): return await try_acquire_memory(per_part) - results = await asyncio.gather(*[one() for _ in range(12)]) - assert len(results) == 12 + results = await asyncio.gather(*[one() for _ in range(count)]) + assert len(results) == count assert sum(results) <= 312 * MB assert all(r == per_part for r in results) finally: