fix: writer preference so budget-monopolizing copies are not starved

s3proxy/concurrency.py

@@ -61,10 +61,28 @@ def __init__(self, limit_mb: int = 128) -> None:
         self._limit_mb = limit_mb
         self._limit_bytes = limit_mb * 1024 * 1024
         self._active_bytes = 0
+        self._pending_exclusive = 0
         self._lock = asyncio.Lock()
         self._condition = asyncio.Condition(self._lock)
         MEMORY_LIMIT_BYTES.set(self._limit_bytes)
 
+    def _can_admit(self, to_reserve: int, exclusive: bool) -> bool:
+        """Whether a request may reserve right now.
+
+        A request that needs the whole budget (exclusive -- e.g. a multi-GB copy
+        clamped to limit_bytes) is admissible only when nothing else holds
+        memory: active_bytes + limit_bytes > limit_bytes for any active_bytes > 0.
+        Without fairness the limiter never reaches that state under load -- every
+        release wakes all waiters and a small request re-grabs memory before the
+        copy can, so active_bytes never drains to 0 and the copy backpressures
+        until it 503s. Writer preference fixes it: while an exclusive request is
+        waiting, hold back new non-exclusive admissions so active_bytes drains
+        and the copy gets its exclusive slot (bounded by in-flight requests).
+        """
+        if self._active_bytes + to_reserve > self._limit_bytes:
+            return False
+        return exclusive or self._pending_exclusive == 0
+
     @property
     def limit_bytes(self) -> int:
         return self._limit_bytes
@@ -116,38 +134,49 @@ async def try_acquire(self, bytes_needed: int, *, copy: bool = False) -> int:
                 copy=copy,
             )
 
+        exclusive = to_reserve >= self._limit_bytes
         async with self._condition:
             deadline = asyncio.get_event_loop().time() + BACKPRESSURE_TIMEOUT
-            while self._active_bytes + to_reserve > self._limit_bytes:
-                remaining = deadline - asyncio.get_event_loop().time()
-                if remaining <= 0:
-                    active_mb = self._active_bytes / 1024 / 1024
-                    request_mb = to_reserve / 1024 / 1024
-                    limit_mb = self._limit_bytes / 1024 / 1024
-                    logger.warning(
-                        "MEMORY_REJECTED",
-                        active_mb=round(active_mb, 2),
-                        requested_mb=round(request_mb, 2),
-                        limit_mb=round(limit_mb, 2),
-                        waited_sec=BACKPRESSURE_TIMEOUT,
-                    )
-                    MEMORY_REJECTIONS.inc()
-                    raise S3Error.slow_down(
-                        f"Memory limit: {active_mb:.0f}MB + {request_mb:.0f}MB > {limit_mb:.0f}MB"
+            if exclusive:
+                self._pending_exclusive += 1
+            try:
+                while not self._can_admit(to_reserve, exclusive):
+                    remaining = deadline - asyncio.get_event_loop().time()
+                    if remaining <= 0:
+                        active_mb = self._active_bytes / 1024 / 1024
+                        request_mb = to_reserve / 1024 / 1024
+                        limit_mb = self._limit_bytes / 1024 / 1024
+                        logger.warning(
+                            "MEMORY_REJECTED",
+                            active_mb=round(active_mb, 2),
+                            requested_mb=round(request_mb, 2),
+                            limit_mb=round(limit_mb, 2),
+                            waited_sec=BACKPRESSURE_TIMEOUT,
+                        )
+                        MEMORY_REJECTIONS.inc()
+                        raise S3Error.slow_down(
+                            f"Memory limit: {active_mb:.0f}MB + "
+                            f"{request_mb:.0f}MB > {limit_mb:.0f}MB"
+                        )
+                    logger.info(
+                        "MEMORY_BACKPRESSURE",
+                        active_mb=round(self._active_bytes / 1024 / 1024, 2),
+                        requested_mb=round(to_reserve / 1024 / 1024, 2),
+                        limit_mb=round(self._limit_bytes / 1024 / 1024, 2),
+                        remaining_sec=round(remaining, 1),
                     )
-                logger.info(
-                    "MEMORY_BACKPRESSURE",
-                    active_mb=round(self._active_bytes / 1024 / 1024, 2),
-                    requested_mb=round(to_reserve / 1024 / 1024, 2),
-                    limit_mb=round(self._limit_bytes / 1024 / 1024, 2),
-                    remaining_sec=round(remaining, 1),
-                )
-                with contextlib.suppress(TimeoutError):
-                    await asyncio.wait_for(self._condition.wait(), timeout=remaining)
-
-            self._active_bytes += to_reserve
-            MEMORY_RESERVED_BYTES.set(self._active_bytes)
-            return to_reserve
+                    with contextlib.suppress(TimeoutError):
+                        await asyncio.wait_for(self._condition.wait(), timeout=remaining)
+
+                self._active_bytes += to_reserve
+                MEMORY_RESERVED_BYTES.set(self._active_bytes)
+                return to_reserve
+            finally:
+                if exclusive:
+                    self._pending_exclusive -= 1
+                    # Whether we got the slot or gave up, release the non-exclusive
+                    # requests we were holding back.
+                    self._condition.notify_all()
 
     async def release(self, bytes_reserved: int) -> None:
         """Release reserved memory and wake waiting requests."""

tests/unit/test_copy_clamp_and_concurrency.py

@@ -12,6 +12,7 @@
 import pytest
 
 from s3proxy import concurrency, crypto
+from s3proxy.concurrency import MIN_RESERVATION
 from s3proxy.errors import S3Error
 from tests.unit.test_copy_reservation_vs_real import _measure_peak
 
@@ -144,6 +145,53 @@ async def test_mixed_three_scylla_uploads_and_manifest_copy_limits_concurrency()
         concurrency.reset_state()
 
 
+@pytest.mark.asyncio
+async def test_pending_exclusive_copy_not_starved_by_small_requests():
+    """Writer preference: a budget-monopolizing copy must not be starved by a
+    stream of small requests.
+
+    The copy reserves the whole budget, so it can only be admitted when
+    active_bytes == 0. A small request that easily fits must NOT jump the queue
+    while the copy is waiting -- otherwise active_bytes never drains and the copy
+    backpressures forever (prod: requested_mb == limit_mb that never clears, then
+    503). It must instead queue behind the pending copy.
+    """
+    original_timeout = concurrency.BACKPRESSURE_TIMEOUT
+    concurrency.BACKPRESSURE_TIMEOUT = 5
+    concurrency.reset_state()
+    concurrency.set_memory_limit(192)
+    try:
+        # An in-flight ~25MB upload holds memory, so the copy cannot start yet.
+        blocker = await concurrency.try_acquire_memory(crypto.governor_memory_footprint(50 * MB))
+        assert 0 < blocker < 192 * MB
+
+        copy_task = asyncio.create_task(
+            concurrency.try_acquire_copy_memory(crypto.copy_pipeline_peak(SCYLLA_MANIFEST_BYTES))
+        )
+        # A tiny request that trivially fits alongside the blocker (25 + 0.06 <<
+        # 192) but must be held back because an exclusive copy is pending.
+        small_task = asyncio.create_task(concurrency.try_acquire_memory(MIN_RESERVATION))
+        await asyncio.sleep(0.05)
+
+        assert not copy_task.done(), "copy must wait for exclusive access"
+        assert not small_task.done(), "small request must queue behind the pending copy"
+
+        # Drain the blocker: active_bytes -> 0 -> the copy wins the slot, not the
+        # queued small request.
+        await concurrency.release_memory(blocker)
+        reserved = await asyncio.wait_for(copy_task, timeout=2)
+        assert reserved == 192 * MB
+        assert concurrency.get_active_memory() == 192 * MB
+        assert not small_task.done(), "copy holds the whole budget; small still waits"
+
+        # Copy releases -> the small request finally proceeds.
+        await concurrency.release_memory(reserved)
+        assert await asyncio.wait_for(small_task, timeout=2) == MIN_RESERVATION
+    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()