Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

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

def __repr__(self):
res = super().__repr__()
Expand All @@ -515,13 +516,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 +571,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 +587,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
29 changes: 29 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,35 @@ 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)

if __name__ == '__main__':
unittest.main()
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fix the returned value of the :meth:`asyncio.Barrier.wait`. This value is now calculated on demand during the *draining* phasis.
Loading