Skip to content
9 changes: 7 additions & 2 deletions Lib/asyncio/locks.py
Original file line number Diff line number Diff line change
Expand Up @@ -489,6 +489,8 @@ def __init__(self, parties):
self._parties = parties
self._state = _BarrierState.FILLING
self._count = 0 # count tasks in Barrier
# index number returned when barrier drains
self._index = self._parties - 1

def __repr__(self):
res = super().__repr__()
Expand All @@ -515,13 +517,14 @@ async def wait(self):
async with self._cond:
await self._block() # Block while the barrier drains or resets.
try:
index = self._count
self._count += 1
if index + 1 == self._parties:
if self._count == self._parties:
# We release the barrier
await self._release()
else:
await self._wait()
index = self._index
self._index = (self._index + 1) % self._parties
return index
finally:
self._count -= 1
Expand Down Expand Up @@ -569,6 +572,7 @@ def _exit(self):
if self._count == 0:
if self._state in (_BarrierState.RESETTING, _BarrierState.DRAINING):
self._state = _BarrierState.FILLING
self._index = self._parties - 1
self._cond.notify_all()

async def reset(self):
Expand All @@ -584,6 +588,7 @@ async def reset(self):
self._state = _BarrierState.RESETTING
else:
self._state = _BarrierState.FILLING
self._index = self._parties - 1
self._cond.notify_all()

async def abort(self):
Expand Down
30 changes: 30 additions & 0 deletions Lib/test/test_asyncio/test_locks.py
Original file line number Diff line number Diff line change
Expand Up @@ -1820,6 +1820,36 @@ async def coro():

self.assertEqual(barrier1.n_waiting, 0)

async def test_filling_tasks_cancel_one_index_still_unique(self):
# See gh-155233: a task cancelled while the barrier is still
# filling used to leave its index available for reuse.
# Index counter is now calculated on demand during the draining phasis.
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)
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)
Comment thread
asvetlov marked this conversation as resolved.
Resolved


if __name__ == '__main__':
unittest.main()
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fix :meth:`asyncio.Barrier.wait` returning duplicate indices when a waiting task is cancelled before the barrier is released.```
Loading