diff --git a/Lib/asyncio/locks.py b/Lib/asyncio/locks.py index c59ea15111bef2b..9aa0313ecfc8744 100644 --- a/Lib/asyncio/locks.py +++ b/Lib/asyncio/locks.py @@ -488,7 +488,13 @@ def __init__(self, parties): self._parties = parties self._state = _BarrierState.FILLING - self._count = 0 # count tasks in Barrier + # Tickets of tasks currently in the barrier, in arrival order for + # the current round. A task's index is only assigned once the + # round's full set of tickets is known (see _release()), so that + # a task leaving early (e.g. via cancellation) can't cause a + # later arrival to reuse its index. + self._present = [] + self._release_index = {} def __repr__(self): res = super().__repr__() @@ -514,17 +520,20 @@ async def wait(self): """ async with self._cond: await self._block() # Block while the barrier drains or resets. + ticket = object() + self._present.append(ticket) try: - index = self._count - self._count += 1 - if index + 1 == self._parties: + if len(self._present) == self._parties: # We release the barrier await self._release() else: await self._wait() - return index + return self._release_index[ticket] finally: - self._count -= 1 + try: + self._present.remove(ticket) + except ValueError: + pass # Wake up any tasks waiting for barrier to drain. self._exit() @@ -547,6 +556,13 @@ async def _block(self): async def _release(self): # Release the tasks waiting in the barrier. + # Assign each currently-present party (including this task, the + # last arrival) a unique index in arrival order, before any of + # them get a chance to leave self._present. + self._release_index = { + ticket: i for i, ticket in enumerate(self._present) + } + # Enter draining state. # Next waiting tasks will be blocked until the end of draining. self._state = _BarrierState.DRAINING @@ -566,9 +582,10 @@ async def _wait(self): def _exit(self): # If we are the last tasks to exit the barrier, signal any tasks # waiting for the barrier to drain. - if self._count == 0: + if not self._present: if self._state in (_BarrierState.RESETTING, _BarrierState.DRAINING): self._state = _BarrierState.FILLING + self._release_index = {} self._cond.notify_all() async def reset(self): @@ -578,7 +595,7 @@ async def reset(self): raised. """ async with self._cond: - if self._count > 0: + if self._present: if self._state is not _BarrierState.RESETTING: #reset the barrier, waking up tasks self._state = _BarrierState.RESETTING @@ -605,7 +622,7 @@ def parties(self): def n_waiting(self): """Return the number of tasks currently waiting at the barrier.""" if self._state is _BarrierState.FILLING: - return self._count + return len(self._present) return 0 @property diff --git a/Lib/test/test_asyncio/test_locks.py b/Lib/test/test_asyncio/test_locks.py index 320a933e144813d..1f103794af29744 100644 --- a/Lib/test/test_asyncio/test_locks.py +++ b/Lib/test/test_asyncio/test_locks.py @@ -1527,6 +1527,39 @@ async def coro(): self.assertEqual(barrier.n_waiting, 0) self.assertFalse(barrier.broken) + async def test_filling_tasks_cancel_one_index_not_reused(self): + # See gh-155233: a task cancelled while the barrier is still + # filling used to leave its index available for reuse, so a + # later arrival could be assigned the same index as an + # already-waiting task from the same round. + self.N = 3 + barrier = asyncio.Barrier(self.N) + results = [] + + async def coro(): + i = await barrier.wait() + results.append(i) + + t1 = asyncio.create_task(coro()) + t2 = asyncio.create_task(coro()) + await asyncio.sleep(0) + await asyncio.sleep(0) + self.assertEqual(barrier.n_waiting, 2) + + t1.cancel() + with self.assertRaises(asyncio.CancelledError): + await t1 + await asyncio.sleep(0) + self.assertEqual(barrier.n_waiting, 1) + + t3 = asyncio.create_task(coro()) + t4 = asyncio.create_task(coro()) + await asyncio.gather(t2, t3, t4) + + self.assertEqual(sorted(results), list(range(self.N))) + self.assertEqual(barrier.n_waiting, 0) + self.assertFalse(barrier.broken) + async def test_reset_barrier(self): barrier = asyncio.Barrier(1) @@ -1576,7 +1609,7 @@ async def coro(): results1.append(True) else: # here drained task outside the barrier - if rest_of_tasks == barrier._count: + if rest_of_tasks == len(barrier._present): # tasks outside the barrier await barrier.reset() diff --git a/Misc/NEWS.d/next/Library/2026-08-05-11-15-07.gh-issue-155233.Qm3xLp.rst b/Misc/NEWS.d/next/Library/2026-08-05-11-15-07.gh-issue-155233.Qm3xLp.rst new file mode 100644 index 000000000000000..0f3a6e43329b05c --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-08-05-11-15-07.gh-issue-155233.Qm3xLp.rst @@ -0,0 +1,3 @@ +Fix :meth:`asyncio.Barrier.wait` reusing the numeric index of a task that +was cancelled while the barrier was still filling, which could cause two +tasks in the same successful release to receive the same index.