diff --git a/.env.example b/.env.example index c0be9a0..ff6cfc6 100644 --- a/.env.example +++ b/.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= diff --git a/compose.yml b/compose.yml index f50fefd..7d7e8f8 100644 --- a/compose.yml +++ b/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 diff --git a/payphone/cogs/calls/cog.py b/payphone/cogs/calls/cog.py index 9ba3e47..97a6782 100644 --- a/payphone/cogs/calls/cog.py +++ b/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() diff --git a/payphone/cogs/mods/cog.py b/payphone/cogs/mods/cog.py index e7aff61..f3a808a 100644 --- a/payphone/cogs/mods/cog.py +++ b/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 tortoise.Tortoise.close_connections() + await flush_messages.flush() await ctx.send("⚡ Messages were flushed, payphone will shut down.") diff --git a/procs/__init__.py b/procs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/procs/flush_messages.py b/procs/flush_messages.py new file mode 100644 index 0000000..c61543f --- /dev/null +++ b/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() diff --git a/run.py b/run.py index ed1def5..12d82bb 100644 --- a/run.py +++ b/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.")