class SetBasedAsyncPopPool(collabc.AsyncGenerator, AsyncPopulationPool[H]):
"""An asynchronous population pool backed by a set."""
def __init__(self, ivals: Optional[AbstractSet[H]] = None):
self._stopped = False
self._set = Set[H]() if ivals is None else {ivals}
self._ac = None
self._stopped = False
async def apopulate(self, val: H, *args: H) -> None:
#Can't populate a closed pool:
if self._stopped: raise StopAsyncIteration
if self._ac is None: self._ac = ACondition()
if args is None or len(args) == 0:
self._set.add(val)
count = 1
else:
argset = {args}
argset.add(val)
count = len(argset)
self._set |= argset
async with self._ac:
self._ac.notify(count)
async def asend(self, value: Optional[H]) -> Optional[H]:
if self._stopped: raise StopAsyncIteration
if self._ac is None: self._ac = ACondition()
if value is None:
async with self._ac:
while len(self._set) == 0:
self._ac.wait()
if self._stopped: raise StopAsyncIteration
return self._set.pop()
else:
if value not in self._set:
self._set.add(value)
await self._ac.acquire()
self._ac.notify()
self._ac.release()
async def athrow(self, typ, val=None, tb=None) -> None:
try:
return await super().athrow(typ, val, tb)
except (Exception, GeneratorExit) as exc:
self._stopped = True
if self._ac is not None:
await self._ac.acquire()
self._ac.notify_all()
self._ac.release()
raise StopAsyncIteration from exc
async def wait_condition_or_timeout(condition: asyncio.Condition,
timeout: float) -> None:
"""Wait for a condition or timeout."""
loop = asyncio.get_event_loop()
future = loop.create_future()
def _handle_timeout() -> None:
if not future.done():
future.set_result(None)
timer_handle = loop.call_later(timeout, _handle_timeout)
condition_wait = loop.create_task(condition.wait())
def _handle_wait_complete(_: asyncio.Task) -> None:
if not future.done():
future.set_result(None)
condition_wait.add_done_callback(_handle_wait_complete)
try:
await future
finally:
timer_handle.cancel()
if not condition_wait.done():
condition_wait.cancel()
with contextlib.suppress(asyncio.CancelledError):
await condition_wait
async def wait_on_condition_with_timeout(self,
condition: asyncio.Condition,
timeout: float) -> bool:
loop = self._loop
# Create a future that will be triggered by either completion or timeout.
waiter = loop.create_future()
# Callback to trigger the future. The varargs are there to consume and void any arguments passed.
# This allows the same callback to be used in loop.call_later and wait_task.add_done_callback,
# which automatically passes the finished future in.
def release_waiter(*_):
if not waiter.done():
waiter.set_result(None)
# Set up the timeout
timeout_handle = loop.call_later(timeout, release_waiter)
# Launch the wait task
wait_task = loop.create_task(condition.wait())
wait_task.add_done_callback(release_waiter)
try:
await waiter # Returns on wait complete or timeout
if wait_task.done():
return True
else:
raise asyncio.TimeoutError()
except (asyncio.TimeoutError, asyncio.CancelledError):
# If timeout or cancellation occur, clean up, cancel the wait, let it handle the cancellation,
# then re-raise.
wait_task.remove_done_callback(release_waiter)
wait_task.cancel()
await asyncio.wait([wait_task])
raise
finally:
timeout_handle.cancel()
