diff --git a/s3proxy/concurrency.py b/s3proxy/concurrency.py index 103ebc0..bda4f3b 100644 --- a/s3proxy/concurrency.py +++ b/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.""" diff --git a/tests/unit/test_copy_clamp_and_concurrency.py b/tests/unit/test_copy_clamp_and_concurrency.py index 4f066af..e1d8e10 100644 --- a/tests/unit/test_copy_clamp_and_concurrency.py +++ b/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()