feat(bot): ✨ Terminate all processes after p.shutdown.

payphone/bot.py

@@ -2,6 +2,8 @@ from typing import Literal
 
 from datetime import datetime, timezone, timedelta
 
+from multiprocessing import managers
+
 import sys
 
 import discord
@@ -20,6 +22,13 @@ from payphone.core import constants, exceptions, formatting, help_command, confi
 
 
 class Payphone(commands.AutoShardedBot):
+    shutdown: managers.ValueProxy[bool]
+    """A value proxy that connects with all other shard processes.
+    
+    This value is shared between shard processes, and if any of them sets this proxy's value to
+    `True`, then all shard processes will be terminated by the main process (`run.py`).
+    """
+
     redis: aredis.Redis = aredis.Redis(host="localhost" if config.IS_DEBUG else "redis")
 
     def __init__(self) -> None:
@@ -398,7 +407,12 @@ class Payphone(commands.AutoShardedBot):
     ) -> None:
         _, exception, _ = sys.exc_info()
 
-        if isinstance(exception, asyncpg.exceptions.InterfaceError):
+        if isinstance(
+            exception,
+            (asyncpg.exceptions.InterfaceError, asyncpg.exceptions.PostgresError),
+        ):
+            # IMPORTANT! These are just spam at the moment, but may cause trouble in the future as
+            # the bot grows.
             return
 
         return await super().on_error(event_method, *args, **kwargs)

payphone/cogs/calls/cog.py

@@ -9,8 +9,6 @@ from discord.utils import find
 
 from tortoise import transactions, exceptions
 
-import asyncpg.exceptions
-
 import aiohttp
 
 import orjson
@@ -19,7 +17,7 @@ from payphone import bot
 from payphone.db import utils as dbutils, models
 from payphone.cogs import base
 from payphone.cogs.calls import embeds
-from payphone.core import constants, checks, censorship, formatting, utils
+from payphone.core import constants, checks, censorship, formatting
 
 
 class Calls(base.Cog):
@@ -877,7 +875,10 @@ class Calls(base.Cog):
             orjson.dumps(message_dict),
         )
 
-        other_channel = await models.Channel.get(id=message.propagation_channel_id)
+        other_channel = await models.Channel.get_or_none(id=message.propagation_channel_id)
+
+        if other_channel is None:
+            return
 
         if other_channel.call_id != message.call_id:
             return

payphone/cogs/mods/cog.py

@@ -243,7 +243,9 @@ class Mods(base.Cog):
         - `[<reason>]`: The reason of the shutdown.
         """
 
-        oncall_channel_ids = list(await self.bot.redis.smembers(constants.CHANNELS_ONCALL_SET))
+        oncall_channel_ids = list(
+            await self.bot.redis.smembers(constants.CHANNELS_ONCALL_SET)
+        )
         searching_channel_ids = list(
             await self.bot.redis.smembers(constants.CHANNELS_SEARCH_QUEUE)
         )
@@ -270,7 +272,9 @@ class Mods(base.Cog):
 
         await asyncio.sleep(seconds)
 
-        oncall_channel_ids = list(await self.bot.redis.smembers(constants.CHANNELS_ONCALL_SET))
+        oncall_channel_ids = list(
+            await self.bot.redis.smembers(constants.CHANNELS_ONCALL_SET)
+        )
         searching_channel_ids = list(
             await self.bot.redis.smembers(constants.CHANNELS_SEARCH_QUEUE)
         )
@@ -310,7 +314,7 @@ class Mods(base.Cog):
 
         await ctx.send("⚡ Messages were flushed, payphone will shut down.")
 
-        sys.exit()
+        self.bot.shutdown.value = True
 
     @commands.command()
     @commands.is_owner()

run.py

@@ -1,12 +1,15 @@
+from multiprocessing import managers
+
 from datetime import datetime, timezone
 
 
-def shard(shard_id: int, shard_count: int) -> None:
+def shard(shard_id: int, shard_count: int, shutdown: managers.ValueProxy[bool]) -> None:
     """Starts a new shard.
 
     Args:
         shard_id (int): The shard's ID.
         shard_count (int): The total amount of shards.
+        shutdown (ValueProxy[bool]): The flag for the main process to shut down all shards.
     """
 
     from payphone.bot import bot
@@ -14,6 +17,7 @@ def shard(shard_id: int, shard_count: int) -> None:
 
     bot.shard_ids = [shard_id]
     bot.shard_count = shard_count
+    bot.shutdown = shutdown
 
     try:
         import uvloop  # pylint: disable=import-error
@@ -46,6 +50,10 @@ def flush_process():
         asyncio.run(loop())
 
 
+class ShutdownException(Exception):
+    """Raised when a shard has triggered the shutdown of all shards."""
+
+
 if __name__ == "__main__":
     import multiprocessing as mp
 
@@ -87,10 +95,14 @@ if __name__ == "__main__":
 
     shard_count = info["shards"]
 
+    manager = mp.Manager()
+    shutdown = manager.Value("b", False)
+
     print(f"Starting {shard_count} shard{"s" if shard_count != 1 else ""}.")
 
     ps: list[mp.Process] = [
-        mp.Process(target=shard, args=(i, shard_count)) for i in range(shard_count)
+        mp.Process(target=shard, args=(i, shard_count, shutdown))
+        for i in range(shard_count)
     ]
     flushing_p = mp.Process(target=flush_process)
 
@@ -101,12 +113,12 @@ if __name__ == "__main__":
             f"Shard {i}'s process started."
             + (
                 " Waiting 5 seconds start initialize next shard..."
-                if len(ps) - 1 < i
+                if i != len(ps) - 1
                 else " All shards were started."
             )
         )
 
-        if i + 1 != len(ps):
+        if i != len(ps) - 1:
             time.sleep(5)
 
     print("Starting messages flushing process...")
@@ -117,11 +129,16 @@ if __name__ == "__main__":
 
     try:
         while True:
+            if shutdown.value:
+                raise ShutdownException()
+
             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] = mp.Process(
+                        target=shard, args=(i, shard_count), name=f"payphone-shard-{i}"
+                    )
                     ps[i].start()
 
                     print(
@@ -132,7 +149,7 @@ if __name__ == "__main__":
             if not flushing_p.is_alive():
                 print("Messages flushing process was stopped. Restarting it...")
 
-                flushing_p = mp.Process(target=flush_process)
+                flushing_p = mp.Process(target=flush_process, name="payphone-flushing")
                 flushing_p.start()
 
                 print("Messages flushing process started.")
@@ -140,7 +157,13 @@ if __name__ == "__main__":
             time.sleep(5)
 
     except KeyboardInterrupt:
-        print(f"Stopping {shard_count} shards...")
+        print("Ctrl-C was pressed.")
+
+    except ShutdownException:
+        print("A shard shut down the bot.")
+
+    finally:
+        print(f"Terminating {shard_count} shards...")
 
         for i, p in enumerate(ps):
             p.terminate()