async def echo(term: str, limit: asyncio.Semaphore): async with limit: print(term) await asyncio.sleep(0.1) if limit.locked(): print(f"limit crossed, sleeping for {2} seconds") await asyncio.sleep(2)
async def safe_make_request(url: str, limit: asyncio.Semaphore) -> int: async with limit: result = await make_request(url) if limit.locked(): print("\nlimit reached, sleeping for 2 seconds...\n") await asyncio.sleep(1.5) return result
async def consumer(result_queue, event: asyncio.Event, limit: asyncio.Semaphore) -> None: async with limit: while True: await event.wait() result = await result_queue.get() print(result) result_queue.task_done() await asyncio.sleep(0.1) if limit.locked(): print("sleeping for 2 sec") await asyncio.sleep(2)
class PriorityQueue: def __init__(self): self.queue = [] self.items = Semaphore(value=0) async def push(self, data, priority=0): self.queue.append((priority, json.dumps(data))) self.queue.sort() self.items.release() async def pop(self, timeout: int = 1) -> Any: try: await wait_for(self.items.acquire(), timeout) return json.loads(self.queue.pop(-1)[1]) except TimeoutError: return None async def pop_ready(self) -> Any: if self.items.locked(): return None await self.items.acquire() return json.loads(self.queue.pop(-1)[1]) async def score(self, data): data = json.dumps(data) for priority, item in self.queue: if data == item: return priority return None async def rank(self, data): data = json.dumps(data) for index, (_, item) in enumerate(self.queue): if data == item: return len(self.queue) - index - 1 return None async def clear(self): self.queue = [] self.items = Semaphore(value=0) async def length(self): return len(self.queue)
async def _drain_semaphore(semaphore: asyncio.Semaphore): while not semaphore.locked(): try: await asyncio.wait_for(semaphore.acquire(), 0.1) except asyncio.TimeoutError: break
class Queue: def __init__(self, n_slots): self._n_slots = n_slots self._wait_tx = OrderedDict() self._wait_tx_sem = Semaphore(0) self._wait_rx = OrderedDict() @property def n_slots(self): return self._n_slots @property def is_full(self): return len(self._wait_tx) + len(self._wait_rx) >= self._n_slots @property def qsize(self): return self._n_slots def add(self, request): instance_id = request.instanceId analysis_id = request.analysisId key = (instance_id, analysis_id) existing = self._wait_tx.get(key) if existing is not None: ex_request, ex_stream = existing log.debug('%s %s', 'cancelling', req_str(ex_request)) ex_stream.cancel() else: existing = self._wait_rx.get(key) if existing is not None: ex_request, ex_stream = existing log.debug('%s %s', 'cancelling', req_str(ex_request)) ex_stream.cancel() if self.is_full: raise QueueFull stream = Stream() log.debug('%s %s', 'queueing', req_str(request)) self._wait_tx[key] = (request, stream) if self._wait_tx_sem.locked(): self._wait_tx_sem.release() stream.add_complete_listener(self._stream_complete) return stream def get(self, key): value = self._wait_tx.get(key) if value is not None: return value value = self._wait_rx.get(key) if value is not None: return value return None def _stream_complete(self): for key, value in self._wait_rx.items(): request, stream = value if stream.is_complete: del self._wait_rx[key] break log.debug('%s %s', 'removing', req_str(request)) def stream(self): # # this isn't compatible with python 3.5: # while True: # await self._wait_tx_sem.acquire() # while len(self._wait_tx) > 0: # key, value = self._wait_tx.popitem() # self._wait_rx[key] = value # request, stream = value # log.debug('%s %s', 'yielding', req_str(request)) # yield value # # so we had to do this: class AsyncGenerator: def __init__(self, parent): self._parent = parent def __aiter__(self): return self async def __anext__(self): if len(self._parent._wait_tx) == 0: await self._parent._wait_tx_sem.acquire() key, value = self._parent._wait_tx.popitem() self._parent._wait_rx[key] = value request, stream = value log.debug('%s %s', 'yielding', req_str(request)) return value return AsyncGenerator(self) def __contains__(self, value): return value in self._wait_tx or value in self._wait_rx
class Pool: def __init__(self, n_slots): self._n_slots = n_slots self._wait_tx = OrderedDict() self._wait_tx_sem = Semaphore(0) self._wait_rx = OrderedDict() self._not_full = Event() self._not_full.set() @property def n_slots(self): return self._n_slots @property def is_full(self): return len(self._wait_tx) + len(self._wait_rx) >= self._n_slots def full(self): return len(self._wait_tx) + len(self._wait_rx) >= self._n_slots async def wait_not_full(self): await self._not_full.wait() @property def qsize(self): return self._n_slots def put_nowait(self, request): instance_id = request.instanceId analysis_id = request.analysisId key = (instance_id, analysis_id) existing = self._wait_tx.get(key) if existing is not None: ex_request, ex_stream = existing log.debug('%s %s', 'cancelling', req_str(ex_request)) ex_stream.cancel() else: existing = self._wait_rx.get(key) if existing is not None: ex_request, ex_stream = existing log.debug('%s %s', 'cancelling', req_str(ex_request)) ex_stream.cancel() if self.full(): raise QueueFull stream = Stream() log.debug('%s %s', 'queueing', req_str(request)) self._wait_tx[key] = (request, stream) if self._wait_tx_sem.locked(): self._wait_tx_sem.release() stream.add_complete_listener(self._stream_complete) if self.is_full: self._not_full.clear() return stream def add(self, request): return self.put_nowait(request) def cancel(self, key): existing = self._wait_tx.get(key) if existing is not None: ex_request, ex_stream = existing log.debug('%s %s', 'cancelling', req_str(ex_request)) ex_stream.cancel() else: existing = self._wait_rx.get(key) if existing is not None: ex_request, ex_stream = existing log.debug('%s %s', 'cancelling', req_str(ex_request)) ex_stream.cancel() else: raise KeyError def get(self, key): value = self._wait_tx.get(key) if value is not None: return value value = self._wait_rx.get(key) if value is not None: return value return None def _stream_complete(self): # iterate through a copied list so we can delete from the original for key, value in list(self._wait_rx.items()): request, stream = value if stream.is_complete: del self._wait_rx[key] log.debug('%s %s', 'removing', req_str(request)) for key, value in list(self._wait_tx.items()): request, stream = value if stream.is_complete: del self._wait_tx[key] log.debug('%s %s', 'removing', req_str(request)) if not self.is_full: self._not_full.set() async def stream(self): while True: await self._wait_tx_sem.acquire() while len(self._wait_tx) > 0: key, value = self._wait_tx.popitem() self._wait_rx[key] = value request, stream = value log.debug('%s %s', 'yielding', req_str(request)) yield value def __contains__(self, value): return value in self._wait_tx or value in self._wait_rx