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

fix(db): ⚡ Add a messages flushing queue.

Browse Source
This commit is contained in:
Nyeki.py
2024-10-18 19:44:12 -03:00
parent db1dbb4349
commit e7b553662d
8 changed files with 158 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
payphone/bot.py
+11 -14
View File
@@ -9,16 +9,13 @@ from discord.ext.commands import errors
from tortoise import expressions, transactions, exceptions as tortoise_exc
import redis.asyncio as aredis
from payphone.db import setup, models
from payphone.core import config, constants, exceptions, formatting, help_command
from payphone.core import constants, exceptions, formatting, help_command
from payphone.core.redis import redis
# TODO: Make it an AutoShardedBot on production.
class Payphone(commands.AutoShardedBot):
redis = aredis.Redis(host="localhost" if config.IS_DEBUG else "redis")
def __init__(self) -> None:
intents = discord.Intents.default()
intents.message_content = True
@@ -43,9 +40,9 @@ class Payphone(commands.AutoShardedBot):
await setup.init()
await self.redis.ping()
await redis.ping()
await self.redis.delete(
await redis.delete(
constants.CHANNELS_ONCALL_SET, constants.CHANNELS_SEARCH_QUEUE
)
@@ -125,7 +122,7 @@ class Payphone(commands.AutoShardedBot):
user_id = user if isinstance(user, int) else user.id
key = constants.BAN_TEMPLATE % user_id
is_banned = await self.redis.get(key)
is_banned = await redis.get(key)
if is_banned is not None:
return bool(int(is_banned))
@@ -139,7 +136,7 @@ class Payphone(commands.AutoShardedBot):
),
).exists()
await self.redis.set(key, int(is_banned), ex=300)
await redis.set(key, int(is_banned), ex=300)
return is_banned
@@ -204,7 +201,7 @@ class Payphone(commands.AutoShardedBot):
banned_by_id=mod_id,
)
await self.redis.set(key, "1", ex=300)
await redis.set(key, "1", ex=300)
return row
@@ -231,7 +228,7 @@ class Payphone(commands.AutoShardedBot):
user_id=user_id,
).update(expires_at=datetime.now(timezone.utc))
await self.redis.set(key, "0", ex=300)
await redis.set(key, "0", ex=300)
async def get_prefix(
self, message: discord.Message | None = None, server_id: int | None = None
@@ -274,7 +271,7 @@ class Payphone(commands.AutoShardedBot):
key = constants.PREFIX_TEMPLATE % server_id
prefix: bytes | None = await self.redis.get(key)
prefix: bytes | None = await redis.get(key)
if prefix is not None:
prefix = prefix.decode()
@@ -282,7 +279,7 @@ class Payphone(commands.AutoShardedBot):
else:
server, _ = await models.Server.get_or_create(id=server_id)
await self.redis.set(key, server.prefix, ex=60 * 60)
await redis.set(key, server.prefix, ex=60 * 60)
prefix = server.prefix
@@ -314,7 +311,7 @@ class Payphone(commands.AutoShardedBot):
await models.Server.filter(id=server_id).update(prefix=prefix)
await self.redis.set(constants.PREFIX_TEMPLATE % server_id, prefix, ex=60 * 60)
await redis.set(constants.PREFIX_TEMPLATE % server_id, prefix, ex=60 * 60)
async def on_command_error(
self, context: commands.Context, exception: commands.CommandError
payphone/cogs/calls/cog.py
+102 -69
View File
@@ -9,15 +9,16 @@ from discord.utils import find
from tortoise import transactions, exceptions
import asyncpg.exceptions
import aiohttp
import orjson
from payphone import bot
from payphone.db import models, utils
from payphone.cogs import base
from payphone.cogs.calls import embeds
from payphone.core import constants, checks, censorship, formatting
from payphone.core.redis import redis
class Calls(base.Cog):
@@ -27,11 +28,44 @@ class Calls(base.Cog):
self.hangup_calls.add_exception_type(exceptions.ConfigurationError)
self.hangup_calls.start()
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.smembers(constants.MESSAGES_QUEUE)
pipe.smembers(constants.MESSAGE_CONTENTS_QUEUE)
pipe.delete(
constants.MESSAGES_QUEUE,
constants.MESSAGE_CONTENTS_QUEUE,
constants.MESSAGES_TO_FLUSH_IDS,
)
messages, message_contents, _ = await pipe.execute()
messages = [models.Message(**orjson.loads(msg)) for msg in messages]
message_contents = [
models.MessageContent(**orjson.loads(msg)) for msg in message_contents
]
await models.Message.bulk_create(messages, batch_size=500)
await models.MessageContent.bulk_create(message_contents, batch_size=500)
print(
f"DONE! Flushed {len(messages):,} messages and {len(message_contents):,} message contents."
)
@tasks.loop(seconds=15)
async def hangup_calls(self) -> None:
"""Hangs up inactive calls.
@@ -74,9 +108,7 @@ class Calls(base.Cog):
if len(hung_up_channel_ids) > 0:
await models.Channel.filter(id__in=hung_up_channel_ids).update(call=None)
await self.bot.redis.srem(
constants.CHANNELS_ONCALL_SET, *hung_up_channel_ids
)
await redis.srem(constants.CHANNELS_ONCALL_SET, *hung_up_channel_ids)
@commands.command(aliases=["c"])
@commands.guild_only()
@@ -94,7 +126,7 @@ class Calls(base.Cog):
user_id = ctx.author.id
channel_id = ctx.channel.id
is_on_call = await self.bot.redis.sismember(
is_on_call = await redis.sismember(
constants.CHANNELS_ONCALL_SET, str(channel_id)
)
@@ -104,7 +136,7 @@ class Calls(base.Cog):
)
return
is_searching = await self.bot.redis.sismember(
is_searching = await redis.sismember(
constants.CHANNELS_SEARCH_QUEUE, str(channel_id)
)
@@ -153,9 +185,7 @@ class Calls(base.Cog):
if last_calls_rss_channels_map[call_rs.call_id] != call_rs.channel_id:
last_calls_other_channels_ids.add(call_rs.channel_id)
other_channel_id: int = await self.bot.redis.spop(
constants.CHANNELS_SEARCH_QUEUE
)
other_channel_id: int = await redis.spop(constants.CHANNELS_SEARCH_QUEUE)
i = 0
while True:
@@ -173,13 +203,9 @@ class Calls(base.Cog):
i += 1
if other_channel_id in last_calls_other_channels_ids:
await self.bot.redis.sadd(
constants.CHANNELS_SEARCH_QUEUE, str(other_channel_id)
)
await redis.sadd(constants.CHANNELS_SEARCH_QUEUE, str(other_channel_id))
other_channel_id = await self.bot.redis.spop(
constants.CHANNELS_SEARCH_QUEUE
)
other_channel_id = await redis.spop(constants.CHANNELS_SEARCH_QUEUE)
if other_channel_id is None:
break
@@ -191,7 +217,7 @@ class Calls(base.Cog):
channel = await utils.get_channel(ctx.channel)
if other_channel_id is not None:
await self.bot.redis.sadd(
await redis.sadd(
constants.CHANNELS_ONCALL_SET, str(channel_id), str(other_channel_id)
)
@@ -220,14 +246,12 @@ class Calls(base.Cog):
).update(call=call)
except errors.DiscordException:
await self.bot.redis.srem(
await redis.srem(
constants.CHANNELS_ONCALL_SET,
str(channel_id),
str(other_channel_id),
)
await self.bot.redis.sadd(
constants.CHANNELS_SEARCH_QUEUE, str(channel_id)
)
await redis.sadd(constants.CHANNELS_SEARCH_QUEUE, str(channel_id))
msg_phone_answered = ""
await ctx.send(
@@ -239,7 +263,7 @@ class Calls(base.Cog):
)
else:
await self.bot.redis.sadd(constants.CHANNELS_SEARCH_QUEUE, str(channel_id))
await redis.sadd(constants.CHANNELS_SEARCH_QUEUE, str(channel_id))
await ctx.send(
(
@@ -264,7 +288,7 @@ class Calls(base.Cog):
user_id = ctx.author.id
channel_id = ctx.channel.id
is_on_call = await self.bot.redis.sismember(
is_on_call = await redis.sismember(
constants.CHANNELS_ONCALL_SET, str(channel_id)
)
@@ -292,7 +316,7 @@ class Calls(base.Cog):
except errors.DiscordException:
pass
await self.bot.redis.srem(
await redis.srem(
constants.CHANNELS_ONCALL_SET, str(channel_id), str(other_channel.id)
)
@@ -335,9 +359,7 @@ class Calls(base.Cog):
if last_calls_rss_channels_map[call_rs.call_id] != call_rs.channel_id:
last_calls_other_channels_ids.add(call_rs.channel_id)
other_channel_id: int = await self.bot.redis.spop(
constants.CHANNELS_SEARCH_QUEUE
)
other_channel_id: int = await redis.spop(constants.CHANNELS_SEARCH_QUEUE)
i = 0
while True:
@@ -355,13 +377,9 @@ class Calls(base.Cog):
i += 1
if other_channel_id in last_calls_other_channels_ids:
await self.bot.redis.sadd(
constants.CHANNELS_SEARCH_QUEUE, str(other_channel_id)
)
await redis.sadd(constants.CHANNELS_SEARCH_QUEUE, str(other_channel_id))
other_channel_id = await self.bot.redis.spop(
constants.CHANNELS_SEARCH_QUEUE
)
other_channel_id = await redis.spop(constants.CHANNELS_SEARCH_QUEUE)
if other_channel_id is None:
break
@@ -371,7 +389,7 @@ class Calls(base.Cog):
break
if other_channel_id is not None:
await self.bot.redis.sadd(
await redis.sadd(
constants.CHANNELS_ONCALL_SET, str(channel_id), str(other_channel_id)
)
@@ -400,14 +418,12 @@ class Calls(base.Cog):
).update(call=call)
except errors.DiscordException:
await self.bot.redis.srem(
await redis.srem(
constants.CHANNELS_ONCALL_SET,
str(channel_id),
str(other_channel_id),
)
await self.bot.redis.sadd(
constants.CHANNELS_SEARCH_QUEUE, str(channel_id)
)
await redis.sadd(constants.CHANNELS_SEARCH_QUEUE, str(channel_id))
msg_phone_answered = ""
await ctx.send(
@@ -423,7 +439,7 @@ class Calls(base.Cog):
)
else:
await self.bot.redis.sadd(constants.CHANNELS_SEARCH_QUEUE, str(channel_id))
await redis.sadd(constants.CHANNELS_SEARCH_QUEUE, str(channel_id))
channel.call = None
await channel.save()
@@ -447,7 +463,7 @@ class Calls(base.Cog):
channel_id = ctx.channel.id
is_on_call = await self.bot.redis.sismember(
is_on_call = await redis.sismember(
constants.CHANNELS_ONCALL_SET, str(channel_id)
)
@@ -462,7 +478,7 @@ class Calls(base.Cog):
lambda c: c.id != channel_id, call.channels
)
await self.bot.redis.srem(
await redis.srem(
constants.CHANNELS_ONCALL_SET,
*(
[str(channel_id)]
@@ -493,12 +509,12 @@ class Calls(base.Cog):
call.ended_at = datetime.now(timezone.utc)
await call.save()
is_searching = await self.bot.redis.sismember(
is_searching = await redis.sismember(
constants.CHANNELS_SEARCH_QUEUE, str(channel_id)
)
if is_searching:
await self.bot.redis.srem(constants.CHANNELS_SEARCH_QUEUE, str(channel_id))
await redis.srem(constants.CHANNELS_SEARCH_QUEUE, str(channel_id))
await ctx.send(f"☎️ {secrets.choice(constants.MSGS_PHONE_NO_ANSWER)}")
@@ -571,7 +587,7 @@ class Calls(base.Cog):
await models.Channel.bulk_update(updated_channels, ["call_id"])
await self.bot.redis.srem(
await redis.srem(
constants.CHANNELS_ONCALL_SET,
str(updated_channels[0]),
str(updated_channels[1]),
@@ -581,7 +597,7 @@ class Calls(base.Cog):
other_channel = find(lambda c: c.id != channel_id, updated_channels)
await self.bot.redis.srem(
await redis.srem(
constants.CHANNELS_ONCALL_SET, str(other_channel.id), str(channel_id)
)
@@ -651,7 +667,7 @@ class Calls(base.Cog):
user_id = message.author.id
channel_id = message.channel.id
is_on_call = await self.bot.redis.sismember(
is_on_call = await redis.sismember(
constants.CHANNELS_ONCALL_SET, str(channel_id)
)
@@ -671,7 +687,7 @@ class Calls(base.Cog):
)
except exceptions.DoesNotExist:
await self.bot.redis.srem(constants.CHANNELS_ONCALL_SET, str(channel_id))
await redis.srem(constants.CHANNELS_ONCALL_SET, str(channel_id))
return
other_channel: models.Channel = find(
@@ -747,28 +763,45 @@ class Calls(base.Cog):
wait=True,
)
while True:
try:
message_row = await models.Message.create(
id=message.id,
call=call,
channel=channel,
user=user,
profile=profile,
mask=call_rs.mask if call_rs.is_masked else None,
propagation_id=webhook_message.id if webhook_message else None,
propagation_channel=other_channel if webhook_message else None,
)
break
pipe = redis.pipeline()
except asyncpg.exceptions.DeadlockDetectedError:
continue
await models.MessageContent.create(
message=message_row,
content=cleaned_content,
was_propagated=not bool(blocked_word),
pipe.sadd(
constants.MESSAGES_QUEUE,
orjson.dumps(
{
"id": message.id,
"call_id": call.id,
"channel_id": channel.id,
"user_id": user.id,
"profile_id": profile.id,
"mask": call_rs.mask if call_rs.is_masked else None,
"propagation_id": webhook_message.id if webhook_message else None,
"propagation_channel_id": (
other_channel.id if webhook_message else None
),
"replying_to_id": (
message.reference.message_id
if message.reference is not None
else None
),
"created_at": datetime.now(timezone.utc),
}
),
)
pipe.sadd(
constants.MESSAGE_CONTENTS_QUEUE,
orjson.dumps(
{
"message_id": message.id,
"content": cleaned_content,
"was_propagated": not bool(blocked_word),
"created_at": datetime.now(timezone.utc),
}
),
)
pipe.sadd(constants.MESSAGES_TO_FLUSH_IDS, message.id)
await pipe.execute()
@commands.Cog.listener()
@commands.guild_only()
@@ -785,7 +818,7 @@ class Calls(base.Cog):
message_id = payload.message_id
channel_id = payload.channel_id
is_on_call = await self.bot.redis.sismember(
is_on_call = await redis.sismember(
constants.CHANNELS_ONCALL_SET, str(channel_id)
)
@@ -852,7 +885,7 @@ class Calls(base.Cog):
if censorship.moderate(content) is not None:
return
is_on_call = await self.bot.redis.sismember(
is_on_call = await redis.sismember(
constants.CHANNELS_ONCALL_SET, str(channel_id)
)
payphone/cogs/mods/cog.py
+7 -6
View File
@@ -16,6 +16,7 @@ from payphone.db import models, utils
from payphone.cogs import base
from payphone.cogs.mods import embeds, views
from payphone.core import checks, converters, constants, formatting as fmt
from payphone.core.redis import redis
class Mods(base.Cog):
@@ -243,10 +244,10 @@ class Mods(base.Cog):
"""
oncall_channel_ids = list(
await self.bot.redis.smembers(constants.CHANNELS_ONCALL_SET)
await redis.smembers(constants.CHANNELS_ONCALL_SET)
)
searching_channel_ids = list(
await self.bot.redis.smembers(constants.CHANNELS_SEARCH_QUEUE)
await redis.smembers(constants.CHANNELS_SEARCH_QUEUE)
)
channel_ids = [int(c) for c in oncall_channel_ids] + [
@@ -272,10 +273,10 @@ class Mods(base.Cog):
await asyncio.sleep(seconds)
oncall_channel_ids = list(
await self.bot.redis.smembers(constants.CHANNELS_ONCALL_SET)
await redis.smembers(constants.CHANNELS_ONCALL_SET)
)
searching_channel_ids = list(
await self.bot.redis.smembers(constants.CHANNELS_SEARCH_QUEUE)
await redis.smembers(constants.CHANNELS_SEARCH_QUEUE)
)
channel_ids = [int(c) for c in oncall_channel_ids] + [
@@ -311,11 +312,11 @@ class Mods(base.Cog):
stats = [
(
"On-call Channels",
lambda: self.bot.redis.scard(constants.CHANNELS_ONCALL_SET),
lambda: redis.scard(constants.CHANNELS_ONCALL_SET),
),
(
"Searching Channels",
lambda: self.bot.redis.scard(constants.CHANNELS_SEARCH_QUEUE),
lambda: redis.scard(constants.CHANNELS_SEARCH_QUEUE),
),
(
"Messages",
payphone/core/checks.py
+2 -1
View File
@@ -2,6 +2,7 @@ from discord.ext import commands
from payphone.db import utils
from payphone.core import exceptions, constants
from payphone.core.redis import redis
def is_moderator():
@@ -39,7 +40,7 @@ def is_on_call():
async def predicate(ctx: commands.Context) -> bool:
channel_id = ctx.channel.id
is_on_call = await ctx.bot.redis.sismember(
is_on_call = await redis.sismember(
constants.CHANNELS_ONCALL_SET, str(channel_id)
)
payphone/core/constants.py
+3
View File
@@ -170,6 +170,9 @@ MSGS_PHONE_NOT_IN_CALL = [
CHANNELS_SEARCH_QUEUE = "channels:searching"
CHANNELS_ONCALL_SET = "channels:on-call"
MESSAGES_QUEUE = "messages"
MESSAGE_CONTENTS_QUEUE = "message-contents"
MESSAGES_TO_FLUSH_IDS = "message-ids"
BAN_TEMPLATE = "b:%s"
PREFIX_TEMPLATE = "p:%s"
payphone/core/redis.py
+6
View File
@@ -0,0 +1,6 @@
from redis import asyncio as aredis
from payphone.core import config
redis = aredis.Redis(host="localhost" if config.IS_DEBUG else "redis")