asyncio CapacityLimiter lost token on cancellation #947

try:
    self.acquire_on_behalf_of_nowait(borrower)
except WouldBlock:
    # 1. limiter is full -> put ourself in queue
    event = asyncio.Event()
    self._wait_queue[borrower] = event
    try:
        # 2. return control to event loop and wait
        await event.wait()
    except BaseException:
        # 3. Imagine 
        #    - the waiting task gets cancelled (from outside or by timeout), and 
        #    - current borrower finished -> calls return() -> calls event.set()
        #  so, event loop resumed on `await event.wait()` with `CancelError` and the event is set
        # I assume event loop has a choice here on which one of these to resume execution, and it chooses cancel
        # so we were given a token, but been cancelled
        # we pop ourselves from the waiting queue, but as the event is set - token has been given to us, 
        # but we aren't using it, so it will be lost if we don't pass it to the next waiter
        self._wait_queue.pop(borrower, None)
        raise

    self._borrowers.add(borrower)

def _wake_next_waiter(self) -> None:
    """
    Wake up the next waiter in the queue if there is capacity.
    This is the core logic that should be called whenever a token slot might have become free.
    """
    if self._wait_queue and len(self._borrowers) < self._total_tokens:
        # Pop the waiter from the queue to "reserve" the slot for them
        event = self._wait_queue.popitem(last=False)[1]
        event.set()

async def acquire_on_behalf_of(self, borrower: object) -> None:
    await AsyncIOBackend.checkpoint_if_cancelled()
    try:
        self.acquire_on_behalf_of_nowait(borrower)
    except WouldBlock:
        event = asyncio.Event()
        self._wait_queue[borrower] = event
        try:
            await event.wait()
        except BaseException:
            self._wait_queue.pop(borrower, None)

            # If the event is set, it means we were given a token by a
            # release() call, but we were cancelled before we could claim it
            # We are now responsible for passing the token on to the next waiter.
            if event.is_set():
                self._wake_next_waiter()

            raise

        self._borrowers.add(borrower)
    else:
        try:
            await AsyncIOBackend.cancel_shielded_checkpoint()
        except BaseException:
            self.release_on_behalf_of(borrower)
            raise

import asyncio

import anyio
import pytest
from anyio._backends._asyncio import CapacityLimiter
from anyio.abc import CancelScope
from more_itertools import one
from pytest import MonkeyPatch


@pytest.fixture
def anyio_backend():
    return "asyncio"


VIRTUAL_LOG: list[str] = []


class SlowEvent(asyncio.Event):
    async def wait(self):
        VIRTUAL_LOG.append("SlowEvent.wait")
        await super().wait()
        await asyncio.sleep(0)


async def releaser_task(
    limiter: CapacityLimiter, token_acquired: asyncio.Event, release_token: asyncio.Event
) -> None:
    """Acquires the token, waits for a signal, then releases it."""

    VIRTUAL_LOG.append("Releaser: trying to acquire...")
    async with limiter:
        VIRTUAL_LOG.append("Releaser: ...acquired.")
        token_acquired.set()
        await release_token.wait()
        VIRTUAL_LOG.append("Releaser: ...releasing token now.")

    # The token is released on __aexit__
    # Yield control to allow the controller to cancel the waiter
    await asyncio.sleep(0)
    VIRTUAL_LOG.append("Releaser: finished.")


async def waiter_task(
    limiter: CapacityLimiter, token_waiting: asyncio.Event, waiter_scope_box: list[CancelScope]
) -> None:
    """The task that will be cancelled at the critical moment."""

    with anyio.CancelScope() as scope:
        waiter_scope_box.append(scope)
        VIRTUAL_LOG.append("Waiter: trying to acquire...")
        token_waiting.set()
        async with limiter:
            pytest.fail("Waiter task was not cancelled and acquired the lock.")

    VIRTUAL_LOG.append("Waiter: finished (was cancelled).")


async def stranded_task(limiter: CapacityLimiter, token_acquired: asyncio.Event) -> None:
    """This task should get stuck waiting for the 'lost' token."""

    VIRTUAL_LOG.append("Stranded: trying to acquire...")
    async with limiter:
        VIRTUAL_LOG.append("Stranded: ...acquired. The test has failed.")
        token_acquired.set()

    pytest.fail("Stranded task should never finish.")


@pytest.mark.asyncio
async def test_capacity_limiter_has_cancel_race_deadlock(
    monkeypatch: MonkeyPatch
) -> None:
    """
    Tests for the specific deadlock scenario in CapacityLimiter where a waiting
    task is cancelled just after it has been woken up by a release, but before
    it has a chance to run and claim the token.

    This causes the token to be "lost", stranding other waiting tasks.
    """

    limiter = CapacityLimiter(1)
    monkeypatch.setattr("anyio._backends._asyncio.asyncio.Event", SlowEvent)

    # Events to control the task execution order
    releaser_has_acquired = asyncio.Event()
    waiter_is_waiting = asyncio.Event()
    release_the_token = asyncio.Event()

    # A place to store the waiter's cancel scope so we can cancel it from the test
    waiter_scope_box: list[CancelScope] = []

    # The task that will be stranded if the deadlock occurs
    stranded_task_finished = asyncio.Event()

    anyio.create_task_group
    async with anyio.create_task_group() as tg:
        # 1. Start the releaser and wait for it to acquire the only token.
        tg.start_soon(releaser_task, limiter, releaser_has_acquired, release_the_token)
        await releaser_has_acquired.wait()
        assert limiter.borrowed_tokens == 1

        # 2. Start the waiter task and wait for it to be blocked.
        tg.start_soon(waiter_task, limiter, waiter_is_waiting, waiter_scope_box)
        await waiter_is_waiting.wait()

        # Give the waiter a chance to actually enter the limiter's wait queue
        await asyncio.sleep(0)
        assert limiter.borrowed_tokens == 1
        assert len(limiter._wait_queue) == 1

        # 3. Start the stranded task, which will wait behind the first waiter.
        tg.start_soon(stranded_task, limiter, stranded_task_finished)
        await asyncio.sleep(0)
        assert len(limiter._wait_queue) == 2

        # 4. THE RACE: Tell the releaser to release the token.
        # The limiter will wake up the waiter_task.
        # Then, immediately cancel the waiter_task before it gets a chance to run.
        VIRTUAL_LOG.append("Controller: Triggering release and immediate cancel...")
        release_the_token.set()

        # Yield control. The releaser will run, release the token (waking the waiter),
        # and then yield again. The waiter is now "ready", but we get control back.
        await asyncio.sleep(0)
        # Now, cancel the waiter task.
        one(waiter_scope_box).cancel()
        VIRTUAL_LOG.append("Controller: ...waiter has been cancelled.")

        # 5. Check if the stranded_task is stuck.
        # We give it a short timeout.
        # If it were going to acquire the lock, it would do so almost instantly.
        VIRTUAL_LOG.append("Controller: Waiting to see if stranded_task acquires the token...")
        with anyio.move_on_after(0.1) as scope:
            await stranded_task_finished.wait()

        # The deadlock occurred if the timeout was reached.
        assert scope.cancelled_caught is True, "Deadlock did not occur, stranded_task finished."

        VIRTUAL_LOG.append("Controller: Stranded_task did not finish, as expected.")

        # 6. Final state check.
        #  The token was lost. No one is borrowing, but the stranded_task is still waiting.
        assert limiter.borrowed_tokens == 0
        assert len(limiter._wait_queue) == 1
        VIRTUAL_LOG.append("Controller: limiter state is corrupted")

        # exit the test
        tg.cancel_scope.cancel()