diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 9d4fdb6..29637bf 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -30,7 +30,11 @@ jobs: integration: runs-on: ubuntu-latest - timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + shard: [memory_usage, memory_leak, memory_copy, core, multipart, copy_range, misc] + timeout-minutes: ${{ contains(fromJSON('["memory_usage", "memory_leak"]'), matrix.shard) && 15 || 10 }} steps: - uses: actions/checkout@v7 @@ -46,5 +50,5 @@ jobs: - name: Install dependencies run: uv sync --extra dev - - name: Run integration tests - run: make test-integration + - name: Run integration tests (${{ matrix.shard }}) + run: make test-integration-shard SHARD=${{ matrix.shard }} diff --git a/Makefile b/Makefile index 29369f2..50c37b5 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: test test-all test-unit test-integration test-run test-oom e2e cluster lint +.PHONY: test test-all test-unit test-integration test-integration-shard test-run test-oom e2e cluster lint # Lint: ruff check + format check lint: @@ -12,6 +12,38 @@ test: test-unit test-unit: uv run pytest -m "not e2e and not ha" -v -n auto +# Integration shards for parallel CI (make test-integration-shard SHARD=memory_usage) +INTEGRATION_memory_usage_TESTS = tests/integration/test_memory_usage.py +INTEGRATION_memory_leak_TESTS = tests/integration/test_memory_leak.py +INTEGRATION_memory_copy_TESTS = tests/integration/test_copy_memory_governor.py +INTEGRATION_memory_copy_PYTEST_OPTS = -n0 +INTEGRATION_core_TESTS = \ + tests/integration/test_integration.py \ + tests/integration/test_handlers.py \ + tests/integration/test_concurrent_operations.py \ + tests/integration/test_per_key_encryption.py \ + tests/integration/test_redis_coordination.py +INTEGRATION_multipart_TESTS = \ + tests/integration/test_upload_part_copy.py \ + tests/integration/test_part_ordering.py \ + tests/integration/test_sequential_part_numbering.py \ + tests/integration/test_sequential_part_numbering_e2e.py \ + tests/integration/test_large_file_streaming.py \ + tests/integration/test_entity_too_small_errors.py \ + tests/integration/test_entity_too_small_fix.py \ + tests/integration/test_partial_complete_fix.py \ + tests/integration/test_multipart_range_validation.py +INTEGRATION_copy_range_TESTS = \ + tests/integration/test_copy_passthrough.py \ + tests/integration/test_download_range_requests.py \ + tests/integration/test_get_prefetch.py \ + tests/integration/test_elasticsearch_range_scenario.py \ + tests/integration/test_state_recovery.py \ + tests/integration/test_state_recovery_e2e.py +INTEGRATION_misc_TESTS = \ + tests/integration/test_delete_objects_errors.py \ + tests/integration/test_metadata_and_errors.py + # Run integration tests (needs minio/redis containers) test-integration: @docker compose -f tests/docker-compose.yml down 2>/dev/null || true @@ -22,6 +54,22 @@ test-integration: docker compose -f tests/docker-compose.yml down; \ exit $$EXIT_CODE +# Run one integration shard (CI matrix). SHARD=memory_usage|memory_leak|memory_copy|... +test-integration-shard: +ifndef SHARD + $(error SHARD is required, e.g. make test-integration-shard SHARD=memory) +endif + @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 \ + $(if $(INTEGRATION_$(SHARD)_PYTEST_OPTS),$(INTEGRATION_$(SHARD)_PYTEST_OPTS),-n auto --dist loadgroup) \ + $(INTEGRATION_$(SHARD)_TESTS); \ + 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 diff --git a/s3proxy/concurrency.py b/s3proxy/concurrency.py index 0b45be1..103ebc0 100644 --- a/s3proxy/concurrency.py +++ b/s3proxy/concurrency.py @@ -12,7 +12,11 @@ import structlog -from s3proxy.crypto import governor_memory_footprint, streaming_governor_clamped_reserve +from s3proxy.crypto import ( + copy_governor_clamped_reserve, + governor_memory_footprint, + streaming_governor_clamped_reserve, +) from s3proxy.errors import S3Error from s3proxy.metrics import MEMORY_LIMIT_BYTES, MEMORY_REJECTIONS, MEMORY_RESERVED_BYTES @@ -80,7 +84,7 @@ def set_memory_limit(self, limit_mb: int) -> None: self._limit_bytes = limit_mb * 1024 * 1024 MEMORY_LIMIT_BYTES.set(self._limit_bytes) - async def try_acquire(self, bytes_needed: int) -> int: + async def try_acquire(self, bytes_needed: int, *, copy: bool = False) -> int: """Reserve memory, waiting up to BACKPRESSURE_TIMEOUT if at capacity.""" if self._limit_bytes <= 0: return 0 @@ -93,14 +97,23 @@ async def try_acquire(self, bytes_needed: int) -> int: # parts are not starved behind one clamped slot. Rare huge uploads may use # more RSS than reserved when run alone; the pod memory limit is the # backstop. + # + # Server-side copies use copy_governor_clamped_reserve instead: their + # honest peak already reflects real chunk work and must not be crushed to + # the routine upload peak (~59MB) or several multi-GB manifest copies run + # concurrently and OOM the pod. if to_reserve > self._limit_bytes: honest = to_reserve - to_reserve = streaming_governor_clamped_reserve(honest, self._limit_bytes) + if copy: + to_reserve = copy_governor_clamped_reserve(honest, self._limit_bytes) + else: + to_reserve = streaming_governor_clamped_reserve(honest, self._limit_bytes) logger.info( "MEMORY_CLAMPED_TO_BUDGET", requested_mb=round(honest / 1024 / 1024, 2), reserved_mb=round(to_reserve / 1024 / 1024, 2), limit_mb=round(self._limit_bytes / 1024 / 1024, 2), + copy=copy, ) async with self._condition: @@ -197,6 +210,10 @@ async def try_acquire_memory(bytes_needed: int) -> int: return await _default.try_acquire(bytes_needed) +async def try_acquire_copy_memory(bytes_needed: int) -> int: + return await _default.try_acquire(bytes_needed, copy=True) + + @contextlib.asynccontextmanager async def reserve_memory(bytes_needed: int): """Reserve memory for the duration of a block, releasing on exit. @@ -213,6 +230,17 @@ async def reserve_memory(bytes_needed: int): await _default.release(reserved) +@contextlib.asynccontextmanager +async def reserve_copy_memory(bytes_needed: int): + """Reserve memory for a server-side copy (upload-style clamp does not apply).""" + reserved = await _default.try_acquire(bytes_needed, copy=True) + try: + yield + finally: + if reserved > 0: + await _default.release(reserved) + + async def release_memory(bytes_reserved: int) -> None: await _default.release(bytes_reserved) diff --git a/s3proxy/crypto.py b/s3proxy/crypto.py index ad51996..5d499b5 100644 --- a/s3proxy/crypto.py +++ b/s3proxy/crypto.py @@ -207,6 +207,17 @@ def streaming_governor_clamped_reserve(honest_peak: int, budget_bytes: int) -> i return min(honest_peak, routine_peak, budget_bytes) +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. + """ + 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. @@ -233,9 +244,16 @@ def copy_pipeline_peak(plaintext_size: int) -> int: adds one pipeline stage the body-fed upload path lacks: the source is read in MAX_BUFFER_SIZE chunks into a _PlaintextReader buffer, so reserve two extra buffers for it. Small copies buffer the whole object and re-encrypt it (~3x). + + Large internal parts (multi-GB manifest copies): aiobotocore copies the + ciphertext body for signing; tracemalloc shows ~part/7 slack at 238MB parts. """ if plaintext_size > STREAMING_THRESHOLD: - return streaming_upload_peak(plaintext_size) + 2 * MAX_BUFFER_SIZE + part = memory_bounded_part_size(plaintext_size) + peak = streaming_upload_peak(plaintext_size) + 2 * MAX_BUFFER_SIZE + if part > 32 * 1024 * 1024: + peak += part // 7 + return peak return max(MAX_BUFFER_SIZE, 3 * plaintext_size) diff --git a/s3proxy/handlers/multipart/copy.py b/s3proxy/handlers/multipart/copy.py index b84afff..161df1b 100644 --- a/s3proxy/handlers/multipart/copy.py +++ b/s3proxy/handlers/multipart/copy.py @@ -62,7 +62,8 @@ async def handle_upload_part_copy(self, request: Request, creds: S3Credentials) # Small copies buffer the whole object + re-encrypt it; gate them # by the limiter too (they carry no body, so the request-level # reservation was ~nothing and a small-object flood ran unbounded). - async with concurrency.reserve_memory(crypto.copy_pipeline_peak(plaintext_size)): + peak = crypto.copy_pipeline_peak(plaintext_size) + async with concurrency.reserve_copy_memory(peak): return await self._simple_copy_part( client, bucket, @@ -221,7 +222,8 @@ async def _streaming_copy_part( # reserved ~nothing -- but this streams the source through decrypt + # re-encrypt. Reserve the pipeline peak so concurrent copies are bounded # (a dedup flood otherwise runs unbounded and OOMs the pod). - async with concurrency.reserve_memory(crypto.copy_pipeline_peak(plaintext_size)): + peak = crypto.copy_pipeline_peak(plaintext_size) + async with concurrency.reserve_copy_memory(peak): return await self._streaming_copy_part_inner( client, bucket, diff --git a/s3proxy/handlers/objects/misc.py b/s3proxy/handlers/objects/misc.py index f47afa0..ab44229 100644 --- a/s3proxy/handlers/objects/misc.py +++ b/s3proxy/handlers/objects/misc.py @@ -386,7 +386,7 @@ async def _copy_encrypted( else: _s = head_resp.get("Metadata", {}).get("plaintext-size") pt_size = int(_s) if _s else crypto.plaintext_size(head_resp.get("ContentLength", 0)) - async with concurrency.reserve_memory(crypto.copy_pipeline_peak(pt_size)): + async with concurrency.reserve_copy_memory(crypto.copy_pipeline_peak(pt_size)): return await self._copy_encrypted_inner( client, bucket, diff --git a/tests/integration/test_copy_memory_governor.py b/tests/integration/test_copy_memory_governor.py new file mode 100644 index 0000000..8ed1c12 --- /dev/null +++ b/tests/integration/test_copy_memory_governor.py @@ -0,0 +1,241 @@ +"""Integration: copy-path memory governor under a real s3proxy subprocess. + +These tests need MinIO on localhost:9000 (same as other e2e tests). +""" + +from __future__ import annotations + +import concurrent.futures +import socket +import uuid + +import boto3 +import pytest +from botocore.exceptions import ClientError + +from tests.integration.conftest import run_s3proxy + +MB = 1024 * 1024 + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +def _client(endpoint: str): + return boto3.client( + "s3", + endpoint_url=endpoint, + aws_access_key_id="minioadmin", + aws_secret_access_key="minioadmin", + region_name="us-east-1", + config=boto3.session.Config( + retries={"max_attempts": 0}, + connect_timeout=10, + read_timeout=300, + ), + ) + + +def _put_large_source(client, bucket: str, key: str, size: int, chunk: int = 8 * MB) -> None: + """Stream a large plaintext object through the proxy (encrypted on store).""" + body = _RepeatingBody(size, chunk) + client.put_object(Bucket=bucket, Key=key, Body=body) + + +class _RepeatingBody: + """File-like stream without holding *size* bytes (boto3 checksum needs tell/seek).""" + + def __init__(self, size: int, chunk: int): + self._size = size + self._chunk = chunk + self._pos = 0 + + def tell(self) -> int: + return self._pos + + def seek(self, offset: int, whence: int = 0) -> int: + if whence == 0: + self._pos = offset + elif whence == 1: + self._pos += offset + elif whence == 2: + self._pos = self._size + offset + else: + raise ValueError(f"invalid whence: {whence}") + self._pos = max(0, min(self._pos, self._size)) + return self._pos + + def read(self, amt=-1): + if self._pos >= self._size: + return b"" + if amt is None or amt < 0: + amt = self._chunk + take = min(amt, self._chunk, self._size - self._pos) + self._pos += take + return b"x" * take + + +def _upload_part_copy(client, bucket: str, dest_key: str, source_key: str) -> dict: + resp = client.create_multipart_upload(Bucket=bucket, Key=dest_key) + upload_id = resp["UploadId"] + try: + copy_resp = client.upload_part_copy( + Bucket=bucket, + Key=dest_key, + PartNumber=1, + UploadId=upload_id, + CopySource={"Bucket": bucket, "Key": source_key}, + ) + client.complete_multipart_upload( + Bucket=bucket, + Key=dest_key, + UploadId=upload_id, + MultipartUpload={ + "Parts": [{"PartNumber": 1, "ETag": copy_resp["CopyPartResult"]["ETag"]}] + }, + ) + return {"success": True} + except ClientError as e: + code = e.response.get("Error", {}).get("Code", "") + return {"success": False, "code": code, "error": str(e)} + except Exception as e: + return {"success": False, "code": "error", "error": str(e)} + + +def _upload_part(client, bucket: str, key: str, part_number: int, upload_id: str, size: int): + body = _RepeatingBody(size, 8 * MB) + try: + resp = client.upload_part( + Bucket=bucket, + Key=key, + PartNumber=part_number, + UploadId=upload_id, + Body=body, + ) + return {"success": True, "etag": resp["ETag"]} + except ClientError as e: + code = e.response.get("Error", {}).get("Code", "") + return {"success": False, "code": code, "error": str(e)} + + +# 64MB source: copy_pipeline_peak ~48MB. Governor must exceed the clamped copy +# slot plus the POST gate's MIN_RESERVATION (64KB); at 32MB the clamp monopolizes +# the budget and even a single copy fails (64KB+32MB > 32MB while the gate is held). +# 96MB fits one ~48MB copy plus concurrent ~20MB upload reservations, and still +# rejects a second concurrent copy (48+48+gate > 96). +SOURCE_SIZE = 64 * MB +GOVERNOR_MB = "96" +SOURCE_KEY = "shared-large-source.bin" + + +@pytest.fixture(scope="module") +def copy_stress_server(): + port = _find_free_port() + with run_s3proxy( + port, + log_output=False, + S3PROXY_MEMORY_LIMIT_MB=GOVERNOR_MB, + # 0 = reject immediately (SlowDown), don't queue — otherwise all 3 copies + # succeed serially within the 2s wait and the concurrency assertion fails. + S3PROXY_BACKPRESSURE_TIMEOUT="0", + S3PROXY_MAX_PART_SIZE_MB="0", + ) as (endpoint, proc): + yield endpoint, proc + + +@pytest.fixture(scope="module") +def copy_bucket(copy_stress_server): + endpoint, _ = copy_stress_server + client = _client(endpoint) + bucket = f"copy-gov-{uuid.uuid4().hex[:8]}" + try: + client.create_bucket(Bucket=bucket) + except ClientError as exc: + pytest.skip(f"MinIO/S3 backend not available for e2e: {exc}") + yield client, bucket + try: + resp = client.list_objects_v2(Bucket=bucket) + if "Contents" in resp: + client.delete_objects( + Bucket=bucket, + Delete={"Objects": [{"Key": o["Key"]} for o in resp["Contents"]]}, + ) + client.delete_bucket(Bucket=bucket) + except Exception: + pass + + +@pytest.fixture(scope="module") +def large_source(copy_bucket): + """Upload once for all copy-governor tests (avoids 3× redundant 256MB puts).""" + client, bucket = copy_bucket + _put_large_source(client, bucket, SOURCE_KEY, SOURCE_SIZE) + return SOURCE_KEY + + +@pytest.mark.e2e +class TestCopyMemoryGovernorSubprocess: + """Real subprocess + boto3: copies must not OOM the server.""" + + def test_three_concurrent_large_copies_do_not_oom_server( + self, copy_stress_server, copy_bucket, large_source + ): + """Prod regression: 3 concurrent large copies must not all run + OOM.""" + endpoint, proc = copy_stress_server + client, bucket = copy_bucket + source_key = large_source + + def one_copy(i: int) -> dict: + return _upload_part_copy(client, bucket, f"dest-{i}.bin", source_key) + + with concurrent.futures.ThreadPoolExecutor(max_workers=3) as pool: + results = list(pool.map(one_copy, range(3))) + + assert proc.poll() is None, "s3proxy OOMKilled or crashed (exit 137?)" + succeeded = sum(1 for r in results if r["success"]) + slowed = sum(1 for r in results if r.get("code") == "SlowDown") + assert succeeded >= 1, f"expected at least one copy to succeed: {results}" + # 96MB budget fits two ~48MB copy slots; the third must be rejected, not queued. + assert slowed >= 1, f"expected backpressure on concurrent copies: {results}" + assert succeeded <= 2, f"more than two copies ran concurrently: {results}" + assert succeeded + slowed == len(results), f"unexpected errors: {results}" + + def test_mixed_scylla_uploads_and_large_copy_server_survives( + self, copy_stress_server, copy_bucket, large_source + ): + """Prod mix: concurrent small parts + large manifest-style copy.""" + endpoint, proc = copy_stress_server + client, bucket = copy_bucket + source_key = large_source + + def scylla_part(i: int) -> dict: + key = f"sst-part-{i}.bin" + resp = client.create_multipart_upload(Bucket=bucket, Key=key) + upload_id = resp["UploadId"] + # 5MB parts reserve ~20MB (not 32MB like 16MB parts) so two can run + # alongside one ~48MB copy within the 96MB mixed-workload governor. + r = _upload_part(client, bucket, key, 1, upload_id, 5 * MB) + if r["success"]: + client.complete_multipart_upload( + Bucket=bucket, + Key=key, + UploadId=upload_id, + MultipartUpload={"Parts": [{"PartNumber": 1, "ETag": r["etag"]}]}, + ) + return r + + def manifest_copy() -> dict: + return _upload_part_copy(client, bucket, "manifest-copy.bin", source_key) + + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: + upload_futs = [pool.submit(scylla_part, i) for i in range(3)] + copy_fut = pool.submit(manifest_copy) + copy_result = copy_fut.result() + upload_results = [f.result() for f in upload_futs] + + assert proc.poll() is None, "server crashed under mixed workload" + assert sum(1 for r in upload_results if r["success"]) >= 2 + assert copy_result["success"] or copy_result.get("code") == "SlowDown" diff --git a/tests/unit/test_copy_clamp_and_concurrency.py b/tests/unit/test_copy_clamp_and_concurrency.py new file mode 100644 index 0000000..4f066af --- /dev/null +++ b/tests/unit/test_copy_clamp_and_concurrency.py @@ -0,0 +1,157 @@ +"""Regression tests for copy-path governor clamp and concurrent manifest copies. + +Prod OOM: three concurrent 4.7GB Scylla .sm_ manifest UploadPartCopy ops on +one pod each reserved ~59MB (upload-style routine clamp) while holding ~175MB +ciphertext at copy.py:420 — governor active_mb ~178, RSS ~790, OOMKilled. +""" + +from __future__ import annotations + +import asyncio + +import pytest + +from s3proxy import concurrency, crypto +from s3proxy.errors import S3Error +from tests.unit.test_copy_reservation_vs_real import _measure_peak + +MB = 1024 * 1024 + +# Observed prod manifest copy size. +SCYLLA_MANIFEST_BYTES = 4_767 * MB + + +def test_copy_clamp_monopolizes_budget_not_routine_upload_peak(): + """4.7GB copy honest peak ~500MB at 192MB budget must reserve 192, not ~59.""" + honest = crypto.copy_pipeline_peak(SCYLLA_MANIFEST_BYTES) + budget = 192 * MB + assert honest > budget + routine = crypto.streaming_upload_peak(crypto.STREAMING_GOVERNOR_CLIENT_PART_BYTES) + assert crypto.streaming_governor_clamped_reserve(honest, budget) == routine + assert crypto.copy_governor_clamped_reserve(honest, budget) == budget + + +@pytest.mark.asyncio +async def test_three_manifest_copies_do_not_run_concurrently_at_192mb(): + """Prod failure mode: 3× 4.7GB copies must not all hold slots on one pod.""" + original_timeout = concurrency.BACKPRESSURE_TIMEOUT + concurrency.BACKPRESSURE_TIMEOUT = 0 + concurrency.reset_state() + concurrency.set_memory_limit(192) + per_copy = crypto.copy_pipeline_peak(SCYLLA_MANIFEST_BYTES) + assert per_copy > 192 * MB + + max_inside = 0 + inside = 0 + admitted = 0 + rejected = 0 + + async def one_copy(): + nonlocal max_inside, inside, admitted, rejected + try: + async with concurrency.reserve_copy_memory(per_copy): + inside += 1 + max_inside = max(max_inside, inside) + admitted += 1 + await asyncio.sleep(0.05) + inside -= 1 + except S3Error: + rejected += 1 + + try: + await asyncio.gather(*[one_copy() for _ in range(3)]) + finally: + concurrency.BACKPRESSURE_TIMEOUT = original_timeout + concurrency.reset_state() + + assert max_inside == 1, f"expected exclusive large copy, saw {max_inside} concurrent" + assert admitted == 1 + assert rejected == 2 + + +@pytest.mark.asyncio +async def test_manifest_copy_plus_scylla_upload_serialized_at_192mb(): + """Active ~25MB Scylla upload must block a budget-monopolizing manifest copy.""" + original_timeout = concurrency.BACKPRESSURE_TIMEOUT + concurrency.BACKPRESSURE_TIMEOUT = 0 + concurrency.reset_state() + concurrency.set_memory_limit(192) + try: + scylla = crypto.governor_memory_footprint(50 * MB) + await concurrency.try_acquire_memory(scylla) + + with pytest.raises(S3Error): + await concurrency.try_acquire_copy_memory( + crypto.copy_pipeline_peak(SCYLLA_MANIFEST_BYTES) + ) + finally: + concurrency.BACKPRESSURE_TIMEOUT = original_timeout + concurrency.reset_state() + + +@pytest.mark.parametrize("mb", [512, 1024, 4767]) +@pytest.mark.asyncio +async def test_copy_reservation_bounds_real_peak_large_manifests(mb: int): + """copy_pipeline_peak must cover tracemalloc peak for prod-sized copies.""" + plaintext_size = mb * MB + real_peak = await _measure_peak(plaintext_size) + reserved = crypto.copy_pipeline_peak(plaintext_size) + assert reserved >= real_peak, ( + f"{mb}MB copy: reserved {reserved / MB:.1f}MB < real {real_peak / MB:.1f}MB" + ) + + +@pytest.mark.asyncio +async def test_old_upload_clamp_would_admit_three_concurrent_manifest_copies(): + """Regression guard: upload-style clamp is what caused prod OOM.""" + original_timeout = concurrency.BACKPRESSURE_TIMEOUT + concurrency.BACKPRESSURE_TIMEOUT = 0 + concurrency.reset_state() + concurrency.set_memory_limit(192) + per_copy = crypto.copy_pipeline_peak(SCYLLA_MANIFEST_BYTES) + try: + + async def acquire_upload_clamp(): + return await concurrency.try_acquire_memory(per_copy) + + results = await asyncio.gather(*[acquire_upload_clamp() for _ in range(3)]) + assert all(r < 80 * MB for r in results) + assert sum(results) < 192 * MB + assert len(results) == 3 + finally: + concurrency.BACKPRESSURE_TIMEOUT = original_timeout + concurrency.reset_state() + + +@pytest.mark.asyncio +async def test_mixed_three_scylla_uploads_and_manifest_copy_limits_concurrency(): + """Prod replay: 3×50MB uploads + manifest copy — copy cannot take 4th slot.""" + original_timeout = concurrency.BACKPRESSURE_TIMEOUT + concurrency.BACKPRESSURE_TIMEOUT = 0 + concurrency.reset_state() + concurrency.set_memory_limit(192) + per_part = crypto.governor_memory_footprint(50 * MB) + manifest = crypto.copy_pipeline_peak(SCYLLA_MANIFEST_BYTES) + try: + uploads = await asyncio.gather( + *[concurrency.try_acquire_memory(per_part) for _ in range(3)] + ) + assert sum(uploads) + 192 * MB > 192 * MB + with pytest.raises(S3Error): + await concurrency.try_acquire_copy_memory(manifest) + finally: + concurrency.BACKPRESSURE_TIMEOUT = original_timeout + concurrency.reset_state() + + +@pytest.mark.asyncio +async def test_copy_acquire_uses_copy_clamp_not_upload_clamp(): + concurrency.reset_state() + concurrency.set_memory_limit(192) + try: + honest = crypto.copy_pipeline_peak(SCYLLA_MANIFEST_BYTES) + reserved = await concurrency.try_acquire_copy_memory(honest) + assert reserved == 192 * MB + assert reserved > crypto.streaming_governor_clamped_reserve(honest, 192 * MB) + finally: + concurrency.reset_state() diff --git a/tests/unit/test_copy_memory_governing.py b/tests/unit/test_copy_memory_governing.py index 82f7a06..a2a28c1 100644 --- a/tests/unit/test_copy_memory_governing.py +++ b/tests/unit/test_copy_memory_governing.py @@ -27,11 +27,11 @@ def test_copy_pipeline_peak_matches_framed_upload(): # size-independent 4*MAX_BUFFER_SIZE under-counted the real peak ~6x -- when # it un-framed a whole part -- and OOMed the pod. for size in (64 * MB, 512 * MB, 5 * 1024 * MB): - # Same framed writer peak, plus two MAX_BUFFER_SIZE buffers for the copy - # source read pipeline the body-fed upload path doesn't have. - assert crypto.copy_pipeline_peak(size) == ( - crypto.streaming_upload_peak(size) + 2 * crypto.MAX_BUFFER_SIZE - ) + part = crypto.memory_bounded_part_size(size) + base = crypto.streaming_upload_peak(size) + 2 * crypto.MAX_BUFFER_SIZE + if part > 32 * MB: + base += part // 7 + assert crypto.copy_pipeline_peak(size) == base def test_copy_pipeline_peak_small_is_three_x(): @@ -63,7 +63,7 @@ async def test_reserve_memory_bounds_concurrent_copies(): async def one_copy(): nonlocal peak_active, inside, max_inside - async with concurrency.reserve_memory(per_copy): + async with concurrency.reserve_copy_memory(per_copy): inside += 1 max_inside = max(max_inside, inside) peak_active = max(peak_active, limiter.active_bytes) diff --git a/tests/unit/test_copy_object_memory_governing.py b/tests/unit/test_copy_object_memory_governing.py new file mode 100644 index 0000000..24731d1 --- /dev/null +++ b/tests/unit/test_copy_object_memory_governing.py @@ -0,0 +1,71 @@ +"""CopyObject must use reserve_copy_memory (copy clamp), not upload clamp.""" + +from __future__ import annotations + +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from s3proxy import concurrency, crypto +from s3proxy.handlers.objects.misc import MiscObjectMixin + +MB = 1024 * 1024 + + +class _Mgr: + async def create_upload(self, *a, **k): + return None + + async def get_upload(self, *a, **k): + return None + + +def _handler(): + h = MiscObjectMixin.__new__(MiscObjectMixin) + h.multipart_manager = _Mgr() + h.keyring = MagicMock() + h.keyring.key_for.return_value = ("kid", b"k" * 32) + h.settings = MagicMock() + h.settings.dektag_name = "x-amz-meta-dek" + h.settings.kidtag_name = "x-amz-meta-kid" + return h + + +@pytest.mark.asyncio +async def test_copy_encrypted_streaming_uses_reserve_copy_memory(): + """Large CopyObject must not go through upload-style routine clamp.""" + h = _handler() + pt_size = 512 * MB + acquired: list[int] = [] + + @asynccontextmanager + async def spy_reserve(needed: int): + acquired.append(needed) + yield + + meta = type("M", (), {"total_plaintext_size": pt_size})() + head = {"Metadata": {}, "ContentLength": pt_size, "ContentType": "application/octet-stream"} + + with ( + patch.object(concurrency, "reserve_copy_memory", side_effect=spy_reserve), + patch.object(h, "_copy_encrypted_inner", new_callable=AsyncMock) as mock_inner, + ): + mock_inner.return_value = MagicMock(status_code=200) + await h._copy_encrypted( + MagicMock(), + "b", + "dst", + None, + "b", + "src", + head, + None, + meta, + "COPY", + None, + ) + + expected = crypto.copy_pipeline_peak(pt_size) + assert acquired == [expected] + assert expected > crypto.governor_memory_footprint(pt_size) diff --git a/tests/unit/test_governor_peak_and_clamp.py b/tests/unit/test_governor_peak_and_clamp.py index afae313..1c69612 100644 --- a/tests/unit/test_governor_peak_and_clamp.py +++ b/tests/unit/test_governor_peak_and_clamp.py @@ -378,14 +378,15 @@ async def one(): reset_state() -def test_honest_peak_still_available_for_copy_pipeline(): - """copy_pipeline_peak uses honest streaming_upload_peak (not governor cap). +def test_honest_copy_peak_exceeds_upload_governor_cap(): + """copy_pipeline_peak uses per-chunk honest peak, not governor_memory_footprint. - Copies reserve separately in the handler; the upload gate cap must not - shrink copy_pipeline_peak or copies would under-reserve and OOM. + Upload gate caps multi-GB Content-Length at the routine workload peak; copies + reserve separately via reserve_copy_memory with copy_governor_clamped_reserve. """ huge = 5 * 1024 * MB assert crypto.copy_pipeline_peak(huge) > estimate_memory_footprint("PUT", huge) + part = crypto.memory_bounded_part_size(huge) assert crypto.copy_pipeline_peak(huge) == ( - crypto.streaming_upload_peak(huge) + 2 * crypto.MAX_BUFFER_SIZE + crypto.streaming_upload_peak(huge) + 2 * crypto.MAX_BUFFER_SIZE + part // 7 )