From 0789747aeceee4b06e239f4d9e8eb5bceb160857 Mon Sep 17 00:00:00 2001 From: Justin Joseph Date: Thu, 30 Apr 2026 14:51:01 -0700 Subject: [PATCH] Add ResizableSemaphore for adaptive pipeline concurrency (#1389) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Add `ResizableSemaphore` — an asyncio-compatible semaphore whose max permit count can be adjusted at runtime. This is the foundational primitive for dynamic concurrency control in SPDL pipelines. Key features: - `resize(new_max)` adjusts permits at runtime; resize-up wakes blocked waiters immediately, resize-down drains gracefully (no preemption) - `acquire()`/`release()` semantics match `asyncio.Semaphore` - `max_value` and `active` properties for observability - Thread-safe within asyncio's single-threaded model - Comprehensive error handling for invalid values This is Diff 1 of a 5-diff series implementing a unified adaptive scheduler for SPDL pipelines (T262755626). Differential Revision: D99920401 --- src/spdl/pipeline/_components/__init__.py | 2 + src/spdl/pipeline/_components/_semaphore.py | 189 +++++ tests/pipeline/resizable_semaphore_test.py | 754 ++++++++++++++++++++ 3 files changed, 945 insertions(+) create mode 100644 src/spdl/pipeline/_components/_semaphore.py create mode 100644 tests/pipeline/resizable_semaphore_test.py diff --git a/src/spdl/pipeline/_components/__init__.py b/src/spdl/pipeline/_components/__init__.py index 437020fc0..0cbe6ced1 100644 --- a/src/spdl/pipeline/_components/__init__.py +++ b/src/spdl/pipeline/_components/__init__.py @@ -20,6 +20,7 @@ set_default_queue_class, StatsQueue, ) +from ._semaphore import ResizableSemaphore __all__ = [ "_build_pipeline_coro", @@ -30,6 +31,7 @@ "is_eof", "is_epoch_end", "PipelineFailure", + "ResizableSemaphore", "set_default_hook_class", "set_default_queue_class", "TaskHook", diff --git a/src/spdl/pipeline/_components/_semaphore.py b/src/spdl/pipeline/_components/_semaphore.py new file mode 100644 index 000000000..c19ba9f18 --- /dev/null +++ b/src/spdl/pipeline/_components/_semaphore.py @@ -0,0 +1,189 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""Resizable asyncio semaphore for dynamic concurrency control.""" + +__all__ = ["ResizableSemaphore"] + +import asyncio +from collections import deque + + +class ResizableSemaphore: + """asyncio.Semaphore variant whose max value can be changed at runtime. + + Thread safety: asyncio is single-threaded per event loop. All methods + MUST be called from coroutines in the same event loop. No locks needed. + + Resize semantics: + - Increase: immediately wake up to ``new_max - old_max`` blocked + waiters. + - Decrease: no preemption. Currently acquired permits continue. + New ``acquire()`` calls block until active count drops below + the new max. ``_current_value`` may go negative during drain. + + Invariant: at any moment, the number of "active" (acquired but not + yet released) permits equals ``max_value - _current_value``. When + ``_current_value`` is negative, more permits are outstanding than + the current max allows -- they drain naturally as tasks + ``release()``. + """ + + def __init__(self, value: int) -> None: + """Create a semaphore with *value* initial permits. + + Args: + value: Initial max permits. Must be >= 1. + + Raises: + ValueError: If value < 1. + """ + if value < 1: + raise ValueError(f"value must be >= 1, got {value}") + self._max_value: int = value + self._current_value: int = value + self._waiters: deque[asyncio.Future[None]] = deque() + + @property + def max_value(self) -> int: + """Current max permits (may differ from initial after resize).""" + return self._max_value + + @property + def active(self) -> int: + """Number of currently acquired (outstanding) permits. + + Can exceed ``max_value`` temporarily after a resize-down. + """ + return self._max_value - self._current_value + + async def acquire(self) -> None: + """Acquire one permit. Blocks if no permits available. + + Raises: + asyncio.CancelledError: If the waiting coroutine is + cancelled while blocked. + """ + # Fast path: permit available and no one queued ahead of us. + if self._current_value > 0 and not self._waiters: + self._current_value -= 1 + return + + fut: asyncio.Future[None] = asyncio.get_running_loop().create_future() + self._waiters.append(fut) + try: + await fut + except asyncio.CancelledError: + # PERMIT-LEAK FIX (V5.1): + # Three states are possible at this point: + # (a) fut not done: nobody granted us a permit yet. Just + # remove from the queue. + # (b) fut done with result: release()/resize() handed us a + # permit (direct hand-off — no _current_value increment + # happened). We are about to NOT enter the critical + # section, so we MUST give the permit back. Call + # release() to wake the next waiter (or restore the + # permit to the pool if no waiters remain). + # (c) fut already cancelled before we entered the await: + # same shape as (a) — `fut in self._waiters` is True + # and `self._waiters.remove(fut)` covers it. No permit + # was granted, so nothing to give back. + if fut in self._waiters: + # Case (a) or (c): pre-grant cancellation. No permit + # was handed off, so nothing to release. + self._waiters.remove(fut) + elif fut.done() and not fut.cancelled() and fut.exception() is None: + # Case (b): post-grant cancellation. release() handed us + # a permit via set_result(None) but we're not going to + # use it. Hand it back so the next waiter (or the pool) + # gets it. + self.release() + raise + # Granted via direct hand-off from release()/resize(). + # _current_value was NOT decremented (the permit transferred in + # flight from the previous holder), so we are already accounted + # for as "active". + + def release(self) -> None: + """Return a permit. Wakes one blocked waiter if any. + + Direct hand-off semantics: when waiters are queued, the permit + transfers from the releaser to the next non-cancelled waiter + without round-tripping through ``_current_value``. This avoids + a window where two concurrent callers could observe + ``_current_value > 0`` between waiter-pop and decrement. + + After a resize-down, when no waiters remain, released permits + that would push ``_current_value`` above ``max_value`` are + absorbed (clamped). This is correct: the permit belonged to + the old, larger max. + """ + # Try to hand the permit directly to the next non-cancelled + # waiter. The skip-loop drains cancelled waiters whose futures + # are already done() — protecting against the release/cancel + # race where a waiter is cancelled mid-iteration. + while self._waiters: + # pyre-ignore[1001]: Future is granted via set_result(), not awaited. + waiter = self._waiters.popleft() + if not waiter.done(): + # Permit transfers atomically: stays "in flight" with + # the new owner. Do NOT touch _current_value. + waiter.set_result(None) + return + # No live waiters; restore one permit to the pool (clamped to + # the current max so resize-down clamps don't drift over). + self._current_value = min(self._current_value + 1, self._max_value) + + def resize(self, new_max: int) -> None: + """Change the maximum number of permits. + + Args: + new_max: New maximum. Must be >= 1. + + When increasing (``new_max > old_max``): + Additional permits become immediately available. Blocked + waiters are woken (via direct hand-off) to fill the new + capacity. Any leftover permits go to the pool, clamped to + the new max. + + When decreasing (``new_max < old_max``): + No preemption -- currently active tasks continue. + ``_current_value`` is reduced by the delta, which may make + it negative. Future ``acquire()`` calls block until enough + releases bring ``_current_value`` back above 0. + + Raises: + ValueError: If new_max < 1. + """ + if new_max < 1: + raise ValueError(f"new_max must be >= 1, got {new_max}") + delta = new_max - self._max_value + self._max_value = new_max + if delta > 0: + # Increase: hand `delta` permits directly to waiters first. + # Skip cancelled waiters (their futures are already done()). + granted = 0 + while self._waiters and granted < delta: + # pyre-ignore[1001]: Future is granted via set_result(), not awaited. + waiter = self._waiters.popleft() + if not waiter.done(): + # Direct hand-off: permit transfers in flight; do + # NOT touch _current_value here. + waiter.set_result(None) + granted += 1 + # Any permits not handed off go to the pool, clamped to max. + leftover = delta - granted + if leftover > 0: + self._current_value = min( + self._current_value + leftover, + self._max_value, + ) + elif delta < 0: + # Decrease: subtract from available pool (may go negative). + # delta is already negative. + self._current_value += delta diff --git a/tests/pipeline/resizable_semaphore_test.py b/tests/pipeline/resizable_semaphore_test.py new file mode 100644 index 000000000..9a3ae7192 --- /dev/null +++ b/tests/pipeline/resizable_semaphore_test.py @@ -0,0 +1,754 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import asyncio +import unittest + +from spdl.pipeline._components._semaphore import ResizableSemaphore + + +class ResizableSemaphoreConstructionTest(unittest.TestCase): + def test_init_valid(self) -> None: + sem = ResizableSemaphore(5) + self.assertEqual(sem.max_value, 5) + self.assertEqual(sem.active, 0) + + def test_init_one(self) -> None: + sem = ResizableSemaphore(1) + self.assertEqual(sem.max_value, 1) + self.assertEqual(sem.active, 0) + + def test_init_zero_raises(self) -> None: + with self.assertRaises(ValueError): + ResizableSemaphore(0) + + def test_init_negative_raises(self) -> None: + with self.assertRaises(ValueError): + ResizableSemaphore(-1) + + +class ResizableSemaphoreAcquireReleaseTest(unittest.TestCase): + def test_acquire_decrements_permits(self) -> None: + async def run() -> None: + sem = ResizableSemaphore(3) + await sem.acquire() + self.assertEqual(sem.active, 1) + await sem.acquire() + self.assertEqual(sem.active, 2) + await sem.acquire() + self.assertEqual(sem.active, 3) + + asyncio.run(run()) + + def test_release_increments_permits(self) -> None: + async def run() -> None: + sem = ResizableSemaphore(3) + await sem.acquire() + await sem.acquire() + self.assertEqual(sem.active, 2) + sem.release() + self.assertEqual(sem.active, 1) + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + def test_acquire_blocks_when_exhausted(self) -> None: + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + acquired: asyncio.Event = asyncio.Event() + + async def try_acquire() -> None: + await sem.acquire() + acquired.set() + + task = asyncio.create_task(try_acquire()) + # Yield to let the task enter acquire() and block. + await asyncio.sleep(0) + self.assertFalse(acquired.is_set()) + + # Release unblocks the waiter. + sem.release() + await asyncio.sleep(0) + self.assertTrue(acquired.is_set()) + + # Clean up. + sem.release() + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(run()) + + def test_release_without_acquire_clamps(self) -> None: + """Release without prior acquire should not exceed max_value.""" + + async def run() -> None: + sem = ResizableSemaphore(3) + # active is 0, current_value == max_value == 3 + sem.release() + # Should clamp: active stays 0, not -1. + self.assertEqual(sem.active, 0) + self.assertEqual(sem.max_value, 3) + + asyncio.run(run()) + + def test_fifo_wake_order(self) -> None: + """Waiters should be woken in FIFO order.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + order: list[int] = [] + + async def waiter(idx: int) -> None: + await sem.acquire() + order.append(idx) + sem.release() + + t1 = asyncio.create_task(waiter(1)) + await asyncio.sleep(0) + t2 = asyncio.create_task(waiter(2)) + await asyncio.sleep(0) + t3 = asyncio.create_task(waiter(3)) + await asyncio.sleep(0) + + # Release the initial acquire — should wake waiter 1 first. + sem.release() + + # Wait for all waiters to complete. + await asyncio.wait_for(asyncio.gather(t1, t2, t3), timeout=5.0) + self.assertEqual(order, [1, 2, 3]) + + asyncio.run(run()) + + +class ResizableSemaphoreResizeUpTest(unittest.TestCase): + def test_resize_up_increases_max(self) -> None: + async def run() -> None: + sem = ResizableSemaphore(2) + sem.resize(5) + self.assertEqual(sem.max_value, 5) + + asyncio.run(run()) + + def test_resize_up_wakes_waiters(self) -> None: + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + acquired_events: list[asyncio.Event] = [] + + async def waiter() -> None: + await sem.acquire() + evt = acquired_events[len(acquired_events)] + evt.set() + + e1 = asyncio.Event() + e2 = asyncio.Event() + acquired_events.extend([e1, e2]) + + # Simpler: track via counter + acquired_count = 0 + + async def counting_waiter() -> None: + nonlocal acquired_count + await sem.acquire() + acquired_count += 1 + + t1 = asyncio.create_task(counting_waiter()) + await asyncio.sleep(0) + t2 = asyncio.create_task(counting_waiter()) + await asyncio.sleep(0) + self.assertEqual(acquired_count, 0) + + # Resize from 1 -> 3: adds 2 permits, should wake both waiters. + sem.resize(3) + await asyncio.sleep(0) + self.assertEqual(acquired_count, 2) + self.assertEqual(sem.active, 3) + + # Clean up — release all 3 acquired permits. + sem.release() + sem.release() + sem.release() + + # Await tasks to prevent warnings. + await asyncio.wait_for(asyncio.gather(t1, t2), timeout=5.0) + + asyncio.run(run()) + + def test_resize_up_partial_wake(self) -> None: + """When resize adds fewer permits than waiters, only some wake.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + acquired: list[int] = [] + + async def waiter(idx: int) -> None: + await sem.acquire() + acquired.append(idx) + + t1 = asyncio.create_task(waiter(1)) + await asyncio.sleep(0) + t2 = asyncio.create_task(waiter(2)) + await asyncio.sleep(0) + t3 = asyncio.create_task(waiter(3)) + await asyncio.sleep(0) + + # Resize from 1 -> 2: adds 1 permit, wakes 1 waiter. + sem.resize(2) + await asyncio.sleep(0) + self.assertEqual(len(acquired), 1) + self.assertEqual(acquired[0], 1) # FIFO + + # Clean up remaining waiters. + for t in (t1, t2, t3): + t.cancel() + try: + await t + except asyncio.CancelledError: + pass + + asyncio.run(run()) + + +class ResizableSemaphoreResizeDownTest(unittest.TestCase): + def test_resize_down_no_preemption(self) -> None: + """Active tasks continue after resize down.""" + + async def run() -> None: + sem = ResizableSemaphore(3) + await sem.acquire() + await sem.acquire() + await sem.acquire() + self.assertEqual(sem.active, 3) + + # Resize down to 1. All 3 are still active. + sem.resize(1) + self.assertEqual(sem.max_value, 1) + self.assertEqual(sem.active, 3) # No preemption. + + asyncio.run(run()) + + def test_resize_down_blocks_new_acquires(self) -> None: + """After resize down, new acquires block until drain completes.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(3) + await sem.acquire() + await sem.acquire() + await sem.acquire() + + sem.resize(1) + + acquired: asyncio.Event = asyncio.Event() + + async def try_acquire() -> None: + await sem.acquire() + acquired.set() + + task = asyncio.create_task(try_acquire()) + await asyncio.sleep(0) + self.assertFalse(acquired.is_set()) + + # Release 3 permits (drain from 3 active -> 0). + # First two releases bring active from 3->2->1 (at max). + # Third release frees a permit for the waiter. + sem.release() + sem.release() + sem.release() + await asyncio.sleep(0) + self.assertTrue(acquired.is_set()) + + # Clean up. + sem.release() + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(run()) + + def test_resize_down_drain_releases_clamp(self) -> None: + """Releases during drain clamp to max_value, not old max.""" + + async def run() -> None: + sem = ResizableSemaphore(5) + await sem.acquire() + await sem.acquire() + # active=2, current_value=3 + sem.resize(2) + # current_value should now be 0 (3 - (5-2) = 0) + self.assertEqual(sem.active, 2) + self.assertEqual(sem.max_value, 2) + + # Release one: active -> 1. + sem.release() + self.assertEqual(sem.active, 1) + + # Release another: active -> 0. + sem.release() + self.assertEqual(sem.active, 0) + + # Extra release should clamp. + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + +class ResizableSemaphoreResizeEdgeCasesTest(unittest.TestCase): + def test_resize_to_same_value(self) -> None: + async def run() -> None: + sem = ResizableSemaphore(3) + await sem.acquire() + sem.resize(3) + self.assertEqual(sem.max_value, 3) + self.assertEqual(sem.active, 1) + + asyncio.run(run()) + + def test_resize_to_zero_raises(self) -> None: + sem = ResizableSemaphore(3) + with self.assertRaises(ValueError): + sem.resize(0) + + def test_resize_to_negative_raises(self) -> None: + sem = ResizableSemaphore(3) + with self.assertRaises(ValueError): + sem.resize(-5) + + def test_resize_preserves_max_after_error(self) -> None: + """Failed resize should not change max_value.""" + sem = ResizableSemaphore(3) + try: + sem.resize(0) + except ValueError: + pass + self.assertEqual(sem.max_value, 3) + + def test_resize_while_waiters_pending_noop(self) -> None: + """Resize to same value while tasks are waiting should not wake them.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + acquired: asyncio.Event = asyncio.Event() + + async def waiter() -> None: + await sem.acquire() + acquired.set() + + task = asyncio.create_task(waiter()) + await asyncio.sleep(0) + + # Resize to same value — waiter should stay blocked. + sem.resize(1) + await asyncio.sleep(0) + self.assertFalse(acquired.is_set()) + + # Clean up. + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + asyncio.run(run()) + + def test_multiple_resizes(self) -> None: + """Multiple sequential resizes should work correctly.""" + + async def run() -> None: + sem = ResizableSemaphore(1) + sem.resize(5) + self.assertEqual(sem.max_value, 5) + sem.resize(2) + self.assertEqual(sem.max_value, 2) + sem.resize(10) + self.assertEqual(sem.max_value, 10) + + asyncio.run(run()) + + +class ResizableSemaphoreCancellationTest(unittest.TestCase): + def test_cancelled_waiter_is_removed(self) -> None: + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + async def waiter() -> None: + await sem.acquire() + + task = asyncio.create_task(waiter()) + await asyncio.sleep(0) + + # Cancel the waiter. + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + # Release — should not raise even though waiter was cancelled. + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + def test_cancel_one_of_multiple_waiters(self) -> None: + """Cancel middle waiter; remaining waiters still get served.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() + + order: list[int] = [] + + async def waiter(idx: int) -> None: + await sem.acquire() + order.append(idx) + sem.release() + + t1 = asyncio.create_task(waiter(1)) + await asyncio.sleep(0) + t2 = asyncio.create_task(waiter(2)) + await asyncio.sleep(0) + t3 = asyncio.create_task(waiter(3)) + await asyncio.sleep(0) + + # Cancel the second waiter. + t2.cancel() + try: + await t2 + except asyncio.CancelledError: + pass + + sem.release() + await asyncio.wait_for(asyncio.gather(t1, t3), timeout=5.0) + self.assertEqual(order, [1, 3]) + + asyncio.run(run()) + + +class ResizableSemaphorePermitLeakTest(unittest.TestCase): + """V5.1 (DESIGN.md): regression coverage for the permit-leak fix in + ``acquire()``'s cancellation handler. + + Without the fix in case (b) — when ``release()`` direct-hands a + permit via ``set_result(None)`` and the waiter is then cancelled + before resuming — the permit was lost. Each cancellation under + contention leaked one permit; eventually the semaphore deadlocked. + """ + + def test_acquire_cancel_after_grant_releases_permit(self) -> None: + """Case (b): waiter is granted a permit via direct hand-off but + cancelled before resuming. The permit must be released back so + another acquirer can make progress. Without the fix, the second + acquirer would deadlock. + """ + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() # holder takes the only permit. + + # First waiter: will be granted the permit by release(), + # then cancelled before its acquire() returns. + granted_w1: asyncio.Event = asyncio.Event() + + async def w1() -> None: + # Acquire will receive set_result(None) from the + # holder's release(); cancellation arrives during + # the await window before acquire() returns. + try: + await sem.acquire() + except asyncio.CancelledError: + granted_w1.set() + raise + + t1 = asyncio.create_task(w1()) + # Yield until W1 has enqueued its waiter future. + await asyncio.sleep(0) + self.assertEqual(len(sem._waiters), 1) + + # Direct-hand the permit to W1: this calls + # waiter.set_result(None) on W1's future. W1 is now in + # case (b): future is done with a result, but W1 hasn't + # resumed yet. + sem.release() + self.assertEqual(len(sem._waiters), 0) + + # Cancel W1 before its acquire() resumes. Permit-leak fix + # must detect case (b) and call self.release() to give + # the permit back. + t1.cancel() + try: + await t1 + except asyncio.CancelledError: + pass + self.assertTrue(granted_w1.is_set()) + + # Now a new acquirer should be able to acquire the permit + # immediately. WITHOUT THE FIX this would deadlock here. + second_acquired: asyncio.Event = asyncio.Event() + + async def w2() -> None: + await sem.acquire() + second_acquired.set() + + t2 = asyncio.create_task(w2()) + await asyncio.wait_for(t2, timeout=2.0) + self.assertTrue(second_acquired.is_set()) + + # Pool accounting: we have one active permit (W2's). + self.assertEqual(sem.active, 1) + + asyncio.run(run()) + + def test_release_grants_to_next_when_first_waiter_cancelled( + self, + ) -> None: + """Direct-hand-off race: ``release()`` pops a waiter whose + future is already cancelled (done()) and must skip it to hand + the permit to the next non-cancelled waiter. + + Sequence: + 1. sem with value=1, acquired by holder. + 2. Two waiters W1, W2 enqueued. + 3. W1.cancel() — its future becomes done() (cancelled). + 4. holder.release() — must pop W1 (skip), pop W2 (grant). + """ + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() # holder + + w1_started: asyncio.Event = asyncio.Event() + w2_started: asyncio.Event = asyncio.Event() + w2_acquired: asyncio.Event = asyncio.Event() + + async def w1() -> None: + w1_started.set() + await sem.acquire() + + async def w2() -> None: + w2_started.set() + await sem.acquire() + w2_acquired.set() + + t1 = asyncio.create_task(w1()) + await w1_started.wait() + t2 = asyncio.create_task(w2()) + await w2_started.wait() + # Force both to enqueue their waiter futures. + await asyncio.sleep(0) + await asyncio.sleep(0) + self.assertEqual(len(sem._waiters), 2) + + # Cancel W1 — its waiter future transitions to done() but + # remains in the deque (acquire()'s except handler removes + # it after the await re-raises CancelledError). + t1.cancel() + try: + await t1 + except asyncio.CancelledError: + pass + + # release() must pop W1 (skip — done()) then pop W2 (grant). + sem.release() + await asyncio.wait_for(t2, timeout=2.0) + self.assertTrue(w2_acquired.is_set()) + self.assertFalse(t2.cancelled()) + + asyncio.run(run()) + + def test_acquire_pre_cancelled_future_fast_path(self) -> None: + """Case (c) documentation: when the future has been cancelled + before the await even completes (e.g., the task was cancelled + before it ran past the ``self._waiters.append(fut)`` line), the + cancellation handler treats this the same as case (a) — the + future is still in ``self._waiters``, so ``self._waiters.remove + (fut)`` covers it. No permit was granted, so nothing to release + back. + + This test verifies the case (c) fast path via cancelling at + scheduling time rather than mid-await: pool accounting must + remain consistent (active=1 from the holder, max=1). + """ + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(1) + await sem.acquire() # holder + + async def w1() -> None: + # Will be cancelled before release() ever fires, while + # waiter future is still pending in the deque. + await sem.acquire() + + t1 = asyncio.create_task(w1()) + # Let W1 enqueue its waiter future. + await asyncio.sleep(0) + self.assertEqual(len(sem._waiters), 1) + + # Cancel W1 immediately — future is still pending (case (a)), + # which shares the cleanup path with case (c). + t1.cancel() + try: + await t1 + except asyncio.CancelledError: + pass + + # The waiter future must be cleaned up from the deque so + # subsequent release() doesn't try to grant to a dead + # waiter. + self.assertEqual(len(sem._waiters), 0) + + # Pool accounting: holder still owns the permit. + self.assertEqual(sem.active, 1) + self.assertEqual(sem.max_value, 1) + + # Release returns the permit cleanly. + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + +class ResizableSemaphoreConcurrencyTest(unittest.TestCase): + def test_concurrent_acquire_release(self) -> None: + """Many tasks acquiring and releasing concurrently.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(3) + completed = 0 + + async def worker() -> None: + nonlocal completed + await sem.acquire() + # Yield to let other tasks proceed. + await asyncio.sleep(0) + sem.release() + completed += 1 + + tasks = [asyncio.create_task(worker()) for _ in range(20)] + await asyncio.wait_for(asyncio.gather(*tasks), timeout=5.0) + self.assertEqual(completed, 20) + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + def test_concurrent_acquire_with_resize(self) -> None: + """Resize during concurrent acquire/release operations.""" + + async def run() -> None: + sem: ResizableSemaphore = ResizableSemaphore(2) + completed = 0 + + async def worker() -> None: + nonlocal completed + await sem.acquire() + await asyncio.sleep(0) + sem.release() + completed += 1 + + tasks = [asyncio.create_task(worker()) for _ in range(10)] + + # Yield a few times to let some workers start. + await asyncio.sleep(0) + await asyncio.sleep(0) + sem.resize(5) # Expand. + await asyncio.sleep(0) + sem.resize(1) # Contract. + + await asyncio.wait_for(asyncio.gather(*tasks), timeout=5.0) + self.assertEqual(completed, 10) + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + def test_active_never_exceeds_max_under_load(self) -> None: + """active should never exceed max_value when max is stable.""" + + async def run() -> None: + max_permits = 4 + sem: ResizableSemaphore = ResizableSemaphore(max_permits) + max_seen = 0 + + async def worker() -> None: + nonlocal max_seen + await sem.acquire() + current = sem.active + if current > max_seen: + max_seen = current + await asyncio.sleep(0) + sem.release() + + tasks = [asyncio.create_task(worker()) for _ in range(30)] + await asyncio.wait_for(asyncio.gather(*tasks), timeout=5.0) + self.assertLessEqual(max_seen, max_permits) + + asyncio.run(run()) + + +class ResizableSemaphorePropertiesTest(unittest.TestCase): + def test_max_value_reflects_resize(self) -> None: + sem = ResizableSemaphore(3) + self.assertEqual(sem.max_value, 3) + sem.resize(7) + self.assertEqual(sem.max_value, 7) + sem.resize(1) + self.assertEqual(sem.max_value, 1) + + def test_active_reflects_acquire_release(self) -> None: + async def run() -> None: + sem = ResizableSemaphore(5) + self.assertEqual(sem.active, 0) + await sem.acquire() + self.assertEqual(sem.active, 1) + await sem.acquire() + self.assertEqual(sem.active, 2) + sem.release() + self.assertEqual(sem.active, 1) + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + def test_active_exceeds_max_after_resize_down(self) -> None: + """active can temporarily exceed max_value after resize down.""" + + async def run() -> None: + sem = ResizableSemaphore(5) + for _ in range(5): + await sem.acquire() + self.assertEqual(sem.active, 5) + + sem.resize(2) + self.assertEqual(sem.max_value, 2) + self.assertEqual(sem.active, 5) # Exceeds max_value. + + # Drain. + for _ in range(5): + sem.release() + self.assertEqual(sem.active, 0) + + asyncio.run(run()) + + +if __name__ == "__main__": + unittest.main()