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
