diff --git a/payphone/bot.py b/payphone/bot.py index 1e21183..808757d 100644 --- a/payphone/bot.py +++ b/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 diff --git a/payphone/cogs/calls/cog.py b/payphone/cogs/calls/cog.py index 3716410..339bb43 100644 --- a/payphone/cogs/calls/cog.py +++ b/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 diff --git a/payphone/cogs/mods/cog.py b/payphone/cogs/mods/cog.py index fbb2641..17a93f4 100644 --- a/payphone/cogs/mods/cog.py +++ b/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. diff --git a/payphone/db/setup.py b/payphone/db/setup.py index dcf9400..40bad08 100644 --- a/payphone/db/setup.py +++ b/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) diff --git a/run.py b/run.py index 971fe17..ed1def5 100644 --- a/run.py +++ b/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.")