bpo-43352: Add a Barrier object in asyncio lib (GH-24903)

Co-authored-by: Yury Selivanov <[email protected]>
Co-authored-by: Andrew Svetlov <[email protected]>

Author: 3 people

Doc/library/asyncio-api-index.rst

@@ -186,11 +186,16 @@ Threading-like synchronization primitives that can be used in Tasks.
     * - :class:`BoundedSemaphore`
       - A bounded semaphore.
 
+    * - :class:`Barrier`
+      - A barrier object.
+
 
 .. rubric:: Examples
 
 * :ref:`Using asyncio.Event <asyncio_example_sync_event>`.
 
+* :ref:`Using asyncio.Barrier <asyncio_example_barrier>`.
+
 * See also the documentation of asyncio
   :ref:`synchronization primitives <asyncio-sync>`.
 
@@ -206,6 +211,9 @@ Exceptions
     * - :exc:`asyncio.CancelledError`
       - Raised when a Task is cancelled. See also :meth:`Task.cancel`.
 
+    * - :exc:`asyncio.BrokenBarrierError`
+      - Raised when a Barrier is broken. See also :meth:`Barrier.wait`.
+
 
 .. rubric:: Examples
 

Doc/library/asyncio-sync.rst

@@ -28,6 +28,7 @@ asyncio has the following basic synchronization primitives:
 * :class:`Condition`
 * :class:`Semaphore`
 * :class:`BoundedSemaphore`
+* :class:`Barrier`
 
 
 ---------
@@ -340,6 +341,115 @@ BoundedSemaphore
    .. versionchanged:: 3.10
       Removed the *loop* parameter.
 
+
+Barrier
+=======
+
+.. class:: Barrier(parties, action=None)
+
+   A barrier object.  Not thread-safe.
+
+   A barrier is a simple synchronization primitive that allows to block until
+   *parties* number of tasks are waiting on it.
+   Tasks can wait on the :meth:`~Barrier.wait` method and would be blocked until
+   the specified number of tasks end up waiting on :meth:`~Barrier.wait`.
+   At that point all of the waiting tasks would unblock simultaneously.
+
+   :keyword:`async with` can be used as an alternative to awaiting on
+   :meth:`~Barrier.wait`.
+
+   The barrier can be reused any number of times.
+
+   .. _asyncio_example_barrier:
+
+   Example::
+
+      async def example_barrier():
+         # barrier with 3 parties
+         b = asyncio.Barrier(3)
+
+         # create 2 new waiting tasks
+         asyncio.create_task(b.wait())
+         asyncio.create_task(b.wait())
+
+         await asyncio.sleep(0)
+         print(b)
+
+         # The third .wait() call passes the barrier
+         await b.wait()
+         print(b)
+         print("barrier passed")
+
+         await asyncio.sleep(0)
+         print(b)
+
+      asyncio.run(example_barrier())
+
+   Result of this example is::
+
+      <asyncio.locks.Barrier object at 0x... [filling, waiters:2/3]>
+      <asyncio.locks.Barrier object at 0x... [draining, waiters:0/3]>
+      barrier passed
+      <asyncio.locks.Barrier object at 0x... [filling, waiters:0/3]>
+
+   .. versionadded:: 3.11
+
+   .. coroutinemethod:: wait()
+
+      Pass the barrier. When all the tasks party to the barrier have called
+      this function, they are all unblocked simultaneously.
+
+      When a waiting or blocked task in the barrier is cancelled,
+      this task exits the barrier which stays in the same state.
+      If the state of the barrier is "filling", the number of waiting task
+      decreases by 1.
+
+      The return value is an integer in the range of 0 to ``parties-1``, different
+      for each task. This can be used to select a task to do some special
+      housekeeping, e.g.::
+
+         ...
+         async with barrier as position:
+            if position == 0:
+               # Only one task print this
+               print('End of *draining phasis*')
+
+      This method may raise a :class:`BrokenBarrierError` exception if the
+      barrier is broken or reset while a task is waiting.
+      It could raise a :exc:`CancelledError` if a task is cancelled.
+
+   .. coroutinemethod:: reset()
+
+      Return the barrier to the default, empty state.  Any tasks waiting on it
+      will receive the :class:`BrokenBarrierError` exception.
+
+      If a barrier is broken it may be better to just leave it and create a new one.
+
+   .. coroutinemethod:: abort()
+
+      Put the barrier into a broken state.  This causes any active or future
+      calls to :meth:`wait` to fail with the :class:`BrokenBarrierError`.
+      Use this for example if one of the taks needs to abort, to avoid infinite
+      waiting tasks.
+
+   .. attribute:: parties
+
+      The number of tasks required to pass the barrier.
+
+   .. attribute:: n_waiting
+
+      The number of tasks currently waiting in the barrier while filling.
+
+   .. attribute:: broken
+
+      A boolean that is ``True`` if the barrier is in the broken state.
+
+
+.. exception:: BrokenBarrierError
+
+   This exception, a subclass of :exc:`RuntimeError`, is raised when the
+   :class:`Barrier` object is reset or broken.
+
 ---------
 
 

Lib/asyncio/exceptions.py

@@ -1,7 +1,8 @@
 """asyncio exceptions."""
 
 
-__all__ = ('CancelledError', 'InvalidStateError', 'TimeoutError',
+__all__ = ('BrokenBarrierError',
+           'CancelledError', 'InvalidStateError', 'TimeoutError',
            'IncompleteReadError', 'LimitOverrunError',
            'SendfileNotAvailableError')
 
@@ -55,3 +56,7 @@ def __init__(self, message, consumed):
 
     def __reduce__(self):
         return type(self), (self.args[0], self.consumed)
+
+
+class BrokenBarrierError(RuntimeError):
+    """Barrier is broken by barrier.abort() call."""

Lib/asyncio/locks.py

@@ -1,14 +1,15 @@
 """Synchronization primitives."""
 
-__all__ = ('Lock', 'Event', 'Condition', 'Semaphore', 'BoundedSemaphore')
+__all__ = ('Lock', 'Event', 'Condition', 'Semaphore',
+           'BoundedSemaphore', 'Barrier')
 
 import collections
+import enum
 
 from . import exceptions
 from . import mixins
 from . import tasks
 
-
 class _ContextManagerMixin:
     async def __aenter__(self):
         await self.acquire()
@@ -416,3 +417,155 @@ def release(self):
         if self._value >= self._bound_value:
             raise ValueError('BoundedSemaphore released too many times')
         super().release()
+
+
+
+class _BarrierState(enum.Enum):
+    FILLING = 'filling'
+    DRAINING = 'draining'
+    RESETTING = 'resetting'
+    BROKEN = 'broken'
+
+
+class Barrier(mixins._LoopBoundMixin):
+    """Asyncio equivalent to threading.Barrier
+
+    Implements a Barrier primitive.
+    Useful for synchronizing a fixed number of tasks at known synchronization
+    points. Tasks block on 'wait()' and are simultaneously awoken once they
+    have all made their call.
+    """
+
+    def __init__(self, parties):
+        """Create a barrier, initialised to 'parties' tasks."""
+        if parties < 1:
+            raise ValueError('parties must be > 0')
+
+        self._cond = Condition() # notify all tasks when state changes
+
+        self._parties = parties
+        self._state = _BarrierState.FILLING
+        self._count = 0       # count tasks in Barrier
+
+    def __repr__(self):
+        res = super().__repr__()
+        extra = f'{self._state.value}'
+        if not self.broken:
+            extra += f', waiters:{self.n_waiting}/{self.parties}'
+        return f'<{res[1:-1]} [{extra}]>'
+
+    async def __aenter__(self):
+        # wait for the barrier reaches the parties number
+        # when start draining release and return index of waited task
+        return await self.wait()
+
+    async def __aexit__(self, *args):
+        pass
+
+    async def wait(self):
+        """Wait for the barrier.
+
+        When the specified number of tasks have started waiting, they are all
+        simultaneously awoken.
+        Returns an unique and individual index number from 0 to 'parties-1'.
+        """
+        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:
+                    # We release the barrier
+                    await self._release()
+                else:
+                    await self._wait()
+                return index
+            finally:
+                self._count -= 1
+                # Wake up any tasks waiting for barrier to drain.
+                self._exit()
+
+    async def _block(self):
+        # Block until the barrier is ready for us,
+        # or raise an exception if it is broken.
+        #
+        # It is draining or resetting, wait until done
+        # unless a CancelledError occurs
+        await self._cond.wait_for(
+            lambda: self._state not in (
+                _BarrierState.DRAINING, _BarrierState.RESETTING
+            )
+        )
+
+        # see if the barrier is in a broken state
+        if self._state is _BarrierState.BROKEN:
+            raise exceptions.BrokenBarrierError("Barrier aborted")
+
+    async def _release(self):
+        # Release the tasks waiting in the barrier.
+
+        # Enter draining state.
+        # Next waiting tasks will be blocked until the end of draining.
+        self._state = _BarrierState.DRAINING
+        self._cond.notify_all()
+
+    async def _wait(self):
+        # Wait in the barrier until we are released. Raise an exception
+        # if the barrier is reset or broken.
+
+        # wait for end of filling
+        # unless a CancelledError occurs
+        await self._cond.wait_for(lambda: self._state is not _BarrierState.FILLING)
+
+        if self._state in (_BarrierState.BROKEN, _BarrierState.RESETTING):
+            raise exceptions.BrokenBarrierError("Abort or reset of barrier")
+
+    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 self._state in (_BarrierState.RESETTING, _BarrierState.DRAINING):
+                self._state = _BarrierState.FILLING
+            self._cond.notify_all()
+
+    async def reset(self):
+        """Reset the barrier to the initial state.
+
+        Any tasks currently waiting will get the BrokenBarrier exception
+        raised.
+        """
+        async with self._cond:
+            if self._count > 0:
+                if self._state is not _BarrierState.RESETTING:
+                    #reset the barrier, waking up tasks
+                    self._state = _BarrierState.RESETTING
+            else:
+                self._state = _BarrierState.FILLING
+            self._cond.notify_all()
+
+    async def abort(self):
+        """Place the barrier into a 'broken' state.
+
+        Useful in case of error.  Any currently waiting tasks and tasks
+        attempting to 'wait()' will have BrokenBarrierError raised.
+        """
+        async with self._cond:
+            self._state = _BarrierState.BROKEN
+            self._cond.notify_all()
+
+    @property
+    def parties(self):
+        """Return the number of tasks required to trip the barrier."""
+        return self._parties
+
+    @property
+    def n_waiting(self):
+        """Return the number of tasks currently waiting at the barrier."""
+        if self._state is _BarrierState.FILLING:
+            return self._count
+        return 0
+
+    @property
+    def broken(self):
+        """Return True if the barrier is in a broken state."""
+        return self._state is _BarrierState.BROKEN