fix: bound UploadPart memory with adaptive internal part size (#80)

Author: ServerSideHannes

.python-version

@@ -0,0 +1 @@
+3.14

s3proxy/concurrency.py

@@ -12,6 +12,7 @@
 
 import structlog
 
+from s3proxy.crypto import memory_bounded_part_size
 from s3proxy.errors import S3Error
 from s3proxy.metrics import MEMORY_LIMIT_BYTES, MEMORY_REJECTIONS, MEMORY_RESERVED_BYTES
 
@@ -163,8 +164,11 @@ async def release(self, bytes_reserved: int) -> None:
 def estimate_memory_footprint(method: str, content_length: int) -> int:
     """Estimate memory needed for a request.
 
-    Streaming PUTs hold an 8MB plaintext buffer + 8MB ciphertext simultaneously,
-    so large PUTs need 2x MAX_BUFFER_SIZE. Small PUTs buffer the whole body + ciphertext.
+    Small PUTs buffer the whole body + ciphertext. Larger PUTs stream and buffer
+    one internal part at a time, so reserve exactly that internal part size --
+    the same value the upload path uses (memory_bounded_part_size). This keeps
+    the reservation honest: the limiter then admits only as many concurrent
+    uploads as actually fit the budget, instead of under-counting and OOMing.
     GETs reserve a baseline here; encrypted GETs acquire additional memory in the handler.
     """
     if method in ("HEAD", "DELETE"):
@@ -175,7 +179,7 @@ def estimate_memory_footprint(method: str, content_length: int) -> int:
         return MIN_RESERVATION
     if content_length <= MAX_BUFFER_SIZE:
         return max(MIN_RESERVATION, content_length * 2)
-    return MAX_BUFFER_SIZE * 2
+    return memory_bounded_part_size(content_length)
 
 
 # Module-level convenience functions delegating to the default instance

s3proxy/crypto.py

@@ -37,6 +37,11 @@
 # Sweet spot: 8MB balances part count vs per-part overhead
 MAX_BUFFER_SIZE = 8 * 1024 * 1024  # 8 MB per internal part
 
+# Internal part numbers are allocated in fixed ranges of this size per client
+# part (see state.manager), so a single client part may not expand into more
+# than this many internal parts without colliding into the next part's range.
+MAX_INTERNAL_PARTS_PER_CLIENT = 20
+
 # Framed internal-part format.
 # S3's 10,000-part limit forces large internal parts for big objects, but we do
 # not want to hold a whole part (plaintext + ciphertext) in memory. So a part is
@@ -135,6 +140,32 @@ def calculate_optimal_part_size(content_length: int) -> int:
     return PART_SIZE
 
 
+def memory_bounded_part_size(
+    content_length: int, max_parts: int = MAX_INTERNAL_PARTS_PER_CLIENT
+) -> int:
+    """Smallest internal part size that bounds memory while respecting limits.
+
+    The framed upload path buffers exactly one internal part at a time, so peak
+    memory tracks the part size. We therefore want parts as small as possible —
+    but two limits stop us going arbitrarily small:
+
+      * each internal part must be >= ~MAX_BUFFER_SIZE (avoids tiny <5MB S3 parts
+        and per-part overhead), and
+      * a client part may expand into at most ``max_parts`` internal parts
+        (the per-client part-number allocation range).
+
+    So we use as many parts as those limits allow and split evenly. For barman's
+    512MB client parts this yields ~26MB parts (20 of them) instead of 64MB,
+    roughly halving per-request memory, and it adapts up only for client parts
+    large enough to need it.
+    """
+    if content_length <= MAX_BUFFER_SIZE:
+        return content_length or 1
+    parts = min(max_parts, content_length // MAX_BUFFER_SIZE)
+    parts = max(parts, 1)
+    return -(-content_length // parts)  # ceil: even split, never exceeds `parts`
+
+
 @dataclass(slots=True)
 class EncryptedData:
     """Container for encrypted data and metadata."""

s3proxy/handlers/multipart/upload_part.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import asyncio
-import gc
 import hashlib
 import time
 from collections import deque
@@ -97,9 +96,11 @@ async def handle_upload_part(self, request: Request, creds: S3Credentials) -> Re
                 and content_length > crypto.STREAMING_THRESHOLD
             )
 
-            # Calculate optimal part size
-            optimal_part_size = crypto.calculate_optimal_part_size(content_length)
-            estimated_parts = max(1, (content_length + optimal_part_size - 1) // optimal_part_size)
+            # Smallest internal part that bounds memory while staying within the
+            # per-client part-number allocation range (so we never collide and
+            # never buffer more than necessary).
+            internal_part_size = crypto.memory_bounded_part_size(content_length)
+            estimated_parts = max(1, -(-content_length // internal_part_size))
 
             use_framed = (
                 (is_unsigned or is_large_signed) and not needs_chunked_decode and content_length > 0
@@ -109,7 +110,7 @@ async def handle_upload_part(self, request: Request, creds: S3Credentials) -> Re
                 bucket=bucket,
                 key=key,
                 part_number=part_num,
-                optimal_part_size_mb=f"{optimal_part_size / 1024 / 1024:.2f}MB",
+                internal_part_size_mb=f"{internal_part_size / 1024 / 1024:.2f}MB",
                 estimated_internal_parts=estimated_parts,
                 is_unsigned=is_unsigned,
                 is_large_signed=is_large_signed,
@@ -142,7 +143,7 @@ async def handle_upload_part(self, request: Request, creds: S3Credentials) -> Re
                         part_num,
                         state,
                         content_length,
-                        optimal_part_size,
+                        internal_part_size,
                         internal_part_start,
                     )
                 else:
@@ -160,7 +161,7 @@ async def handle_upload_part(self, request: Request, creds: S3Credentials) -> Re
                         is_streaming_sig,
                         is_large_signed,
                         needs_chunked_decode,
-                        optimal_part_size,
+                        internal_part_size,
                         internal_part_start,
                     )
 
@@ -403,12 +404,9 @@ async def _stream_and_upload_framed(
                 )
                 frame_idx += 1
 
-            part_ciphertext = bytes(ciphertext)
-            del ciphertext
-            gc.collect()
-
             upload_start = time.monotonic()
-            resp = await client.upload_part(bucket, key, upload_id, ipn, part_ciphertext)
+            resp = await client.upload_part(bucket, key, upload_id, ipn, ciphertext)
+            del ciphertext
             etag = resp["ETag"].strip('"')
             logger.info(
                 "INTERNAL_PART_UPLOADED",

s3proxy/state/manager.py

@@ -3,6 +3,7 @@
 import structlog
 from structlog.stdlib import BoundLogger
 
+from ..crypto import MAX_INTERNAL_PARTS_PER_CLIENT
 from .models import (
     MultipartUploadState,
     PartMetadata,
@@ -13,8 +14,9 @@
 
 logger: BoundLogger = structlog.get_logger(__name__)
 
-# Maximum internal parts per client part (for range allocation)
-MAX_INTERNAL_PARTS_PER_CLIENT = 20
+# Re-exported from crypto so part sizing and part-number allocation share one
+# source of truth.
+__all__ = ["MultipartStateManager", "MAX_INTERNAL_PARTS_PER_CLIENT"]
 
 
 class MultipartStateManager:

tests/unit/test_concurrency_limit.py

@@ -331,12 +331,40 @@ def test_estimate_memory_footprint_put_small(self):
         assert footprint == 2 * 1024 * 1024
 
     def test_estimate_memory_footprint_put_large(self):
-        """PUT with large file should use 2x buffer size (buffer + ciphertext)."""
+        """Large PUT reserves the real internal-part buffer the upload holds."""
         import s3proxy.concurrency as concurrency_module
-
-        # 100MB file → 16MB footprint (8MB buffer + 8MB ciphertext simultaneously)
-        footprint = concurrency_module.estimate_memory_footprint("PUT", 100 * 1024 * 1024)
-        assert footprint == concurrency_module.MAX_BUFFER_SIZE * 2
+        from s3proxy import crypto
+
+        for mb in (50, 100, 512, 1024):
+            cl = mb * 1024 * 1024
+            footprint = concurrency_module.estimate_memory_footprint("PUT", cl)
+            assert footprint == crypto.memory_bounded_part_size(cl)
+
+    def test_large_uploads_bounded_below_pod_memory(self):
+        """Regression for the barman OOM. Two linked invariants:
+        1. an internal part never expands beyond the per-client allocation range
+           (or part numbers collide) -- for ANY client part size, and
+        2. the reservation tracks the real buffer, so admitted x footprint never
+           exceeds the budget (limiter guarantee), and barman-scale parts admit
+           only ~2 concurrent (the old flat-16MB estimate admitted ~4 -> OOM).
+        """
+        import s3proxy.concurrency as concurrency_module
+        from s3proxy import crypto
+        from s3proxy.state import MAX_INTERNAL_PARTS_PER_CLIENT
+
+        budget = concurrency_module.get_memory_limit()
+        for mb in (50, 128, 320, 512, 1024, 4096):
+            cl = mb * 1024 * 1024
+            part = crypto.memory_bounded_part_size(cl)
+            internal_parts = -(-cl // part)
+            assert internal_parts <= MAX_INTERNAL_PARTS_PER_CLIENT, "would collide part numbers"
+            footprint = concurrency_module.estimate_memory_footprint("PUT", cl)
+            # limiter guarantee: total admitted memory never exceeds the budget
+            assert (budget // footprint) * footprint <= budget
+
+        # barman-scale parts: bounded to ~2 concurrent on the default 64MB budget
+        footprint_512 = concurrency_module.estimate_memory_footprint("PUT", 512 * 1024 * 1024)
+        assert budget // footprint_512 <= 2
 
     def test_estimate_memory_footprint_get(self):
         """GET should always use fixed buffer size."""

tests/unit/test_framed_crypto.py

@@ -66,3 +66,34 @@ def test_frame_nonces_unique():
     assert len(nonces) == 1000
     # And distinct from the legacy part nonce so an upload never reuses one.
     assert crypto.derive_part_nonce(UPLOAD_ID, PART) not in nonces
+
+
+MB = 1024 * 1024
+
+
+@pytest.mark.parametrize(
+    "content_mb",
+    [9, 16, 50, 100, 160, 200, 512, 1024, 4096, 10_000],
+)
+def test_memory_bounded_part_size_respects_limits(content_mb):
+    """For any client part size: never expand beyond the per-client allocation
+    range (or S3 part numbers collide), and never create a non-final part below
+    S3's 5MB minimum."""
+    cl = content_mb * MB
+    size = crypto.memory_bounded_part_size(cl)
+    parts = -(-cl // size)
+    assert parts <= crypto.MAX_INTERNAL_PARTS_PER_CLIENT
+    # every part except possibly the last is `size`; the last is the remainder
+    last = cl - (parts - 1) * size
+    if parts > 1:
+        assert last >= crypto.MIN_PART_SIZE
+        assert size >= crypto.MIN_PART_SIZE
+
+
+def test_memory_bounded_part_size_is_small_until_forced_larger():
+    """Small/mid client parts stay ~8MB; size grows only when the 20-part cap
+    forces it (e.g. barman's 512MB parts -> ~26MB, not 64MB)."""
+    assert crypto.memory_bounded_part_size(50 * MB) <= 9 * MB
+    assert crypto.memory_bounded_part_size(160 * MB) <= 9 * MB
+    barman = crypto.memory_bounded_part_size(512 * MB)
+    assert 20 * MB <= barman <= 32 * MB

tests/unit/test_memory_concurrency.py

@@ -6,7 +6,8 @@
 
 Memory estimation logic:
 - PUT ≤8MB: content_length * 2 (body + ciphertext buffer)
-- PUT >8MB: MAX_BUFFER_SIZE * 2 (16MB, streaming buffer + ciphertext)
+- PUT >8MB: memory_bounded_part_size(content_length) (the one internal part the
+  streaming/framed upload path actually buffers at a time)
 - GET: MAX_BUFFER_SIZE (8MB baseline, handler acquires more for encrypted decrypts)
 - POST: MIN_RESERVATION (64KB, metadata only)
 - HEAD/DELETE: 0 (no buffering, bypass limit)
@@ -46,12 +47,17 @@ def test_small_file_uses_content_length_x2(self):
         footprint = concurrency_module.estimate_memory_footprint("PUT", 100 * 1024)
         assert footprint == 200 * 1024
 
-    def test_large_file_uses_double_buffer(self):
-        """PUT with 100MB file should reserve 16MB (buffer + ciphertext)."""
+    def test_large_file_reserves_real_internal_part(self):
+        """Large PUTs must reserve the actual internal-part buffer the upload
+        path holds (memory_bounded_part_size), not a flat guess -- otherwise the
+        limiter under-counts and admits too many concurrent uploads (the OOM)."""
         import s3proxy.concurrency as concurrency_module
+        from s3proxy import crypto
 
-        footprint = concurrency_module.estimate_memory_footprint("PUT", 100 * 1024 * 1024)
-        assert footprint == concurrency_module.MAX_BUFFER_SIZE * 2  # 16MB
+        for mb in (50, 100, 512, 1024):
+            cl = mb * 1024 * 1024
+            footprint = concurrency_module.estimate_memory_footprint("PUT", cl)
+            assert footprint == crypto.memory_bounded_part_size(cl)
 
     def test_minimum_reservation_enforced(self):
         """0-byte file should still reserve MIN_RESERVATION (64KB)."""
@@ -263,10 +269,10 @@ async def test_mixed_workload_scenario(self):
 
         reservations = []
 
-        # 2 large streaming uploads (16MB each = 32MB)
+        # 2 large streaming uploads (320MB -> 16MB internal part each = 32MB)
         for _ in range(2):
-            footprint = concurrency_module.estimate_memory_footprint("PUT", 100 * 1024 * 1024)
-            assert footprint == 16 * 1024 * 1024  # buffer + ciphertext
+            footprint = concurrency_module.estimate_memory_footprint("PUT", 320 * 1024 * 1024)
+            assert footprint == 16 * 1024 * 1024  # one 16MB internal part
             reserved = await concurrency_module.try_acquire_memory(footprint)
             reservations.append(reserved)
 

tests/unit/test_streaming_framed_upload.py

@@ -115,7 +115,7 @@ async def _measure_framed_peak(client_part_size: int) -> int:
         1,
         types.SimpleNamespace(dek=crypto.generate_dek()),
         client_part_size,
-        crypto.PART_SIZE,  # same internal part size barman uses for 512MB parts
+        crypto.memory_bounded_part_size(client_part_size),  # size the handler picks
         1,
     )
     _, peak = tracemalloc.get_traced_memory()
@@ -125,13 +125,18 @@ async def _measure_framed_peak(client_part_size: int) -> int:
 
 @pytest.mark.asyncio
 async def test_framed_upload_memory_is_independent_of_part_size():
-    """A 256MB client part peaks no higher than a 64MB one: ciphertext is built
-    one internal part at a time, not scaled by total client part size."""
+    """Peak memory tracks one internal part, not the client part size. With the
+    adaptive sizing a 512MB barman part uses ~26MB internal parts and peaks well
+    under what a single 64MB-part request used to (~80MB+), and a 1GB client
+    part does not peak meaningfully higher."""
     small = await _measure_framed_peak(64 * 1024 * 1024)
-    large = await _measure_framed_peak(256 * 1024 * 1024)
-
-    # Must stay below the old buffered path (~2× part_size ≈ 257MB per request).
-    assert small < 220 * 1024 * 1024, f"small peak {small / 1024 / 1024:.1f}MB"
-    assert large < 220 * 1024 * 1024, f"large peak {large / 1024 / 1024:.1f}MB"
-    # Larger client part must not scale memory linearly with part count.
-    assert large <= small * 1.5 + crypto.FRAME_PLAINTEXT_SIZE
+    barman = await _measure_framed_peak(512 * 1024 * 1024)
+    huge = await _measure_framed_peak(1024 * 1024 * 1024)
+
+    # One internal part + frame overhead. Generous ceiling that still catches a
+    # regression to buffering a whole 64MB part (~80MB) or the client part.
+    ceiling = 70 * 1024 * 1024
+    assert barman < ceiling, f"barman peak {barman / 1024 / 1024:.1f}MB"
+    assert small < ceiling, f"small peak {small / 1024 / 1024:.1f}MB"
+    # Memory must not scale with client part size.
+    assert huge <= barman * 1.5 + crypto.FRAME_PLAINTEXT_SIZE