perf(bot): ⚡ Distribute messages flushing between shards.

This commit is contained in:
Nyeki.py
2024-10-22 17:20:24 -03:00
parent 0fa5f73863
commit edae71827b
+18 -12
View File
@@ -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