File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff 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 ):
Original file line number Diff line number Diff 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
18241854if __name__ == '__main__' :
18251855 unittest .main ()
Original file line number Diff line number Diff line change 1+ Fix :meth: `asyncio.Barrier.wait ` returning duplicate indices when a waiting task is cancelled before the barrier is released.```
You can’t perform that action at this time.
0 commit comments