Merge pull request #43 from Payphone-Bot/38-distribute-call-messages-flushing-among-shard-process-tasks

Distribute call messages flushing among shard process tasks.

payphone/bot.py

@@ -64,7 +64,7 @@ class Payphone(commands.AutoShardedBot):
             shard_id=shard_id,
         )
 
-        print(f"Shard is ready! Shard ID: {shard_id}")
+        print(f"Shard {shard_id} is ready!")
 
     async def on_shard_resumed(self, shard_id: int) -> None:
         """Sets the activity for the shard.
@@ -81,7 +81,7 @@ class Payphone(commands.AutoShardedBot):
             shard_id=shard_id,
         )
 
-        print(f"Shard was resumed! Shard ID: {shard_id}")
+        print(f"Shard {shard_id} was resumed!")
 
     async def get_or_fetch_channel(
         self, channel_id: int

payphone/cogs/calls/cog.py

@@ -86,24 +86,30 @@ class Calls(base.Cog):
             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."""
+    async def flush_messages(self, all_calls: bool = False) -> None:
+        """Fetches all the pending messages from Redis and flushes them to the DB.
 
-        pipe = redis.pipeline()
+        Args:
+            all_calls (bool, optional): Whether to flush the messages for all calls of all shards.
+        """
 
-        pipe.smembers(constants.CALL_IDS_TO_FLUSH)
-        pipe.delete(constants.CALL_IDS_TO_FLUSH)
-
-        call_ids, _ = await pipe.execute()
+        call_ids = await redis.smembers(constants.CALL_IDS_TO_FLUSH)
 
         # Redis returns bytes.
-        call_ids = [int(call_id) for call_id in call_ids]
-
-        pipe = redis.pipeline()
+        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)
@@ -140,7 +146,7 @@ class Calls(base.Cog):
 
             for _ in range(3):
                 try:
-                    await models.Message.bulk_create(message_rows, batch_size=500)
+                    await models.Message.bulk_create(message_rows, batch_size=20)
                     break
 
                 except (
@@ -152,7 +158,7 @@ class Calls(base.Cog):
             for _ in range(3):
                 try:
                     await models.MessageContent.bulk_create(
-                        message_content_rows, batch_size=500
+                        message_content_rows, batch_size=20
                     )
                     break
 

payphone/cogs/mods/cog.py

@@ -407,6 +407,7 @@ class Mods(base.Cog):
                 await asyncio.sleep(wait)
 
     @commands.command()
+    @checks.is_moderator()
     async def user(self, ctx: commands.Context, user: converters.DiscordUserID) -> None:
         """
         Shows all the information about a user.

payphone/db/setup.py

@@ -1,8 +1,5 @@
-from datetime import datetime, timezone
-
 import tortoise
 
-from payphone.db import models
 from payphone.core import config
 
 
@@ -10,8 +7,3 @@ async def init() -> None:
     """Initialize the database and it's connections."""
 
     await tortoise.Tortoise.init(config.TORTOISE_ORM, use_tz=True)
-
-    await tortoise.Tortoise.generate_schemas(safe=True)
-
-    await models.Call.filter(ended_at=None).update(ended_at=datetime.now(timezone.utc))
-    await models.Channel.exclude(call_id=None).update(call_id=None)

run.py

@@ -1,3 +1,6 @@
+from datetime import datetime, timezone
+
+
 def shard(shard_id: int, shard_count: int) -> None:
     """Starts a new shard.
 
@@ -20,9 +23,29 @@ if __name__ == "__main__":
 
     import time
 
+    from tortoise import run_async
+
     import httpx
 
+    from payphone.db import setup, models
     from payphone.core import config
+    from payphone.core.redis import redis
+
+    async def init():
+        await setup.init()
+
+        await models.Call.filter(ended_at=None).update(
+            ended_at=datetime.now(timezone.utc)
+        )
+        await models.Channel.exclude(call_id=None).update(call_id=None)
+
+        await redis.flushall()
+
+    print("Running DB and Redis cleanup...")
+
+    run_async(init())
+
+    print("Cleanup done! Loading recommended shard amount...")
 
     response = httpx.get(
         "https://discord.com/api/v10/gateway/bot",
@@ -39,15 +62,42 @@ if __name__ == "__main__":
         mp.Process(target=shard, args=(i, shard_count)) for i in range(shard_count)
     ]
 
-    for p in ps:
+    for i, p in enumerate(ps):
         p.start()
 
-        time.sleep(5)
+        print(
+            f"Shard {i}'s process started." + (" Waiting 5 seconds start initialize next shard..." if len(ps) - 1 < i else " All shards were started.")
+        )
 
-    while True:
-        for i in range(shard_count):
-            if not ps[i].is_alive():
-                ps[i] = mp.Process(target=shard, args=(i, shard_count))
-                ps[i].start()
+        if i + 1 != len(ps):
+            time.sleep(5)
 
-        time.sleep(5)
+    try:
+        while True:
+            for i in range(shard_count):
+                if not ps[i].is_alive():
+                    print(f"Shard {i}'s process exited. Restarting the shard...")
+
+                    ps[i] = mp.Process(target=shard, args=(i, shard_count))
+                    ps[i].start()
+
+                    print(
+                        f"Shard {i} was restarted. Waiting 5 seconds to check next shard..."
+                    )
+                    time.sleep(5)
+
+            time.sleep(5)
+
+    except KeyboardInterrupt:
+        print(f"Stopping {shard_count} shards...")
+
+        for i, p in enumerate(ps):
+            p.terminate()
+            print(f"Terminated shard {i}.")
+
+        print("Checking that all processes have finished...")
+        
+        for p in ps:
+            p.join()
+
+        print("All processes have been terminated.")