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

feat(bot): ✨ Flush all messages to DB on call hangup.

Browse Source
This commit is contained in:
Nyeki.py
2024-10-20 21:51:48 -03:00
parent e3c58139d2
commit 6fbb65507f
5 changed files with 172 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
payphone/cogs/calls/cog.py
+86 -54
View File
@@ -30,62 +30,15 @@ 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.start()
def cog_unload(self) -> None:
self.hangup_calls.cancel()
self.flush_messages.cancel()
return super().cog_unload()
@tasks.loop(seconds=5)
async def flush_messages(self) -> None:
"""Flush all the in-memory messages and message contents to the DB."""
pipe = redis.pipeline()
pipe.hgetall(constants.MESSAGES_QUEUE)
pipe.hgetall(constants.MESSAGE_CONTENTS_QUEUE)
pipe.delete(
constants.MESSAGES_QUEUE,
constants.MESSAGE_CONTENTS_QUEUE,
constants.MESSAGES_PROPAGATION_MAP,
)
messages, message_contents, _ = await pipe.execute()
messages = sorted(
[models.Message(**orjson.loads(msg)) for msg in messages.values()],
key=lambda m: m.created_at,
)
message_contents = sorted(
[
models.MessageContent(**orjson.loads(msg))
for msg in message_contents.values()
],
key=lambda m: m.created_at,
)
while True:
try:
await models.Message.bulk_create(messages, batch_size=500)
break
except asyncpg.exceptions.DeadlockDetectedError:
continue
while True:
try:
await models.MessageContent.bulk_create(
message_contents, batch_size=500
)
break
except asyncpg.exceptions.DeadlockDetectedError:
continue
@tasks.loop(seconds=15)
async def hangup_calls(self) -> None:
"""Hangs up inactive calls.
@@ -129,9 +82,79 @@ class Calls(base.Cog):
hung_up_channel_ids = [id for id in hung_up_channel_ids if id is not None]
if len(hung_up_channel_ids) > 0:
await models.Channel.filter(id__in=hung_up_channel_ids).update(call=None)
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) -> None:
"""Fetches all the pending messages from Redis and flushes them to the DB."""
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]
pipe = redis.pipeline()
if len(call_ids) == 0:
return
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,
)
await pipe.execute()
for _ in range(3):
try:
await models.Message.bulk_create(message_rows, batch_size=500)
break
except asyncpg.exceptions.DeadlockDetectedError:
continue
for _ in range(3):
try:
await models.MessageContent.bulk_create(
message_content_rows, batch_size=500
)
break
except asyncpg.exceptions.DeadlockDetectedError:
continue
@commands.command(aliases=["c"])
@commands.guild_only()
@checks.is_unbanned()
@@ -322,6 +345,8 @@ class Calls(base.Cog):
"channels"
)
await redis.sadd(constants.CALL_IDS_TO_FLUSH, call.id)
call.ended_at = datetime.now(timezone.utc)
await call.save()
@@ -501,6 +526,8 @@ class Calls(base.Cog):
"channels"
)
await redis.sadd(constants.CALL_IDS_TO_FLUSH, call.id)
other_channel: models.Channel = find(
lambda c: c.id != channel_id, call.channels
)
@@ -606,6 +633,8 @@ class Calls(base.Cog):
call.ended_at = datetime.now(timezone.utc)
await call.save()
await redis.sadd(constants.CALL_IDS_TO_FLUSH, call.id)
updated_channels = []
for channel in call.channels:
@@ -799,7 +828,7 @@ class Calls(base.Cog):
pipe = redis.pipeline()
pipe.hset(
constants.MESSAGES_QUEUE,
constants.CALL_MESSAGES_TEMPLATE % call.id,
message.id,
orjson.dumps(
{
@@ -819,8 +848,8 @@ class Calls(base.Cog):
),
)
pipe.hset(
constants.MESSAGE_CONTENTS_QUEUE,
message.id,
constants.CALL_MESSAGE_CONTENTS_TEMPLATE % call.id,
constants.CALL_MESSAGE_CONTENTS_HASH_TEMPLATE % (message.id, 1),
orjson.dumps(
{
"message_id": message.id,
@@ -830,12 +859,15 @@ class Calls(base.Cog):
}
),
)
pipe.hset(
constants.CALL_MESSAGE_CONTENT_ITERATIONS_TEMPLATE % call.id, message.id, 1
)
if webhook_message is not None:
pipe.hset(
constants.MESSAGES_PROPAGATION_MAP,
constants.CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE % call.id,
webhook_message.id,
message.id
message.id,
)
await pipe.execute()
payphone/cogs/calls/embeds.py
+28 -37
View File
@@ -1,7 +1,5 @@
import discord
from tortoise import expressions
import orjson
from payphone.db import models
@@ -25,51 +23,44 @@ async def ReplyEmbed(
message = None
unflushed_propagation_to_message_id = await redis.hget(
constants.MESSAGES_PROPAGATION_MAP, reference.message_id
unflushed_message = await redis.hget(
constants.CALL_MESSAGES_TEMPLATE % call_id, reference.message_id
)
unflushed_message = None
unflushed_message_content = None
last_content: models.MessageContent | None = None
if unflushed_propagation_to_message_id is not None:
unflushed_message = await redis.hget(
constants.MESSAGES_QUEUE, unflushed_propagation_to_message_id
if unflushed_message is None:
message_id = await redis.hget(
constants.CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE % call_id,
reference.message_id,
)
if unflushed_message is not None:
unflushed_message_content = await redis.hget(
constants.MESSAGE_CONTENTS_QUEUE, unflushed_propagation_to_message_id
print(constants.CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE % call_id)
if message_id is not None:
unflushed_message = await redis.hget(
constants.CALL_MESSAGES_TEMPLATE % call_id, message_id
)
else:
unflushed_message = await redis.hget(
constants.MESSAGES_QUEUE, reference.message_id
)
if unflushed_message is None:
return (None, None)
if unflushed_message is not None:
unflushed_message_content = await redis.hget(
constants.MESSAGE_CONTENTS_QUEUE, reference.message_id
)
message = models.Message(**orjson.loads(unflushed_message))
if unflushed_message is not None:
message = models.Message(**orjson.loads(unflushed_message))
last_content = models.MessageContent(**orjson.loads(unflushed_message_content))
unflushed_message_content_iteration = int(await redis.hget(
constants.CALL_MESSAGE_CONTENT_ITERATIONS_TEMPLATE % call_id, message.id
))
unflushed_message_content = await redis.hget(
constants.CALL_MESSAGE_CONTENTS_TEMPLATE % call_id,
constants.CALL_MESSAGE_CONTENTS_HASH_TEMPLATE
% (message.id, unflushed_message_content_iteration),
)
message.channel = await models.Channel.get(id=message.channel_id)
message.profile = await models.Profile.get(id=message.profile_id)
message.propagation_channel = await models.Channel.get(id=message.propagation_channel_id)
last_content = models.MessageContent(**orjson.loads(unflushed_message_content))
else:
message = (
await models.Message.get_or_none(
expressions.Q(id=reference.message_id)
| expressions.Q(propagation_id=reference.message_id),
call_id=call_id,
)
.prefetch_related("contents")
.select_related("profile", "channel", "propagation_channel")
)
message.channel = await models.Channel.get(id=message.channel_id)
message.profile = await models.Profile.get(id=message.profile_id)
message.propagation_channel = await models.Channel.get(
id=message.propagation_channel_id
)
if message is None:
return None, None
payphone/cogs/mods/cog.py
+11 -1
View File
@@ -289,7 +289,17 @@ class Mods(base.Cog):
except errors.DiscordException:
pass
await ctx.send("⚡ Payphone was shut down.")
await ctx.send(
"⚡ Notified all channels of calls hang up. Flushing messages to the DB..."
)
calls = await models.Call.filter(ended_at=None)
await redis.sadd(constants.CALL_IDS_TO_FLUSH, *(call.id for call in calls))
await self.bot.get_cog("Calls").flush_messages()
await ctx.send("⚡ Messages were flushed, payphone will shut down.")
sys.exit()
payphone/core/constants.py
+19
View File
@@ -173,6 +173,25 @@ CHANNELS_ONCALL_SET = "channels:on-call"
MESSAGES_QUEUE = "messages"
MESSAGE_CONTENTS_QUEUE = "message-contents"
MESSAGES_PROPAGATION_MAP = "messages-propagation"
CALL_MESSAGES_TEMPLATE = "m:%s"
"""`%s`: Call ID. The key template for call messages hash."""
CALL_MESSAGE_CONTENTS_TEMPLATE = "mc:%s"
"""`%s`: Call ID. The key template for call message contents hash."""
CALL_MESSAGE_CONTENTS_HASH_TEMPLATE = "%s-%s"
"""`%s-%s`: Call ID - Iteration. The key template for call message contents hash keys."""
CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE = "mfp:%s"
"""`%s`: Call ID. A map from propagation IDs to message IDs."""
CALL_MESSAGE_CONTENT_ITERATIONS_TEMPLATE = "mci:%s"
"""`%s`: Call ID. A map of messages to amount of iterations of their message contents."""
CALL_IDS_TO_FLUSH = "call-ids-to-flush"
"""A set of IDs of calls to flush to the DB."""
BAN_TEMPLATE = "b:%s"
PREFIX_TEMPLATE = "p:%s"
payphone/core/utils.py
+28 -1
View File
@@ -1,4 +1,4 @@
from typing import Callable, Iterable, TypeVar
from typing import Callable, Iterable, TypeVar, Generator
T = TypeVar("T")
@@ -28,3 +28,30 @@ def deduplicate(i: Iterable[T], key: Callable[[T], U] = lambda i: i) -> list[T]:
found_keys.add(k)
return result
def chunks(i: Iterable[T], size: int) -> Generator[list[T], None, None]:
"""Chunk an iterable.
Args:
i (Iterable[T]): The iterable.
size (int): The size of each chunk.
Returns:
Generator[list[T], None, None]: The iterable's chunks.
"""
if size <= 0:
raise ValueError("Chunks size can't be lower than 1.")
chunk = []
for item in i:
chunk.append(item)
if len(chunk) == size:
yield chunk
chunk = []
if len(chunk) > 0:
yield chunk