From b5872df03419c40eeed4ceb3b844cf2befb3d94a Mon Sep 17 00:00:00 2001 From: "Nyeki.py" <84998222+Nekidev@users.noreply.github.com> Date: Thu, 24 Oct 2024 00:22:07 -0300 Subject: [PATCH] feat(bot): :sparkles: Terminate all processes after `p.shutdown`. --- payphone/bot.py | 16 +++++++++++++++- payphone/cogs/calls/cog.py | 9 +++++---- payphone/cogs/mods/cog.py | 10 +++++++--- run.py | 37 ++++++++++++++++++++++++++++++------- 4 files changed, 57 insertions(+), 15 deletions(-) diff --git a/payphone/bot.py b/payphone/bot.py index 8cfefb7..2bd2888 100644 --- a/payphone/bot.py +++ b/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) diff --git a/payphone/cogs/calls/cog.py b/payphone/cogs/calls/cog.py index 8dd9ea8..a9a478e 100644 --- a/payphone/cogs/calls/cog.py +++ b/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 diff --git a/payphone/cogs/mods/cog.py b/payphone/cogs/mods/cog.py index 1ca67b3..3428224 100644 --- a/payphone/cogs/mods/cog.py +++ b/payphone/cogs/mods/cog.py @@ -243,7 +243,9 @@ class Mods(base.Cog): - `[]`: 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() diff --git a/run.py b/run.py index e0f7e42..08b771e 100644 --- a/run.py +++ b/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()