Skip to content

Commit e35d340

Browse files
miss-islingtonYvesDupblurb-it[bot]asvetlov
authored
[3.14] gh-155233: Fix asyncio.Barrier reusing a cancelled task's index (alternative to GH-155235) (GH-158957) (#161411)
gh-155233: Fix asyncio.Barrier reusing a cancelled task's index (alternative to GH-155235) (GH-158957) * Initial commit: - add a private ``_index`` attr - add a test Co-authored-by: Sreehari Annam * 📜🤖 Added by blurb_it. * Apply suggestions from code review * Update Lib/asyncio/locks.py * Update Lib/asyncio/locks.py --------- (cherry picked from commit 49c43f7) Co-authored-by: Duprat <yduprat@gmail.com> Co-authored-by: blurb-it[bot] <43283697+blurb-it[bot]@users.noreply.github.com> Co-authored-by: Andrew Svetlov <andrew.svetlov@gmail.com>
1 parent 704f0ed commit e35d340

3 files changed

Lines changed: 38 additions & 2 deletions

File tree

‎Lib/asyncio/locks.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -492,6 +492,8 @@ def __init__(self, parties):
492492
self._parties = parties
493493
self._state = _BarrierState.FILLING
494494
self._count = 0 # count tasks in Barrier
495+
# index number returned when barrier drains
496+
self._index = self._parties - 1
495497

496498
def __repr__(self):
497499
res = super().__repr__()
@@ -518,13 +520,14 @@ async def wait(self):
518520
async with self._cond:
519521
await self._block() # Block while the barrier drains or resets.
520522
try:
521-
index = self._count
522523
self._count += 1
523-
if index + 1 == self._parties:
524+
if self._count == self._parties:
524525
# We release the barrier
525526
await self._release()
526527
else:
527528
await self._wait()
529+
index = self._index
530+
self._index = (self._index + 1) % self._parties
528531
return index
529532
finally:
530533
self._count -= 1
@@ -572,6 +575,7 @@ def _exit(self):
572575
if self._count == 0:
573576
if self._state in (_BarrierState.RESETTING, _BarrierState.DRAINING):
574577
self._state = _BarrierState.FILLING
578+
self._index = self._parties - 1
575579
self._cond.notify_all()
576580

577581
async def reset(self):
@@ -587,6 +591,7 @@ async def reset(self):
587591
self._state = _BarrierState.RESETTING
588592
else:
589593
self._state = _BarrierState.FILLING
594+
self._index = self._parties - 1
590595
self._cond.notify_all()
591596

592597
async def abort(self):

‎Lib/test/test_asyncio/test_locks.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1820,6 +1820,36 @@ async def coro():
18201820

18211821
self.assertEqual(barrier1.n_waiting, 0)
18221822

1823+
async def test_filling_tasks_cancel_one_index_still_unique(self):
1824+
# See gh-155233: a task cancelled while the barrier is still
1825+
# filling used to leave its index available for reuse.
1826+
# Index counter is now calculated on demand during the draining phasis.
1827+
self.N = 3
1828+
barrier = asyncio.Barrier(self.N)
1829+
results = []
1830+
1831+
async def coro():
1832+
i = await barrier.wait()
1833+
results.append(i)
1834+
1835+
t1 = asyncio.create_task(coro())
1836+
t2 = asyncio.create_task(coro())
1837+
await asyncio.sleep(0)
1838+
self.assertEqual(barrier.n_waiting, 2)
1839+
1840+
t1.cancel()
1841+
with self.assertRaises(asyncio.CancelledError):
1842+
await t1
1843+
await asyncio.sleep(0)
1844+
self.assertEqual(barrier.n_waiting, 1)
1845+
1846+
t3 = asyncio.create_task(coro())
1847+
t4 = asyncio.create_task(coro())
1848+
await asyncio.gather(t2, t3, t4)
1849+
1850+
self.assertEqual(sorted(results), list(range(self.N)))
1851+
self.assertEqual(barrier.n_waiting, 0)
1852+
18231853

18241854
if __name__ == '__main__':
18251855
unittest.main()
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Fix :meth:`asyncio.Barrier.wait` returning duplicate indices when a waiting task is cancelled before the barrier is released.```

0 commit comments

Comments
 (0)