Merge pull request #56 from Payphone-Bot/53-move-messages-flushing-away-from-bot-processes

Move messages flushing to a separate process.

.env.example

@@ -5,7 +5,3 @@ IS_DEBUG=false
 SUGGESTIONS_CHANNEL_ID=
 
 SERVER_INVITE=
-
-# The CPU cores assigned to each container, e.g. "0,1" for the first two cores.
-# CPUSET_BOT=
-# CPUSET_DATABASE=

compose.yml

@@ -3,7 +3,6 @@ name: "payphone"
 services:
   bot:
     build: .
-    cpuset: ${CPUSET_BOT:-0,1}
     env_file:
       - .env
     environment:
@@ -15,7 +14,6 @@ services:
 
   database:
     image: postgres:16.3
-    cpuset: ${CPUSET_DATABASE:-2,3}
     environment:
       - POSTGRES_DB=payphone
       - POSTGRES_USER=payphone

payphone/cogs/calls/cog.py

@@ -30,11 +30,6 @@ 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.add_exception_type(RuntimeError)
-        self.flush_messages.start()
-
     def cog_unload(self) -> None:
         self.hangup_calls.cancel()
 
@@ -86,89 +81,6 @@ class Calls(base.Cog):
             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, all_calls: bool = False) -> None:
-        """Fetches all the pending messages from Redis and flushes them to the DB.
-
-        Args:
-            all_calls (bool, optional): Whether to flush the messages for all calls of all shards.
-        """
-
-        call_ids = await redis.smembers(constants.CALL_IDS_TO_FLUSH)
-
-        # Redis returns bytes.
-        call_ids = [
-            int(call_id)
-            for call_id in call_ids
-            if all_calls
-            or (int(call_id) % self.bot.shard_count) - self.bot.shard_ids[0] == 0
-        ]
-
-        if len(call_ids) == 0:
-            return
-
-        await redis.srem(constants.CALL_IDS_TO_FLUSH, *call_ids)
-
-        pipe = redis.pipeline()
-
-        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,
-                constants.CALL_MESSAGE_REPLIES_TEMPLATE % call_id,
-            )
-
-            await pipe.execute()
-
-            for _ in range(3):
-                try:
-                    await models.Message.bulk_create(message_rows, batch_size=20)
-                    break
-
-                except (
-                    asyncpg.exceptions.PostgresError,
-                    asyncpg.exceptions.InterfaceError,
-                ):
-                    continue
-
-            for _ in range(3):
-                try:
-                    await models.MessageContent.bulk_create(
-                        message_content_rows, batch_size=20
-                    )
-                    break
-
-                except (
-                    asyncpg.exceptions.PostgresError,
-                    asyncpg.exceptions.InterfaceError,
-                ):
-                    continue
-
     @commands.command(aliases=["c"])
     @commands.guild_only()
     @checks.is_unbanned()

payphone/cogs/mods/cog.py

@@ -12,6 +12,8 @@ from discord.ext import commands
 import tortoise
 from tortoise import transactions, expressions
 
+from procs import flush_messages
+
 from payphone import bot
 from payphone.db import utils, models
 from payphone.cogs import base
@@ -306,9 +308,7 @@ class Mods(base.Cog):
             ended_at=datetime.now(timezone.utc)
         )
 
-        await self.bot.get_cog("Calls").flush_messages(True)
+        await flush_messages.flush()
-
-        await tortoise.Tortoise.close_connections()
 
         await ctx.send("⚡ Messages were flushed, payphone will shut down.")
 

procs/flush_messages.py

@@ -0,0 +1,110 @@
+async def flush():
+    import orjson
+
+    import asyncpg.exceptions
+
+    from payphone.db import models
+
+    from payphone.core import constants, utils
+    from payphone.core.redis import redis
+
+    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]
+
+    if len(call_ids) == 0:
+        print("Skipping flushing. No calls to flush.")
+        return
+
+    print(f"Flushing messages for {len(call_ids)} calls.")
+
+    pipe = redis.pipeline()
+
+    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,
+            constants.CALL_MESSAGE_REPLIES_TEMPLATE % call_id,
+        )
+
+        await pipe.execute()
+
+        for _ in range(5):
+            try:
+                await models.Message.bulk_create(message_rows, batch_size=20)
+                break
+
+            except (
+                asyncpg.exceptions.PostgresError,
+                asyncpg.exceptions.InterfaceError,
+            ) as e:
+                print(
+                    f"[Call ID: {call_id}] An error occured trying to flush messages for call ID. {e}"
+                )
+                continue
+
+        for _ in range(5):
+            try:
+                await models.MessageContent.bulk_create(
+                    message_content_rows, batch_size=20
+                )
+                break
+
+            except (
+                asyncpg.exceptions.PostgresError,
+                asyncpg.exceptions.InterfaceError,
+            ) as e:
+                print(
+                    f"[Call ID: {call_id}] An error occured trying to flush message contents. {e}"
+                )
+                continue
+
+        print(f"[Call ID: {call_id}] Flushed {len(messages)} messages and {len(message_contents)} message contents.")
+
+    print("Finished flushing messages.")
+
+
+async def loop() -> None:
+    """Loops in the background flushing messages to the DB."""
+
+    import asyncio
+
+    from payphone.db import setup
+
+    await setup.init()
+
+    while True:
+        await asyncio.sleep(30)
+
+        await flush()

run.py

@@ -18,6 +18,14 @@ def shard(shard_id: int, shard_count: int) -> None:
     bot.run(config.TOKEN)
 
 
+def flush_process():
+    import asyncio
+
+    from procs.flush_messages import loop
+
+    asyncio.run(loop())
+
+
 if __name__ == "__main__":
     import multiprocessing as mp
 
@@ -61,17 +69,29 @@ if __name__ == "__main__":
     ps: list[mp.Process] = [
         mp.Process(target=shard, args=(i, shard_count)) for i in range(shard_count)
     ]
+    flushing_p = mp.Process(target=flush_process)
 
     for i, p in enumerate(ps):
         p.start()
 
         print(
-            f"Shard {i}'s process started." + (" Waiting 5 seconds start initialize next shard..." if len(ps) - 1 < i else " All shards were started.")
+            f"Shard {i}'s process started."
+            + (
+                " Waiting 5 seconds start initialize next shard..."
+                if len(ps) - 1 < i
+                else " All shards were started."
+            )
         )
 
         if i + 1 != len(ps):
             time.sleep(5)
 
+    print("Starting messages flushing process...")
+
+    flushing_p.start()
+
+    print("Messages flushing process started.")
+
     try:
         while True:
             for i in range(shard_count):
@@ -86,6 +106,14 @@ if __name__ == "__main__":
                     )
                     time.sleep(5)
 
+            if not flushing_p.is_alive():
+                print("Messages flushing process was stopped. Restarting it...")
+
+                flushing_p = mp.Process(target=flush_process)
+                flushing_p.start()
+
+                print("Messages flushing process started.")
+
             time.sleep(5)
 
     except KeyboardInterrupt:
@@ -95,9 +123,15 @@ if __name__ == "__main__":
             p.terminate()
             print(f"Terminated shard {i}.")
 
+        print("Stopping messages flushing process...")
+
+        flushing_p.terminate()
+
         print("Checking that all processes have finished...")
 
         for p in ps:
             p.join()
 
+        flushing_p.join()
+
         print("All processes have been terminated.")