async def before_server_start(_sanic, loop): mpsc = Receiver(loop=loop) redis = await storage.get_async_redis_pool(loop) await redis.subscribe( mpsc.channel('scoreboard'), mpsc.channel('stolen_flags'), ) sio.start_background_task(background_task, mpsc)
async def redis_listener(): redis = await aioredis.create_redis('redis://redishost:6379') receiver = Receiver() await redis.subscribe(receiver.channel('sockets:notification:message'), receiver.channel('sockets:notification:message2')) while await receiver.wait_message(): sender, msg = await receiver.get() for sid in users: await sio.emit('notification:message', {'data': str(msg)}, room=sid)
async def listen_redis(self, websocket, channels): channel = channels.pop(0) try: mpsc = Receiver(loop=asyncio.get_event_loop()) await self.sub.subscribe( mpsc.channel(channel), *(mpsc.channel(channel) for channel in channels)) async for channel, msg in mpsc.iter(): if websocket.client_state == WebSocketState.CONNECTED: await websocket.send_bytes(msg) except: import traceback traceback.print_exc() await self.close_connections() finally: logger.info('Connection closed')
async def test_stopped(create_connection, server, caplog): sub = await create_connection(server.tcp_address) pub = await create_connection(server.tcp_address) mpsc = Receiver() await sub.execute_pubsub("subscribe", mpsc.channel("channel:1")) assert mpsc.is_active mpsc.stop() caplog.clear() with caplog.at_level("DEBUG", "aioredis"): await pub.execute("publish", "channel:1", b"Hello") await asyncio.sleep(0) assert len(caplog.record_tuples) == 1 # Receiver must have 1 EndOfStream message message = ( "Pub/Sub listener message after stop: " "sender: <_Sender name:b'channel:1', is_pattern:False, receiver:" "<Receiver is_active:False, senders:1, qsize:0>>, data: b'Hello'") assert caplog.record_tuples == [ ("aioredis", logging.WARNING, message), ] # assert (await mpsc.get()) is None with pytest.raises(ChannelClosedError): await mpsc.get() res = await mpsc.wait_message() assert res is False
async def test_pubsub_receiver_iter(create_redis, server, event_loop): sub = await create_redis(server.tcp_address) pub = await create_redis(server.tcp_address) mpsc = Receiver() async def coro(mpsc): lst = [] async for msg in mpsc.iter(): lst.append(msg) return lst tsk = asyncio.ensure_future(coro(mpsc)) (snd1, ) = await sub.subscribe(mpsc.channel("chan:1")) (snd2, ) = await sub.subscribe(mpsc.channel("chan:2")) (snd3, ) = await sub.psubscribe(mpsc.pattern("chan:*")) subscribers = await pub.publish_json("chan:1", {"Hello": "World"}) assert subscribers > 1 subscribers = await pub.publish_json("chan:2", ["message"]) assert subscribers > 1 event_loop.call_later(0, mpsc.stop) await asyncio.sleep(0.01) assert await tsk == [ (snd1, b'{"Hello": "World"}'), (snd3, (b"chan:1", b'{"Hello": "World"}')), (snd2, b'["message"]'), (snd3, (b"chan:2", b'["message"]')), ] assert not mpsc.is_active
async def test_pubsub_receiver_iter(create_redis, server, loop): sub = await create_redis(server.tcp_address, loop=loop) pub = await create_redis(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) async def coro(mpsc): lst = [] async for msg in mpsc.iter(): lst.append(msg) return lst tsk = asyncio.ensure_future(coro(mpsc), loop=loop) snd1, = await sub.subscribe(mpsc.channel('chan:1')) snd2, = await sub.subscribe(mpsc.channel('chan:2')) snd3, = await sub.psubscribe(mpsc.pattern('chan:*')) subscribers = await pub.publish_json('chan:1', {'Hello': 'World'}) assert subscribers > 1 subscribers = await pub.publish_json('chan:2', ['message']) assert subscribers > 1 loop.call_later(0, mpsc.stop) # await asyncio.sleep(0, loop=loop) assert await tsk == [ (snd1, b'{"Hello": "World"}'), (snd3, (b'chan:1', b'{"Hello": "World"}')), (snd2, b'["message"]'), (snd3, (b'chan:2', b'["message"]')), ] assert not mpsc.is_active
class Messager: def __init__(self, inbound, outbound, loop): self.loop = loop self.conn = None self.inbound = inbound self.outbound = outbound self.receiver = None self.task = None self.replies = dict() self.expected = set() async def initialize(self): self.conn = await aioredis.create_redis_pool(("localhost", 6379), encoding="utf-8", maxsize=2) self.receiver = Receiver(loop=self.loop) await self.conn.subscribe(self.receiver.channel(self.inbound)) self.task = self.loop.create_task(self.fetcher()) print( f'Redis connection established, listening on {self.inbound}, sending on {self.outbound}' ) # FIXME - propper logging async def terminate(self): # terminate channels and disconnect from redis self.conn.unsubscribe(self.inbound) self.task.cancel() self.receiver.stop() self.conn.close() await self.conn.wait_closed() async def fetcher(self): async for sender, message in self.receiver.iter(encoding='utf-8', decoder=json.loads): channel = sender.name.decode() if channel == self.inbound: uid = message["uid"] if uid not in self.expected: print("Unexpected message!") print(message) else: self.expected.remove(uid) self.replies[uid] = message async def get_reply(self, data): try: return (await asyncio.wait_for(self._get_reply(data), 10))["reply"] except TimeoutError: raise Redisception("No reply received from the bot!") async def _get_reply(self, data): uid = str(uuid.uuid4()) self.expected.add(uid) data["uid"] = uid await self.conn.publish_json(self.outbound, data) while uid not in self.replies: await asyncio.sleep(0.1) reply = self.replies[uid] del self.replies[uid] return reply
async def _update_log_info(self): def process_log_batch(log_batch): ip = log_batch["ip"] pid = str(log_batch["pid"]) if pid != "autoscaler": logs_for_ip = dict(DataSource.ip_and_pid_to_logs.get(ip, {})) logs_for_pid = list(logs_for_ip.get(pid, [])) logs_for_pid.extend(log_batch["lines"]) # Only cache upto MAX_LOGS_TO_CACHE logs_length = len(logs_for_pid) if logs_length > MAX_LOGS_TO_CACHE * LOG_PRUNE_THREASHOLD: offset = logs_length - MAX_LOGS_TO_CACHE del logs_for_pid[:offset] logs_for_ip[pid] = logs_for_pid DataSource.ip_and_pid_to_logs[ip] = logs_for_ip logger.info(f"Received a log for {ip} and {pid}") aioredis_client = self._dashboard_head.aioredis_client receiver = Receiver() channel = receiver.channel(gcs_utils.LOG_FILE_CHANNEL) await aioredis_client.subscribe(channel) logger.info("Subscribed to %s", channel) async for sender, msg in receiver.iter(): try: data = json.loads(ray._private.utils.decode(msg)) data["pid"] = str(data["pid"]) process_log_batch(data) except Exception: logger.exception("Error receiving log from Redis.")
async def test_stopped(create_connection, server, loop): sub = await create_connection(server.tcp_address, loop=loop) pub = await create_connection(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) await sub.execute_pubsub('subscribe', mpsc.channel('channel:1')) assert mpsc.is_active mpsc.stop() with logs('aioredis', 'DEBUG') as cm: await pub.execute('publish', 'channel:1', b'Hello') await asyncio.sleep(0, loop=loop) assert len(cm.output) == 1 # Receiver must have 1 EndOfStream message warn_messaege = ( "WARNING:aioredis:Pub/Sub listener message after stop: " "sender: <_Sender name:b'channel:1', is_pattern:False, receiver:" "<Receiver is_active:False, senders:1, qsize:0>>, data: b'Hello'" ) assert cm.output == [warn_messaege] # assert (await mpsc.get()) is None with pytest.raises(ChannelClosedError): await mpsc.get() res = await mpsc.wait_message() assert res is False
async def test_stopped(create_connection, server, loop): sub = await create_connection(server.tcp_address, loop=loop) pub = await create_connection(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) await sub.execute_pubsub('subscribe', mpsc.channel('channel:1')) assert mpsc.is_active mpsc.stop() with logs('aioredis', 'DEBUG') as cm: await pub.execute('publish', 'channel:1', b'Hello') await asyncio.sleep(0, loop=loop) assert len(cm.output) == 1 # Receiver must have 1 EndOfStream message warn_messaege = ( "WARNING:aioredis:Pub/Sub listener message after stop: " "sender: <_Sender name:b'channel:1', is_pattern:False, receiver:" "<Receiver is_active:False, senders:1, qsize:0>>, data: b'Hello'") assert cm.output == [warn_messaege] # assert (await mpsc.get()) is None with pytest.raises(ChannelClosedError): await mpsc.get() res = await mpsc.wait_message() assert res is False
async def producer(self): channel_names = self.channel_names or [] channel_patterns = self.channel_patterns or [] if not channel_names and not channel_patterns: return mpsc = Receiver() if channel_names: channels = [mpsc.channel(c) for c in channel_names] await self.redis.subscribe(*channels) if channel_patterns: tasks = set() for p in channel_patterns: tasks.add(self.redis.psubscribe(mpsc.pattern(p))) if tasks: await asyncio.wait(tasks) try: await self.receiver_reader(mpsc) finally: if channel_names: await self.redis.unsubscribe(*channel_names) if channel_patterns: tasks = set() for p in channel_patterns: tasks.add(self.redis.punsubscribe(p)) await asyncio.wait(tasks) mpsc.stop()
async def redis_relay(websocket): conn = await aioredis.create_connection(("localhost", 6379)) receiver = Receiver() conn.execute_pubsub("subscribe", receiver.channel("marketupdates")) while await receiver.wait_message(): *_, message = await receiver.get() await websocket.send(message.decode())
async def receive(self): print('Receive from redis') sub = await aioredis.create_connection(('localhost', 6379)) receiver = Receiver() sub.execute_pubsub('subscribe', receiver.channel(self.channel)) while (await receiver.wait_message()): msg = await receiver.get() print("Got Message:", msg)
async def subscribe(self): if not self._subscribe: redis = await self.redis # self._subscribe = await redis.subscribe('ws') r = Receiver(loop=self.loop) for ch in self._channels: await redis.subscribe(r.channel(ch)) self._subscribe = r return self._subscribe
async def subscribe_to_channel(loop, redis_pool): mpsc = Receiver(loop=loop) logger.info("Aquiring redis connection from pool") connection = await redis_pool.acquire() logger.info("Subscribing to redis channel: {}".format(websocket_channel)) await connection.execute_pubsub("subscribe", mpsc.channel(websocket_channel)) return mpsc
async def test_unsubscribe(create_connection, server): sub = await create_connection(server.tcp_address) pub = await create_connection(server.tcp_address) mpsc = Receiver() await sub.execute_pubsub( "subscribe", mpsc.channel("channel:1"), mpsc.channel("channel:3") ) res = await pub.execute("publish", "channel:3", "Hello world") assert res == 1 res = await pub.execute("publish", "channel:1", "Hello world") assert res == 1 assert mpsc.is_active assert (await mpsc.wait_message()) is True ch, msg = await mpsc.get() assert ch.name == b"channel:3" assert not ch.is_pattern assert msg == b"Hello world" assert (await mpsc.wait_message()) is True ch, msg = await mpsc.get() assert ch.name == b"channel:1" assert not ch.is_pattern assert msg == b"Hello world" await sub.execute_pubsub("unsubscribe", "channel:1") assert mpsc.is_active res = await pub.execute("publish", "channel:3", "message") assert res == 1 assert (await mpsc.wait_message()) is True ch, msg = await mpsc.get() assert ch.name == b"channel:3" assert not ch.is_pattern assert msg == b"message" waiter = asyncio.ensure_future(mpsc.get()) await sub.execute_pubsub("unsubscribe", "channel:3") assert not mpsc.is_active assert await waiter is None
async def _event_listener(self, channel): async for _ in AsyncCirculator(): async with self._redis_pool.get_client() as cache: receiver = Receiver() await cache.subscribe(receiver.channel(channel)) async for channel, message in receiver.iter(): await self._event_assigner(channel, message)
async def test_unsubscribe(create_connection, server, loop): sub = await create_connection(server.tcp_address, loop=loop) pub = await create_connection(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) await sub.execute_pubsub('subscribe', mpsc.channel('channel:1'), mpsc.channel('channel:3')) res = await pub.execute("publish", "channel:3", "Hello world") assert res == 1 res = await pub.execute("publish", "channel:1", "Hello world") assert res == 1 assert mpsc.is_active assert (await mpsc.wait_message()) is True ch, msg = await mpsc.get() assert ch.name == b'channel:3' assert not ch.is_pattern assert msg == b"Hello world" assert (await mpsc.wait_message()) is True ch, msg = await mpsc.get() assert ch.name == b'channel:1' assert not ch.is_pattern assert msg == b"Hello world" await sub.execute_pubsub('unsubscribe', 'channel:1') assert mpsc.is_active res = await pub.execute("publish", "channel:3", "message") assert res == 1 assert (await mpsc.wait_message()) is True ch, msg = await mpsc.get() assert ch.name == b'channel:3' assert not ch.is_pattern assert msg == b"message" await sub.execute_pubsub('unsubscribe', 'channel:3') assert not mpsc.is_active res = await mpsc.get() assert res is None
async def main(): await redis.connect() mpsc = Receiver(loop=asyncio.get_event_loop()) asyncio.ensure_future(reader(mpsc)) await redis._redis.subscribe(mpsc.channel('channel:1')) while True: try: await asyncio.sleep(10) print('hearbeat', flush=True) except Exception as ex: print(ex, flush=True) mpsc.stop() await redis.disconnect()
async def test_subscriptions(create_connection, server): sub = await create_connection(server.tcp_address) pub = await create_connection(server.tcp_address) mpsc = Receiver() await sub.execute_pubsub("subscribe", mpsc.channel("channel:1"), mpsc.channel("channel:3")) res = await pub.execute("publish", "channel:3", "Hello world") assert res == 1 res = await pub.execute("publish", "channel:1", "Hello world") assert res == 1 assert mpsc.is_active ch, msg = await mpsc.get() assert ch.name == b"channel:3" assert not ch.is_pattern assert msg == b"Hello world" ch, msg = await mpsc.get() assert ch.name == b"channel:1" assert not ch.is_pattern assert msg == b"Hello world"
async def subscribe_chat(self): redis_for_pubsub = await aioredis.create_redis(self.redis.address, db=self.redis.db) receiver = Receiver() channel = receiver.channel(CHAT_CHANNEL) await redis_for_pubsub.subscribe(channel) async def _get(): return (await receiver.get())[1] yield channel, _get await redis_for_pubsub.unsubscribe(CHAT_CHANNEL)
async def test_subscriptions(create_connection, server, loop): sub = await create_connection(server.tcp_address, loop=loop) pub = await create_connection(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) await sub.execute_pubsub('subscribe', mpsc.channel('channel:1'), mpsc.channel('channel:3')) res = await pub.execute("publish", "channel:3", "Hello world") assert res == 1 res = await pub.execute("publish", "channel:1", "Hello world") assert res == 1 assert mpsc.is_active ch, msg = await mpsc.get() assert ch.name == b'channel:3' assert not ch.is_pattern assert msg == b"Hello world" ch, msg = await mpsc.get() assert ch.name == b'channel:1' assert not ch.is_pattern assert msg == b"Hello world"
async def test_decode_message_error(loop): mpsc = Receiver(loop) ch = mpsc.channel('channel:1') ch.put_nowait(b'{"hello": "world"}') unexpected = (mock.ANY, {'hello': 'world'}) with pytest.raises(TypeError): assert (await mpsc.get(decoder=json.loads)) == unexpected ch = mpsc.pattern('*') ch.put_nowait((b'channel', b'{"hello": "world"}')) unexpected = (mock.ANY, b'channel', {'hello': 'world'}) with pytest.raises(TypeError): assert (await mpsc.get(decoder=json.loads)) == unexpected
async def test_decode_message_error(): mpsc = Receiver() ch = mpsc.channel("channel:1") ch.put_nowait(b'{"hello": "world"}') unexpected = (mock.ANY, {"hello": "world"}) with pytest.raises(TypeError): assert (await mpsc.get(decoder=json.loads)) == unexpected ch = mpsc.pattern("*") ch.put_nowait((b"channel", b'{"hello": "world"}')) unexpected = (mock.ANY, b"channel", {"hello": "world"}) with pytest.raises(TypeError): assert (await mpsc.get(decoder=json.loads)) == unexpected
def test_listener_channel(loop): mpsc = Receiver(loop=loop) assert not mpsc.is_active ch_a = mpsc.channel("channel:1") assert isinstance(ch_a, AbcChannel) assert mpsc.is_active ch_b = mpsc.channel('channel:1') assert ch_a is ch_b assert ch_a.name == ch_b.name assert ch_a.is_pattern == ch_b.is_pattern assert mpsc.is_active # remember id; drop refs to objects and create new one; ch_a.close() assert not ch_a.is_active assert not mpsc.is_active ch = mpsc.channel("channel:1") assert ch is not ch_a assert dict(mpsc.channels) == {b'channel:1': ch} assert dict(mpsc.patterns) == {}
async def test_pubsub_receiver_call_stop_with_empty_queue( create_redis, server, loop): sub = await create_redis(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) # FIXME: currently at least one subscriber is needed snd1, = await sub.subscribe(mpsc.channel('chan:1')) now = loop.time() loop.call_later(.5, mpsc.stop) async for i in mpsc.iter(): # noqa (flake8 bug with async for) assert False, "StopAsyncIteration not raised" dt = loop.time() - now assert dt <= 1.5 assert not mpsc.is_active
async def subscribe_handler(_request, ws): data = await ws.recv() try: decoded_data = ujson.decode(data) except ValueError: await ws.send(ujson.dumps({'error': 'invalid json data'})) return token = decoded_data.get('token', '') loop = asyncio.get_event_loop() redis = await storage.get_async_redis(loop) mpsc = Receiver(loop=loop) await redis.subscribe(mpsc.channel(f'updates:{token}')) async for channel, msg in mpsc.iter(): await ws.send(msg.decode())
async def test_wait_message(create_connection, server, loop): sub = await create_connection(server.tcp_address, loop=loop) pub = await create_connection(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) await sub.execute_pubsub('subscribe', mpsc.channel('channel:1')) fut = asyncio.ensure_future(mpsc.wait_message(), loop=loop) assert not fut.done() await asyncio.sleep(0, loop=loop) assert not fut.done() await pub.execute('publish', 'channel:1', 'hello') await asyncio.sleep(0, loop=loop) # read in connection await asyncio.sleep(0, loop=loop) # call Future.set_result assert fut.done() res = await fut assert res is True
def test_wait_message(create_connection, server, loop): sub = yield from create_connection(server.tcp_address, loop=loop) pub = yield from create_connection(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) yield from sub.execute_pubsub('subscribe', mpsc.channel('channel:1')) fut = async_task(mpsc.wait_message(), loop=loop) assert not fut.done() yield from asyncio.sleep(0, loop=loop) assert not fut.done() yield from pub.execute('publish', 'channel:1', 'hello') yield from asyncio.sleep(0, loop=loop) # read in connection yield from asyncio.sleep(0, loop=loop) # call Future.set_result assert fut.done() res = yield from fut assert res is True
async def work(self, app): """ Start work, fetch tasks :param app: LRS object """ self.app = app redis_pool = await aioredis.create_pool( (app.redis_host, app.redis_port), minsize=5, maxsize=10, password=app.redis_password or None, db=app.redis_db, loop=app.loop) self.redis = await aioredis.Redis(redis_pool) mpsc = Receiver(loop=app.loop) asyncio.ensure_future(self.reader(mpsc), loop=app.loop) await self.redis.subscribe(mpsc.channel('crawl:finish'))
async def test_wait_message(create_connection, server): sub = await create_connection(server.tcp_address) pub = await create_connection(server.tcp_address) mpsc = Receiver() await sub.execute_pubsub("subscribe", mpsc.channel("channel:1")) fut = asyncio.ensure_future(mpsc.wait_message()) assert not fut.done() await asyncio.sleep(0) assert not fut.done() await pub.execute("publish", "channel:1", "hello") await asyncio.sleep(0) # read in connection await asyncio.sleep(0) # call Future.set_result assert fut.done() res = await fut assert res is True
async def test_decode_message(loop): mpsc = Receiver(loop) ch = mpsc.channel('channel:1') ch.put_nowait(b'Some data') res = await mpsc.get(encoding='utf-8') assert isinstance(res[0], _Sender) assert res[1] == 'Some data' ch.put_nowait('{"hello": "world"}') res = await mpsc.get(decoder=json.loads) assert isinstance(res[0], _Sender) assert res[1] == {'hello': 'world'} ch.put_nowait(b'{"hello": "world"}') res = await mpsc.get(encoding='utf-8', decoder=json.loads) assert isinstance(res[0], _Sender) assert res[1] == {'hello': 'world'}
async def test_unsubscribe(create_connection, server, loop): sub = await create_connection(server.tcp_address, loop=loop) pub = await create_connection(server.tcp_address, loop=loop) mpsc = Receiver(loop=loop) await sub.execute_pubsub('subscribe', mpsc.channel('channel:1'), mpsc.channel('channel:3')) res = await pub.execute("publish", "channel:3", "Hello world") assert res == 1 res = await pub.execute("publish", "channel:1", "Hello world") assert res == 1 assert mpsc.is_active assert (await mpsc.wait_message()) is True ch, msg = await mpsc.get() assert ch.name == b'channel:3' assert not ch.is_pattern assert msg == b"Hello world" assert (await mpsc.wait_message()) is True ch, msg = await mpsc.get() assert ch.name == b'channel:1' assert not ch.is_pattern assert msg == b"Hello world" await sub.execute_pubsub('unsubscribe', 'channel:1') assert mpsc.is_active res = await pub.execute("publish", "channel:3", "message") assert res == 1 assert (await mpsc.wait_message()) is True ch, msg = await mpsc.get() assert ch.name == b'channel:3' assert not ch.is_pattern assert msg == b"message" waiter = asyncio.ensure_future(mpsc.get(), loop=loop) await sub.execute_pubsub('unsubscribe', 'channel:3') assert not mpsc.is_active assert await waiter is None
async def test_pubsub_receiver_stop_on_disconnect(create_redis, server, loop): pub = await create_redis(server.tcp_address, loop=loop) sub = await create_redis(server.tcp_address, loop=loop) sub_name = 'sub-{:X}'.format(id(sub)) await sub.client_setname(sub_name) for sub_info in await pub.client_list(): if sub_info.name == sub_name: break assert sub_info.name == sub_name mpsc = Receiver(loop=loop) await sub.subscribe(mpsc.channel('channel:1')) await sub.subscribe(mpsc.channel('channel:2')) await sub.psubscribe(mpsc.pattern('channel:*')) q = asyncio.Queue(loop=loop) EOF = object() async def reader(): async for ch, msg in mpsc.iter(encoding='utf-8'): await q.put((ch.name, msg)) await q.put(EOF) tsk = asyncio.ensure_future(reader(), loop=loop) await pub.publish_json('channel:1', ['hello']) await pub.publish_json('channel:2', ['hello']) # receive all messages assert await q.get() == (b'channel:1', '["hello"]') assert await q.get() == (b'channel:*', (b'channel:1', '["hello"]')) assert await q.get() == (b'channel:2', '["hello"]') assert await q.get() == (b'channel:*', (b'channel:2', '["hello"]')) # XXX: need to implement `client kill` assert await pub.execute('client', 'kill', sub_info.addr) in (b'OK', 1) await asyncio.wait_for(tsk, timeout=1, loop=loop) assert await q.get() is EOF