UiO-Gaming/mustafa
1
0 Fork 0
mirror of https://github.com/UiO-Gaming/mustafa.git synced 2026-11-05 13:56:48 +01:00
Code Issues Packages Projects Releases Wiki Activity

fix cache race condition

Browse Source
This commit is contained in:
leander committed 2026-10-29 22:06:45 +01:00
1 parent f445a62681
commit cecc2f7666
1 file changed
+48 -58
Download Patch File Download Diff File Expand all files Collapse all files
src/cogs/word_cloud.py
+48 -58
View File
@@ -1,7 +1,8 @@
import asyncio
import functools
import itertools
import json
import re
from collections import Counter
from collections import defaultdict
from io import BytesIO
from io import StringIO
@@ -33,13 +34,15 @@ class WordCloud(commands.Cog):
self.bot = bot
# Default dict where all users are empty defaultdicts
# The user's defaultdict has a default value of 0
self.word_freq_cache = defaultdict(lambda: defaultdict(int))
# Word counts per user that haven't been written to the database yet
self.word_freq_cache: defaultdict[int, Counter] = defaultdict(Counter)
# We cache the consenting users to avoid querying the database every 16:36
# It's frail but it works unless you have a skill issue
self.consenting_users = []
# Cached to avoid querying the database on every message
self.consenting_users: set[int] = set()
# Held while writing the cache and while deleting a user's data,
# so a user's words can't be written after their data has been deleted
self.cache_lock = asyncio.Lock()
try:
nltk.data.find("corpora/stopwords")
@@ -96,7 +99,7 @@ class WordCloud(commands.Cog):
SELECT discord_user_id
FROM wordcloud_metadata
""")
self.consenting_users = [row[0] for row in rows]
self.consenting_users = {row[0] for row in rows}
async def insert_cache(self):
"""
@@ -104,34 +107,26 @@ class WordCloud(commands.Cog):
We do this in order to prevent excess database writes
"""
if not self.word_freq_cache:
return
async with self.cache_lock:
if not self.word_freq_cache:
return
# Swap in an empty cache before awaiting, so words from messages sent in the meantime aren't lost
cache, self.word_freq_cache = self.word_freq_cache, defaultdict(lambda: defaultdict(int))
# Swap in an empty cache before awaiting, so words from messages sent in the meantime aren't lost
cache, self.word_freq_cache = self.word_freq_cache, defaultdict(Counter)
word_freqs = [(user_id, word, freq) for user_id, words in cache.items() for word, freq in words.items()]
# Flatten the defaultdict into a list of tuples
# This is a bit of a hack but it works
word_freqs = list(
itertools.chain.from_iterable(
[(user_id, word, freq) for word, freq in user_word_freqs.items()]
for user_id, user_word_freqs in cache.items()
)
)
# Insert cache into database
try:
await self.bot.db.executemany(
"""
INSERT INTO wordcloud_words (discord_user_id, word, frequency)
VALUES (%s, %s, %s)
ON CONFLICT (discord_user_id, word)
DO UPDATE SET frequency = wordcloud_words.frequency + EXCLUDED.frequency
""",
word_freqs,
)
except psycopg.Error as err:
self.bot.logger.error(f"Failed to insert wordcloud cache into database - {err}")
try:
await self.bot.db.executemany(
"""
INSERT INTO wordcloud_words (discord_user_id, word, frequency)
VALUES (%s, %s, %s)
ON CONFLICT (discord_user_id, word)
DO UPDATE SET frequency = wordcloud_words.frequency + EXCLUDED.frequency
""",
word_freqs,
)
except psycopg.Error as err:
self.bot.logger.error(f"Failed to insert wordcloud cache into database - {err}")
@tasks.loop(minutes=20)
@misc_utils.log_errors
@@ -156,12 +151,7 @@ class WordCloud(commands.Cog):
if not self.can_count_message(message):
return
tokens = self.tokenize(message.clean_content)
for token in tokens:
# Enter into cache
# This may take a performance hit but who tf cares. We're using python anyway
self.word_freq_cache[message.author.id][token] += 1
self.word_freq_cache[message.author.id].update(self.tokenize(message.clean_content))
def can_count_message(self, message: discord.Message) -> bool:
"""
@@ -339,7 +329,7 @@ class WordCloud(commands.Cog):
ephemeral=False,
)
self.consenting_users.append(interaction.user.id)
self.consenting_users.add(interaction.user.id)
embed = embed_templates.success("Samtykke registrert!")
await interaction.response.send_message(embed=embed, ephemeral=False)
@@ -356,28 +346,28 @@ class WordCloud(commands.Cog):
interaction (discord.Interaction): Slash command context object
"""
try:
self.consenting_users.remove(interaction.user.id)
except ValueError:
if interaction.user.id not in self.consenting_users:
return await interaction.response.send_message(
embed=embed_templates.error_warning(self.MSG_NO_DATA), ephemeral=False
)
self.word_freq_cache.pop(interaction.user.id, None)
async with self.cache_lock:
try:
await self.bot.db.execute_in_transaction(
[
("DELETE FROM wordcloud_metadata WHERE discord_user_id = %s", (interaction.user.id,)),
("DELETE FROM wordcloud_words WHERE discord_user_id = %s", (interaction.user.id,)),
]
)
except psycopg.Error as err:
self.bot.logger.error(f"Failed to delete wordcloud metadata from database - {err}")
return await interaction.response.send_message(
embed=embed_templates.error_fatal("Klarte ikke å slette fra database"),
ephemeral=False,
)
try:
await self.bot.db.execute_in_transaction(
[
("DELETE FROM wordcloud_metadata WHERE discord_user_id = %s", (interaction.user.id,)),
("DELETE FROM wordcloud_words WHERE discord_user_id = %s", (interaction.user.id,)),
]
)
except psycopg.Error as err:
self.bot.logger.error(f"Failed to delete wordcloud metadata from database - {err}")
return await interaction.response.send_message(
embed=embed_templates.error_fatal("Klarte ikke å slette fra database"),
ephemeral=False,
)
self.consenting_users.discard(interaction.user.id)
self.word_freq_cache.pop(interaction.user.id, None)
embed = embed_templates.success("Meldingsdata er slettet!")
await interaction.response.send_message(embed=embed)