feat(bot): ✨ Flush all messages to DB on call hangup.

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,
+            constants.CALL_MESSAGE_CONTENTS_TEMPLATE % call.id,
-            message.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()

payphone/cogs/calls/embeds.py

@@ -1,7 +1,5 @@
 import discord
 
-from tortoise import expressions
-
 import orjson
 
 from payphone.db import models
@@ -25,50 +23,43 @@ async def ReplyEmbed(
 
     message = None
 
-    unflushed_propagation_to_message_id = await redis.hget(
-        constants.MESSAGES_PROPAGATION_MAP, 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
+        constants.CALL_MESSAGES_TEMPLATE % call_id, reference.message_id
     )
 
-        if unflushed_message is not None:
+    if unflushed_message is None:
-            unflushed_message_content = await redis.hget(
+        message_id = await redis.hget(
-                constants.MESSAGE_CONTENTS_QUEUE, unflushed_propagation_to_message_id
+            constants.CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE % call_id,
+            reference.message_id,
         )
 
-    else:
+        print(constants.CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE % call_id)
+
+        if message_id is not None:
             unflushed_message = await redis.hget(
-            constants.MESSAGES_QUEUE, reference.message_id
+                constants.CALL_MESSAGES_TEMPLATE % call_id, message_id
             )
 
-        if unflushed_message is not None:
+    if unflushed_message is None:
-            unflushed_message_content = await redis.hget(
+        return (None, None)
-                constants.MESSAGE_CONTENTS_QUEUE, reference.message_id
-            )
 
-    if unflushed_message is not None:
     message = models.Message(**orjson.loads(unflushed_message))
+
+    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),
+    )
+
     last_content = models.MessageContent(**orjson.loads(unflushed_message_content))
 
     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)
+    message.propagation_channel = await models.Channel.get(
-
+        id=message.propagation_channel_id
-    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")
     )
 
     if message is None:

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()
 

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"
 

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