async def receive(self, name: str):
"""
Wait for a message on a specific queue
This is a blocking operation
"""
assert name in self.queues, f"Missing queue {name}"
queue = self.queues[name]
logger.debug("Wait for message on bus", queue=name, instance=queue)
if isinstance(queue, RedisQueue):
async with AsyncRedis() as redis:
_, payload = await redis.blpop(queue.name)
assert isinstance(payload, bytes)
try:
return pickle.loads(payload)
except Exception as e:
logger.error("Bad redis payload", error=str(e))
await asyncio.sleep(1)
return
elif isinstance(queue, asyncio.Queue):
return await queue.get()
else:
# Run the synchronous mp queue.get in the asynchronous loop
# but use an asyncio sleep to be able to react to cancellation
async def _get():
while True:
try:
return queue.get(timeout=0)
except Empty:
await asyncio.sleep(1)
return await _get()
async def get(self, key):
if self.redis_enabled and key not in self.cache:
async with AsyncRedis() as redis:
val = await redis.get(self._redis_key(key))
if val is not None:
self.cache[key] = pickle.loads(val)
return self.cache[key]
async def rem(self, key):
if self.redis_enabled:
async with AsyncRedis() as redis:
await redis.delete(self._redis_key(key))
try:
del self.cache[key]
except KeyError:
pass
async def set(self, key, value):
self.cache[key] = value
if self.redis_enabled:
async with AsyncRedis() as redis:
await redis.expire(self.name, self.expiration)
await redis.set(self._redis_key(key),
pickle.dumps(value),
expire=self.expiration)
async def send(self, name: str, payload: Any):
"""
Send a message on a specific queue
"""
assert name in self.queues, f"Missing queue {name}"
queue = self.queues[name]
if isinstance(queue, RedisQueue):
async with AsyncRedis() as redis:
await redis.rpush(queue.name, pickle.dumps(payload))
elif isinstance(queue, asyncio.Queue):
await queue.put(payload)
else:
# Run the synchronous mp queue.put in the asynchronous loop
await asyncio.get_running_loop().run_in_executor(
None, lambda: queue.put(payload)
)