nyeki/payphone
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(bot): ✨ Move messages flushing to a separate process.

Browse Source
This commit is contained in:
Nyeki.py
2024-10-23 21:49:03 -03:00
parent 00404973bf
commit 2ff7021595
7 changed files with 149 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.env.example
-4
View File
@@ -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=
compose.yml
-2
View File
@@ -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
payphone/cogs/calls/cog.py
-88
View File
@@ -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()
payphone/cogs/mods/cog.py
+3 -3
View File
@@ -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.")
procs/__init__.py
View File
procs/flush_messages.py
+110
View File
@@ -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()
run.py
+35 -1
View File
@@ -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.")