def on_message(self,
channel: Union[str, Pattern],
callback: Callable[..., None],
reconnect: bool):
pattern = not isinstance(channel, str)
channel_obj = aioredis.Channel(channel.pattern.replace('.*', '*').replace('.?', '?') if pattern else channel,
is_pattern=pattern)
tup = channel, channel_obj, callback, reconnect
async def message_handling():
try:
async for message in channel_obj.iter(decoder=lambda m: m.decode()):
if pattern:
if not self._call_callback(tup, False, [message[0].decode(), message[1]]):
break
else:
if not self._call_callback(tup, False, [channel, message]):
break
finally:
channel_obj.close()
self.at_end()
asyncio.ensure_future(message_handling())
if self.add_event_wait is None:
self.add_callback_event = asyncio.Future()
self.add_event_wait = asyncio.ensure_future(self.add_callback_event)
if not self.add_callback_event.done():
self.add_callback_event.set_result(tup)
self._add_callback(tup)
async def psubscribe(self, channel_name: str, callback: Callable):
if channel_name in self._psubscribed_channels:
item = self._psubscribed_channels[channel_name]
item.add_callback(callback)
return
channel = aioredis.Channel(channel_name, is_pattern=True)
await self.sub_conn.execute_pubsub('psubscribe', channel)
item = SubscribedItem(channel)
item.add_callback(callback)
self._psubscribed_channels[channel_name] = item
async def redis_pubsub_handler():
redis = await aioredis.create_connection(address=redis_connection_host,
db=redis_connection_db)
channel = aioredis.Channel("websocket", is_pattern=False)
redis.execute_pubsub("subscribe", channel)
while True:
message = await channel.get()
if message:
for queue in socket_queues.values():
queue.put_nowait(message.decode("UTF-8"))
async def _subscription_loop(self):
logger.info("Starting subscription loop...")
loop = asyncio.get_event_loop()
while self._started:
# If there are no subscribers then it's pointless to loop endlessly.
# Just block the execution until there is some subscriber available.
# This piece of code is unblocked by trying to obtain a result of some task.
if not self._subscription_list:
await self._non_empty_sub_list_event.wait()
self._non_empty_sub_list_event.clear()
continue
new_subs = False
for sub in self._subscription_list:
# This is a brand new subscriber if it doesn't have task assigned.
# We'll create a channel for subscriptions, start listening on the channel
# and register the subscription. New subscriber will also be flagged
# as a new subscriber so that we can later recognize it.
if sub._task is None:
sub._channel = aioredis.Channel(
self._task_key(sub.task_id), is_pattern=False
)
sub._task = loop.create_task(sub._channel.wait_message())
sub._new = True
self._subscription_table[sub._task] = sub
self._tasks.add(sub._task)
new_subs = True
else:
sub._new = False
# If there are any new subscribers then resend SUBSCRIBE command.
if new_subs:
logger.debug(
"Subscribing to channels\n{}".format(
"\n".join(
[
sub._channel.name.decode("utf-8")
for sub in self._subscription_list
]
)
)
)
self._pubsub_connection.execute_pubsub(
"subscribe", *[sub._channel for sub in self._subscription_list]
)
# Wait for any of the subscriptions channels to receive any message (or to be cancelled).
# If there are any new subscribers then set the timeout to 0.1s instead of 10s.
# This is because what might have happened is that we've checked the presence of the key
# and it didn't exist. But until we got the point where we started subscription channel,
# there could have been PUBLISH which we might have missed. So we set the really low timeout
# to recheck the presence of keys almost instantly and introduced guarantee of delivery.
done, pending = await asyncio.wait(
self._tasks,
return_when=asyncio.FIRST_COMPLETED,
timeout=10 if not new_subs else 0.1,
)
# Remove tasks which have been finished (by either being accepted or cancelled).
self._tasks.difference_update(done)
to_unsubscribe = []
try:
while done:
finished_task = done.pop()
sub = self._subscription_table[finished_task]
if not finished_task.cancelled():
result = await sub._channel.get()
else:
result = None
logger.debug(result)
# This actually set the future of the object that user is waiting on in Result.get()
if not (sub.result.cancelled() or sub.result.done()):
sub.result.set_result(result)
to_unsubscribe.append(
self._subscription_table[finished_task]._channel
)
self._subscription_list.remove(sub)
del self._subscription_table[finished_task]
except aioredis.errors.ChannelClosedError:
await retry(self._reconnect_pubsub_connection, max_wait_time=60.0)
continue
# Here we resolve the situation mentioned above about missing possible PUBLISH between initial check
# of the key and SUBSCRIBE. Do it only for new subscribers so that we don't issue huge MGET commands
# each iteration.
if new_subs:
new_subs_list = [sub for sub in self._subscription_list if sub._new]
# New subscriber might have been removed before because it obtained the message in the meantime.
if new_subs_list:
values = await self._connection.execute(
"mget", *[self._task_key(sub.task_id) for sub in new_subs_list]
)
for sub, value in filter(
lambda x: x[1] is not None, zip(new_subs_list, values)
):
# If we just cancel the task it will be naturally cancelled in the next iteration of wait()
# will pick it up and it will be removed as if the user have given up or we got PUBLISH.
if not (sub.result.cancelled() or sub.result.done()):
sub.result.set_result(value)
if not sub._task.done():
sub._task.cancel()
# In the end, UNSUBSCRIBE channels which already received their message to clean up Redis.
if to_unsubscribe:
await self._pubsub_connection.execute_pubsub(
"unsubscribe", *to_unsubscribe
)
logger.info("Stopping Subscription loop...")