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