fix: accurate memory accounting for encrypted GETs and streaming PUTs

ServerSideHannes · ServerSideHannes · commit 97e811cfd26a · 2026-02-11T08:02:37.000+01:00
- Fix concurrency limiter to reject requests exceeding budget instead of
  silently capping reservations
- Fix PUT memory estimate: streaming PUTs hold buffer + ciphertext (16MB
  not 8MB)
- Add dynamic memory acquisition in GET handler for encrypted decrypts
  (ciphertext + plaintext buffered simultaneously)
- Add container-based OOM proof test (256MB container, 5GB+ data)
- Split CI into parallel unit/integration jobs, separate OOM workflow
- Add make test-integration target
diff --git a/.github/workflows/oom-test.yml b/.github/workflows/oom-test.yml
@@ -0,0 +1,31 @@
+name: OOM Proof Test
+
+on:
+  push:
+    branches: [main]
+    paths:
+      - 's3proxy/**'
+      - 'tests/**'
+  workflow_dispatch:
+
+jobs:
+  oom-test:
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Set up Python
+        uses: actions/setup-python@v6
+        with:
+          python-version: '3.14'
+          cache: 'pip'
+
+      - name: Install uv
+        run: pip install uv
+
+      - name: Install dependencies
+        run: uv sync --extra dev
+
+      - name: OOM proof test (256MB container)
+        run: make test-oom
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -11,9 +11,9 @@ on:
   workflow_dispatch:
 
 jobs:
-  tests:
+  unit:
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 10
     steps:
       - uses: actions/checkout@v6
 
@@ -29,8 +29,26 @@ jobs:
       - name: Install dependencies
         run: uv sync --extra dev
 
-      - name: Run all tests
-        run: make test-all
+      - name: Run unit tests
+        run: make test-unit
 
-      - name: OOM proof test (128MB container)
-        run: make test-oom
+  integration:
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Set up Python
+        uses: actions/setup-python@v6
+        with:
+          python-version: '3.14'
+          cache: 'pip'
+
+      - name: Install uv
+        run: pip install uv
+
+      - name: Install dependencies
+        run: uv sync --extra dev
+
+      - name: Run integration tests
+        run: make test-integration
diff --git a/Makefile b/Makefile
@@ -1,4 +1,4 @@
-.PHONY: test test-all test-unit test-run test-oom e2e cluster lint
+.PHONY: test test-all test-unit test-integration test-run test-oom e2e cluster lint
 
 # Lint: ruff check + format check
 lint:
@@ -12,7 +12,17 @@ test: test-unit
 test-unit:
 	uv run pytest -m "not e2e and not ha" -v -n auto
 
-# Run all tests with containers (parallel execution)
+# Run integration tests (needs minio/redis containers)
+test-integration:
+	@docker compose -f tests/docker-compose.yml down 2>/dev/null || true
+	@docker compose -f tests/docker-compose.yml up -d
+	@sleep 3
+	@AWS_ACCESS_KEY_ID=minioadmin AWS_SECRET_ACCESS_KEY=minioadmin uv run pytest -m "e2e" -v -n auto --dist loadgroup; \
+		EXIT_CODE=$$?; \
+		docker compose -f tests/docker-compose.yml down; \
+		exit $$EXIT_CODE
+
+# Run all tests with containers (unit + integration)
 test-all:
 	@docker compose -f tests/docker-compose.yml down 2>/dev/null || true
 	@docker compose -f tests/docker-compose.yml up -d
@@ -33,7 +43,7 @@ test-run:
 		docker compose -f tests/docker-compose.yml down; \
 		exit $$EXIT_CODE
 
-# OOM proof test: runs s3proxy in a 128MB container and hammers it
+# OOM proof test: runs s3proxy in a 256MB container and hammers it
 test-oom:
 	@docker compose -f tests/docker-compose.yml --profile oom down 2>/dev/null || true
 	@docker compose -f tests/docker-compose.yml --profile oom up -d --build
diff --git a/s3proxy/concurrency.py b/s3proxy/concurrency.py
@@ -84,7 +84,21 @@ async def try_acquire(self, bytes_needed: int) -> int:
         if self._limit_bytes <= 0:
             return 0
 
-        to_reserve = max(MIN_RESERVATION, min(bytes_needed, self._limit_bytes))
+        to_reserve = max(MIN_RESERVATION, bytes_needed)
+
+        # Single request exceeds entire budget — can never fit, reject immediately
+        if to_reserve > self._limit_bytes:
+            request_mb = to_reserve / 1024 / 1024
+            limit_mb = self._limit_bytes / 1024 / 1024
+            logger.warning(
+                "MEMORY_TOO_LARGE",
+                requested_mb=round(request_mb, 2),
+                limit_mb=round(limit_mb, 2),
+            )
+            MEMORY_REJECTIONS.inc()
+            raise S3Error.slow_down(
+                f"Request needs {request_mb:.0f}MB but budget is {limit_mb:.0f}MB"
+            )
 
         async with self._condition:
             deadline = asyncio.get_event_loop().time() + BACKPRESSURE_TIMEOUT
@@ -147,7 +161,12 @@ async def release(self, bytes_reserved: int) -> None:
 
 
 def estimate_memory_footprint(method: str, content_length: int) -> int:
-    """Estimate memory needed for a request."""
+    """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.
+    GETs reserve a baseline here; encrypted GETs acquire additional memory in the handler.
+    """
     if method in ("HEAD", "DELETE"):
         return 0
     if method == "GET":
@@ -156,7 +175,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
+    return MAX_BUFFER_SIZE * 2
 
 
 # Module-level convenience functions delegating to the default instance
diff --git a/s3proxy/handlers/objects/get.py b/s3proxy/handlers/objects/get.py
@@ -11,7 +11,8 @@
 from fastapi.responses import StreamingResponse
 from structlog.stdlib import BoundLogger
 
-from ... import crypto
+from ... import concurrency, crypto
+from ...concurrency import MAX_BUFFER_SIZE
 from ...errors import S3Error
 from ...s3client import S3Client, S3Credentials
 from ...state import (
@@ -147,37 +148,51 @@ async def _decrypt_single_object(
     ) -> Response:
         logger.info("GET_ENCRYPTED_SINGLE", bucket=bucket, key=key)
         resp = await client.get_object(bucket, key)
-        wrapped_dek = base64.b64decode(wrapped_dek_b64)
-        # Read and close body to release aioboto3/aiohttp resources
-        async with resp["Body"] as body:
-            ciphertext = await body.read()
-        plaintext = crypto.decrypt_object(ciphertext, wrapped_dek, self.settings.kek)
-        del ciphertext  # Free memory
+        content_length = resp.get("ContentLength", 0)
 
-        content_type = head_resp.get("ContentType", "application/octet-stream")
-        cache_control = head_resp.get("CacheControl")
-        expires = head_resp.get("Expires")
+        # Encrypted decrypts buffer ciphertext + plaintext simultaneously.
+        # Acquire additional memory beyond the initial MAX_BUFFER_SIZE reservation.
+        additional = max(0, content_length * 2 - MAX_BUFFER_SIZE)
+        extra_reserved = 0
+        try:
+            if additional > 0:
+                extra_reserved = await concurrency.try_acquire_memory(additional)
+
+            wrapped_dek = base64.b64decode(wrapped_dek_b64)
+            async with resp["Body"] as body:
+                ciphertext = await body.read()
+            plaintext = crypto.decrypt_object(ciphertext, wrapped_dek, self.settings.kek)
+            del ciphertext
+
+            content_type = head_resp.get("ContentType", "application/octet-stream")
+            cache_control = head_resp.get("CacheControl")
+            expires = head_resp.get("Expires")
+
+            if range_header:
+                start, end = self._parse_range(range_header, len(plaintext))
+                headers = self._build_headers(
+                    content_type=content_type,
+                    content_length=end - start + 1,
+                    last_modified=last_modified,
+                    cache_control=cache_control,
+                    expires=expires,
+                )
+                headers["Content-Range"] = f"bytes {start}-{end}/{len(plaintext)}"
+                return Response(
+                    content=plaintext[start : end + 1], status_code=206, headers=headers
+                )
 
-        if range_header:
-            start, end = self._parse_range(range_header, len(plaintext))
             headers = self._build_headers(
                 content_type=content_type,
-                content_length=end - start + 1,
+                content_length=len(plaintext),
                 last_modified=last_modified,
                 cache_control=cache_control,
                 expires=expires,
             )
-            headers["Content-Range"] = f"bytes {start}-{end}/{len(plaintext)}"
-            return Response(content=plaintext[start : end + 1], status_code=206, headers=headers)
-
-        headers = self._build_headers(
-            content_type=content_type,
-            content_length=len(plaintext),
-            last_modified=last_modified,
-            cache_control=cache_control,
-            expires=expires,
-        )
-        return Response(content=plaintext, headers=headers)
+            return Response(content=plaintext, headers=headers)
+        finally:
+            if extra_reserved > 0:
+                await concurrency.release_memory(extra_reserved)
 
     async def _get_multipart(
         self,
@@ -378,13 +393,17 @@ async def _fetch_internal_part(
         ct_end: int,
         dek: bytes,
     ) -> bytes:
+        expected_size = ct_end - ct_start + 1
+        additional = max(0, expected_size * 2 - MAX_BUFFER_SIZE)
+        extra_reserved = 0
         try:
+            if additional > 0:
+                extra_reserved = await concurrency.try_acquire_memory(additional)
+
             resp = await client.get_object(bucket, key, f"bytes={ct_start}-{ct_end}")
-            # Read and close body to release aioboto3/aiohttp resources
             async with resp["Body"] as body:
                 ciphertext = await body.read()
 
-            expected_size = ct_end - ct_start + 1
             if len(ciphertext) < crypto.ENCRYPTION_OVERHEAD or len(ciphertext) != expected_size:
                 logger.error(
                     "GET_CIPHERTEXT_SIZE_MISMATCH",
@@ -419,6 +438,9 @@ async def _fetch_internal_part(
                     f"range {ct_start}-{ct_end} invalid"
                 ) from e
             raise
+        finally:
+            if extra_reserved > 0:
+                await concurrency.release_memory(extra_reserved)
 
     async def _fetch_and_decrypt_part(
         self,
@@ -445,12 +467,21 @@ async def _fetch_and_decrypt_part(
 
         self._validate_ciphertext_range(bucket, key, part_num, 0, ct_end, actual_size)
 
-        resp = await client.get_object(bucket, key, f"bytes={ct_start}-{ct_end}")
-        # Read and close body to release aioboto3/aiohttp resources
-        async with resp["Body"] as body:
-            ciphertext = await body.read()
-        decrypted = crypto.decrypt(ciphertext, dek)
-        return decrypted[off_start : off_end + 1]
+        part_size = part_meta.ciphertext_size
+        additional = max(0, part_size * 2 - MAX_BUFFER_SIZE)
+        extra_reserved = 0
+        try:
+            if additional > 0:
+                extra_reserved = await concurrency.try_acquire_memory(additional)
+
+            resp = await client.get_object(bucket, key, f"bytes={ct_start}-{ct_end}")
+            async with resp["Body"] as body:
+                ciphertext = await body.read()
+            decrypted = crypto.decrypt(ciphertext, dek)
+            return decrypted[off_start : off_end + 1]
+        finally:
+            if extra_reserved > 0:
+                await concurrency.release_memory(extra_reserved)
 
     def _build_response_headers(self, resp: dict, last_modified: str | None) -> dict[str, str]:
         return self._build_headers(
diff --git a/tests/docker-compose.yml b/tests/docker-compose.yml
@@ -22,14 +22,14 @@ services:
       context: ../
       dockerfile: Dockerfile
     container_name: s3proxy-test-server
-    mem_limit: 128m
+    mem_limit: 256m
     ports:
       - "4433:4433"
     environment:
       S3PROXY_ENCRYPT_KEY: "test-encryption-key-32-bytes!!"
       S3PROXY_HOST: "http://minio:9000"
       S3PROXY_REGION: "us-east-1"
-      S3PROXY_MEMORY_LIMIT_MB: "16"
+      S3PROXY_MEMORY_LIMIT_MB: "48"
       S3PROXY_LOG_LEVEL: "WARNING"
       AWS_ACCESS_KEY_ID: minioadmin
       AWS_SECRET_ACCESS_KEY: minioadmin
diff --git a/tests/integration/test_memory_leak.py b/tests/integration/test_memory_leak.py
diff --git a/tests/unit/test_concurrency_limit.py b/tests/unit/test_concurrency_limit.py
diff --git a/tests/unit/test_memory_concurrency.py b/tests/unit/test_memory_concurrency.py
