Stream sealed frames as upload bodies: O(part) → O(frame) memory

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,23 +265,19 @@ 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)
 
 
 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.

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,35 @@ 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,
+            )
+
+        # 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,
@@ -1233,10 +1258,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=source_bytes + len(deferred_tail),
+            reopen_source=open_source,
         )
 
         etag = md5.hexdigest()
@@ -1282,14 +1309,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 +1345,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 +1382,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 +1462,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