From 6fbb65507f709fe22524a813321d7030e4c6140d Mon Sep 17 00:00:00 2001 From: "Nyeki.py" <84998222+Nekidev@users.noreply.github.com> Date: Sun, 20 Oct 2024 21:51:48 -0300 Subject: [PATCH] feat(bot): :sparkles: Flush all messages to DB on call hangup. --- payphone/cogs/calls/cog.py | 140 +++++++++++++++++++++------------- payphone/cogs/calls/embeds.py | 65 +++++++--------- payphone/cogs/mods/cog.py | 12 ++- payphone/core/constants.py | 19 +++++ payphone/core/utils.py | 29 ++++++- 5 files changed, 172 insertions(+), 93 deletions(-) diff --git a/payphone/cogs/calls/cog.py b/payphone/cogs/calls/cog.py index 45a684f..f8ad535 100644 --- a/payphone/cogs/calls/cog.py +++ b/payphone/cogs/calls/cog.py @@ -30,62 +30,15 @@ class Calls(base.Cog): self.hangup_calls.add_exception_type(exceptions.ConfigurationError) self.hangup_calls.start() + self.flush_messages.add_exception_type(exceptions.ConfigurationError) self.flush_messages.add_exception_type(TypeError) self.flush_messages.start() def cog_unload(self) -> None: self.hangup_calls.cancel() - self.flush_messages.cancel() return super().cog_unload() - @tasks.loop(seconds=5) - async def flush_messages(self) -> None: - """Flush all the in-memory messages and message contents to the DB.""" - - pipe = redis.pipeline() - - pipe.hgetall(constants.MESSAGES_QUEUE) - pipe.hgetall(constants.MESSAGE_CONTENTS_QUEUE) - - pipe.delete( - constants.MESSAGES_QUEUE, - constants.MESSAGE_CONTENTS_QUEUE, - constants.MESSAGES_PROPAGATION_MAP, - ) - - messages, message_contents, _ = await pipe.execute() - - messages = sorted( - [models.Message(**orjson.loads(msg)) for msg in messages.values()], - key=lambda m: m.created_at, - ) - message_contents = sorted( - [ - models.MessageContent(**orjson.loads(msg)) - for msg in message_contents.values() - ], - key=lambda m: m.created_at, - ) - - while True: - try: - await models.Message.bulk_create(messages, batch_size=500) - break - - except asyncpg.exceptions.DeadlockDetectedError: - continue - - while True: - try: - await models.MessageContent.bulk_create( - message_contents, batch_size=500 - ) - break - - except asyncpg.exceptions.DeadlockDetectedError: - continue - @tasks.loop(seconds=15) async def hangup_calls(self) -> None: """Hangs up inactive calls. @@ -129,9 +82,79 @@ class Calls(base.Cog): hung_up_channel_ids = [id for id in hung_up_channel_ids if id is not None] if len(hung_up_channel_ids) > 0: - await models.Channel.filter(id__in=hung_up_channel_ids).update(call=None) + await models.Channel.filter(id__in=hung_up_channel_ids).update(call_id=None) await redis.srem(constants.CHANNELS_ONCALL_SET, *hung_up_channel_ids) + @tasks.loop(seconds=15) + async def flush_messages(self) -> None: + """Fetches all the pending messages from Redis and flushes them to the DB.""" + + pipe = redis.pipeline() + + pipe.smembers(constants.CALL_IDS_TO_FLUSH) + pipe.delete(constants.CALL_IDS_TO_FLUSH) + + call_ids, _ = await pipe.execute() + + # Redis returns bytes. + call_ids = [int(call_id) for call_id in call_ids] + + pipe = redis.pipeline() + + if len(call_ids) == 0: + return + + for call_id in call_ids: + pipe.hgetall(constants.CALL_MESSAGES_TEMPLATE % call_id) + pipe.hgetall(constants.CALL_MESSAGE_CONTENTS_TEMPLATE % call_id) + + for call_id, (messages, message_contents) in zip( + call_ids, utils.chunks(await pipe.execute(), 2) + ): + message_rows: list[models.Message] = [] + + for message_json in messages.values(): + message_rows.append(models.Message(**orjson.loads(message_json))) + + message_content_rows: list[models.MessageContent] = [] + + for message_content_json in message_contents.values(): + message_content_rows.append( + models.MessageContent(**orjson.loads(message_content_json)) + ) + + message_rows.sort(key=lambda m: m.id) + message_content_rows.sort(key=lambda mc: mc.message_id) + + pipe = redis.pipeline() + + pipe.delete( + constants.CALL_MESSAGE_CONTENT_ITERATIONS_TEMPLATE % call_id, + constants.CALL_MESSAGE_CONTENTS_TEMPLATE % call_id, + constants.CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE % call_id, + constants.CALL_MESSAGES_TEMPLATE % call_id, + ) + + await pipe.execute() + + for _ in range(3): + try: + await models.Message.bulk_create(message_rows, batch_size=500) + break + + except asyncpg.exceptions.DeadlockDetectedError: + continue + + for _ in range(3): + try: + await models.MessageContent.bulk_create( + message_content_rows, batch_size=500 + ) + break + + except asyncpg.exceptions.DeadlockDetectedError: + continue + @commands.command(aliases=["c"]) @commands.guild_only() @checks.is_unbanned() @@ -322,6 +345,8 @@ class Calls(base.Cog): "channels" ) + await redis.sadd(constants.CALL_IDS_TO_FLUSH, call.id) + call.ended_at = datetime.now(timezone.utc) await call.save() @@ -501,6 +526,8 @@ class Calls(base.Cog): "channels" ) + await redis.sadd(constants.CALL_IDS_TO_FLUSH, call.id) + other_channel: models.Channel = find( lambda c: c.id != channel_id, call.channels ) @@ -606,6 +633,8 @@ class Calls(base.Cog): call.ended_at = datetime.now(timezone.utc) await call.save() + await redis.sadd(constants.CALL_IDS_TO_FLUSH, call.id) + updated_channels = [] for channel in call.channels: @@ -799,7 +828,7 @@ class Calls(base.Cog): pipe = redis.pipeline() pipe.hset( - constants.MESSAGES_QUEUE, + constants.CALL_MESSAGES_TEMPLATE % call.id, message.id, orjson.dumps( { @@ -819,8 +848,8 @@ class Calls(base.Cog): ), ) pipe.hset( - constants.MESSAGE_CONTENTS_QUEUE, - message.id, + constants.CALL_MESSAGE_CONTENTS_TEMPLATE % call.id, + constants.CALL_MESSAGE_CONTENTS_HASH_TEMPLATE % (message.id, 1), orjson.dumps( { "message_id": message.id, @@ -830,12 +859,15 @@ class Calls(base.Cog): } ), ) + pipe.hset( + constants.CALL_MESSAGE_CONTENT_ITERATIONS_TEMPLATE % call.id, message.id, 1 + ) if webhook_message is not None: pipe.hset( - constants.MESSAGES_PROPAGATION_MAP, + constants.CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE % call.id, webhook_message.id, - message.id + message.id, ) await pipe.execute() diff --git a/payphone/cogs/calls/embeds.py b/payphone/cogs/calls/embeds.py index 53298d3..435eab9 100644 --- a/payphone/cogs/calls/embeds.py +++ b/payphone/cogs/calls/embeds.py @@ -1,7 +1,5 @@ import discord -from tortoise import expressions - import orjson from payphone.db import models @@ -25,51 +23,44 @@ async def ReplyEmbed( message = None - unflushed_propagation_to_message_id = await redis.hget( - constants.MESSAGES_PROPAGATION_MAP, reference.message_id + unflushed_message = await redis.hget( + constants.CALL_MESSAGES_TEMPLATE % call_id, reference.message_id ) - unflushed_message = None - unflushed_message_content = None - last_content: models.MessageContent | None = None - if unflushed_propagation_to_message_id is not None: - unflushed_message = await redis.hget( - constants.MESSAGES_QUEUE, unflushed_propagation_to_message_id + if unflushed_message is None: + message_id = await redis.hget( + constants.CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE % call_id, + reference.message_id, ) - if unflushed_message is not None: - unflushed_message_content = await redis.hget( - constants.MESSAGE_CONTENTS_QUEUE, unflushed_propagation_to_message_id + print(constants.CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE % call_id) + + if message_id is not None: + unflushed_message = await redis.hget( + constants.CALL_MESSAGES_TEMPLATE % call_id, message_id ) - else: - unflushed_message = await redis.hget( - constants.MESSAGES_QUEUE, reference.message_id - ) + if unflushed_message is None: + return (None, None) - if unflushed_message is not None: - unflushed_message_content = await redis.hget( - constants.MESSAGE_CONTENTS_QUEUE, reference.message_id - ) + message = models.Message(**orjson.loads(unflushed_message)) - if unflushed_message is not None: - message = models.Message(**orjson.loads(unflushed_message)) - last_content = models.MessageContent(**orjson.loads(unflushed_message_content)) + unflushed_message_content_iteration = int(await redis.hget( + constants.CALL_MESSAGE_CONTENT_ITERATIONS_TEMPLATE % call_id, message.id + )) + unflushed_message_content = await redis.hget( + constants.CALL_MESSAGE_CONTENTS_TEMPLATE % call_id, + constants.CALL_MESSAGE_CONTENTS_HASH_TEMPLATE + % (message.id, unflushed_message_content_iteration), + ) - message.channel = await models.Channel.get(id=message.channel_id) - message.profile = await models.Profile.get(id=message.profile_id) - message.propagation_channel = await models.Channel.get(id=message.propagation_channel_id) + last_content = models.MessageContent(**orjson.loads(unflushed_message_content)) - else: - message = ( - await models.Message.get_or_none( - expressions.Q(id=reference.message_id) - | expressions.Q(propagation_id=reference.message_id), - call_id=call_id, - ) - .prefetch_related("contents") - .select_related("profile", "channel", "propagation_channel") - ) + message.channel = await models.Channel.get(id=message.channel_id) + message.profile = await models.Profile.get(id=message.profile_id) + message.propagation_channel = await models.Channel.get( + id=message.propagation_channel_id + ) if message is None: return None, None diff --git a/payphone/cogs/mods/cog.py b/payphone/cogs/mods/cog.py index 98cf14c..08e5991 100644 --- a/payphone/cogs/mods/cog.py +++ b/payphone/cogs/mods/cog.py @@ -289,7 +289,17 @@ class Mods(base.Cog): except errors.DiscordException: pass - await ctx.send("⚡ Payphone was shut down.") + await ctx.send( + "⚡ Notified all channels of calls hang up. Flushing messages to the DB..." + ) + + calls = await models.Call.filter(ended_at=None) + + await redis.sadd(constants.CALL_IDS_TO_FLUSH, *(call.id for call in calls)) + + await self.bot.get_cog("Calls").flush_messages() + + await ctx.send("⚡ Messages were flushed, payphone will shut down.") sys.exit() diff --git a/payphone/core/constants.py b/payphone/core/constants.py index 31a4789..90b24cb 100644 --- a/payphone/core/constants.py +++ b/payphone/core/constants.py @@ -173,6 +173,25 @@ CHANNELS_ONCALL_SET = "channels:on-call" MESSAGES_QUEUE = "messages" MESSAGE_CONTENTS_QUEUE = "message-contents" MESSAGES_PROPAGATION_MAP = "messages-propagation" + +CALL_MESSAGES_TEMPLATE = "m:%s" +"""`%s`: Call ID. The key template for call messages hash.""" + +CALL_MESSAGE_CONTENTS_TEMPLATE = "mc:%s" +"""`%s`: Call ID. The key template for call message contents hash.""" + +CALL_MESSAGE_CONTENTS_HASH_TEMPLATE = "%s-%s" +"""`%s-%s`: Call ID - Iteration. The key template for call message contents hash keys.""" + +CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE = "mfp:%s" +"""`%s`: Call ID. A map from propagation IDs to message IDs.""" + +CALL_MESSAGE_CONTENT_ITERATIONS_TEMPLATE = "mci:%s" +"""`%s`: Call ID. A map of messages to amount of iterations of their message contents.""" + +CALL_IDS_TO_FLUSH = "call-ids-to-flush" +"""A set of IDs of calls to flush to the DB.""" + BAN_TEMPLATE = "b:%s" PREFIX_TEMPLATE = "p:%s" diff --git a/payphone/core/utils.py b/payphone/core/utils.py index 04cdae0..dd2bffd 100644 --- a/payphone/core/utils.py +++ b/payphone/core/utils.py @@ -1,4 +1,4 @@ -from typing import Callable, Iterable, TypeVar +from typing import Callable, Iterable, TypeVar, Generator T = TypeVar("T") @@ -28,3 +28,30 @@ def deduplicate(i: Iterable[T], key: Callable[[T], U] = lambda i: i) -> list[T]: found_keys.add(k) return result + + +def chunks(i: Iterable[T], size: int) -> Generator[list[T], None, None]: + """Chunk an iterable. + + Args: + i (Iterable[T]): The iterable. + size (int): The size of each chunk. + + Returns: + Generator[list[T], None, None]: The iterable's chunks. + """ + + if size <= 0: + raise ValueError("Chunks size can't be lower than 1.") + + chunk = [] + + for item in i: + chunk.append(item) + + if len(chunk) == size: + yield chunk + chunk = [] + + if len(chunk) > 0: + yield chunk