fix cache race condition

src/cogs/word_cloud.py

@@ -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
+        # Word counts per user that haven't been written to the database yet
-        # The user's defaultdict has a default value of 0
+        self.word_freq_cache: defaultdict[int, Counter] = defaultdict(Counter)
-        self.word_freq_cache = defaultdict(lambda: defaultdict(int))
 
-        # We cache the consenting users to avoid querying the database every 16:36
+        # Cached to avoid querying the database on every message
-        # It's frail but it works unless you have a skill issue
+        self.consenting_users: set[int] = set()
-        self.consenting_users = []
+
+        # 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:
+        async with self.cache_lock:
-            return
+            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
+            # 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))
+            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
+            try:
-        # This is a bit of a hack but it works
+                await self.bot.db.executemany(
-        word_freqs = list(
+                    """
-            itertools.chain.from_iterable(
+                    INSERT INTO wordcloud_words (discord_user_id, word, frequency)
-                [(user_id, word, freq) for word, freq in user_word_freqs.items()]
+                    VALUES (%s, %s, %s)
-                for user_id, user_word_freqs in cache.items()
+                    ON CONFLICT (discord_user_id, word)
-            )
+                    DO UPDATE SET frequency = wordcloud_words.frequency + EXCLUDED.frequency
-        )
+                    """,
-
+                    word_freqs,
-        # Insert cache into database
+                )
-        try:
+            except psycopg.Error as err:
-            await self.bot.db.executemany(
+                self.bot.logger.error(f"Failed to insert wordcloud cache into database - {err}")
-                """
-                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)
+        self.word_freq_cache[message.author.id].update(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
 
     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:
+        if interaction.user.id not in self.consenting_users:
-            self.consenting_users.remove(interaction.user.id)
-        except ValueError:
             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:
+            self.consenting_users.discard(interaction.user.id)
-            await self.bot.db.execute_in_transaction(
+            self.word_freq_cache.pop(interaction.user.id, None)
-                [
-                    ("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,
-            )
 
         embed = embed_templates.success("Meldingsdata er slettet!")
         await interaction.response.send_message(embed=embed)