primitives/barrier.py

   1 # barrier.py
   2 # Copyright (c) 2018-2020 Peter Hinch
   3 # Released under the MIT License (MIT) - see LICENSE file
   4
   5 # Now uses Event rather than polling.
   6
   7 try:
   8     import uasyncio as asyncio
   9 except ImportError:
  10     import asyncio
  11
  12 from . import launch
  13
  14 # A Barrier synchronises N coros. Each issues await barrier.
  15 # Execution pauses until all other participant coros are waiting on it.
  16 # At that point the callback is executed. Then the barrier is 'opened' and
  17 # execution of all participants resumes.
  18
  19 class Barrier():
  20     def __init__(self, participants, func=None, args=()):
  21         self._participants = participants
  22         self._count = participants
  23         self._func = func
  24         self._args = args
  25         self._res = None
  26         self._evt = asyncio.Event()
  27
  28     def __await__(self):
  29         if self.trigger():
  30             return  # Other tasks have already reached barrier
  31         await self._evt.wait()  # Wait until last task reaches it
  32
  33     __iter__ = __await__
  34
  35     def result(self):
  36         return self._res
  37
  38     def trigger(self):
  39         self._count -=1
  40         if self._count < 0:
  41             raise ValueError('Too many tasks accessing Barrier')
  42         if self._count > 0:
  43             return False  # At least 1 other task has not reached barrier
  44         # All other tasks are waiting
  45         if self._func is not None:
  46             self._res = launch(self._func, self._args)
  47         self._count = self._participants
  48         self._evt.set()  # Release others
  49         self._evt.clear()
  50         return True
  51
  52     def busy(self):
  53         return  self._count < self._participants