Merge pull request #987 from python-discord/help_channel_rediscache

Add persistence to the help channel system

3 files changed, 98 insertions, 53 deletions

diff --git a/bot/cogs/help_channels.py b/bot/cogs/help_channels.py
index 6ff285c37..187adfe51 100644
--- a/bot/cogs/help_channels.py
+++ b/bot/cogs/help_channels.py
@@ -5,8 +5,7 @@ import logging
 import random
 import typing as t
 from collections import deque
-from contextlib import suppress
-from datetime import datetime
+from datetime import datetime, timedelta, timezone
 from pathlib import Path
 
 import discord
@@ -15,6 +14,7 @@ from discord.ext import commands
 
 from bot import constants
 from bot.bot import Bot
+from bot.utils import RedisCache
 from bot.utils.checks import with_role_check
 from bot.utils.scheduling import Scheduler
 
@@ -99,13 +99,24 @@ class HelpChannels(Scheduler, commands.Cog):
     Help channels are named after the chemical elements in `bot/resources/elements.json`.
     """
 
+    # This cache tracks which channels are claimed by which members.
+    # RedisCache[discord.TextChannel.id, t.Union[discord.User.id, discord.Member.id]]
+    help_channel_claimants = RedisCache()
+
+    # This cache maps a help channel to whether it has had any
+    # activity other than the original claimant. True being no other
+    # activity and False being other activity.
+    # RedisCache[discord.TextChannel.id, bool]
+    unanswered = RedisCache()
+
+    # This dictionary maps a help channel to the time it was claimed
+    # RedisCache[discord.TextChannel.id, UtcPosixTimestamp]
+    claim_times = RedisCache()
+
     def __init__(self, bot: Bot):
         super().__init__()
 
         self.bot = bot
-        self.help_channel_claimants: (
-            t.Dict[discord.TextChannel, t.Union[discord.Member, discord.User]]
-        ) = {}
 
         # Categories
         self.available_category: discord.CategoryChannel = None
@@ -125,16 +136,6 @@ class HelpChannels(Scheduler, commands.Cog):
         self.on_message_lock = asyncio.Lock()
         self.init_task = self.bot.loop.create_task(self.init_cog())
 
-        # Stats
-
-        # This dictionary maps a help channel to the time it was claimed
-        self.claim_times: t.Dict[int, datetime] = {}
-
-        # This dictionary maps a help channel to whether it has had any
-        # activity other than the original claimant. True being no other
-        # activity and False being other activity.
-        self.unanswered: t.Dict[int, bool] = {}
-
     def cog_unload(self) -> None:
         """Cancel the init task and scheduled tasks when the cog unloads."""
         log.trace("Cog unload: cancelling the init_cog task")
@@ -197,7 +198,7 @@ class HelpChannels(Scheduler, commands.Cog):
 
     async def dormant_check(self, ctx: commands.Context) -> bool:
         """Return True if the user is the help channel claimant or passes the role check."""
-        if self.help_channel_claimants.get(ctx.channel) == ctx.author:
+        if await self.help_channel_claimants.get(ctx.channel.id) == ctx.author.id:
             log.trace(f"{ctx.author} is the help channel claimant, passing the check for dormant.")
             self.bot.stats.incr("help.dormant_invoke.claimant")
             return True
@@ -222,10 +223,11 @@ class HelpChannels(Scheduler, commands.Cog):
         log.trace("close command invoked; checking if the channel is in-use.")
         if ctx.channel.category == self.in_use_category:
             if await self.dormant_check(ctx):
-                with suppress(KeyError):
-                    del self.help_channel_claimants[ctx.channel]
 
+                # Remove the claimant and the cooldown role
+                await self.help_channel_claimants.delete(ctx.channel.id)
                 await self.remove_cooldown_role(ctx.author)
+
                 # Ignore missing task when cooldown has passed but the channel still isn't dormant.
                 self.cancel_task(ctx.author.id, ignore_missing=True)
 
@@ -284,6 +286,15 @@ class HelpChannels(Scheduler, commands.Cog):
             if channel.category_id == category.id and not self.is_excluded_channel(channel):
                 yield channel
 
+    async def get_in_use_time(self, channel_id: int) -> t.Optional[timedelta]:
+        """Return the duration `channel_id` has been in use. Return None if it's not in use."""
+        log.trace(f"Calculating in use time for channel {channel_id}.")
+
+        claimed_timestamp = await self.claim_times.get(channel_id)
+        if claimed_timestamp:
+            claimed = datetime.utcfromtimestamp(claimed_timestamp)
+            return datetime.utcnow() - claimed
+
     @staticmethod
     def get_names() -> t.List[str]:
         """
@@ -386,7 +397,7 @@ class HelpChannels(Scheduler, commands.Cog):
 
         log.trace("Initialising the cog.")
         await self.init_categories()
-        await self.reset_send_permissions()
+        await self.check_cooldowns()
 
         self.channel_queue = self.create_channel_queue()
         self.name_queue = self.create_name_queue()
@@ -546,19 +557,17 @@ class HelpChannels(Scheduler, commands.Cog):
 
         self.bot.stats.incr(f"help.dormant_calls.{caller}")
 
-        if channel.id in self.claim_times:
-            claimed = self.claim_times[channel.id]
-            in_use_time = datetime.now() - claimed
+        in_use_time = await self.get_in_use_time(channel.id)
+        if in_use_time:
             self.bot.stats.timing("help.in_use_time", in_use_time)
 
-        if channel.id in self.unanswered:
-            if self.unanswered[channel.id]:
-                self.bot.stats.incr("help.sessions.unanswered")
-            else:
-                self.bot.stats.incr("help.sessions.answered")
+        unanswered = await self.unanswered.get(channel.id)
+        if unanswered:
+            self.bot.stats.incr("help.sessions.unanswered")
+        elif unanswered is not None:
+            self.bot.stats.incr("help.sessions.answered")
 
         log.trace(f"Position of #{channel} ({channel.id}) is actually {channel.position}.")
-
         log.trace(f"Sending dormant message for #{channel} ({channel.id}).")
         embed = discord.Embed(description=DORMANT_MSG)
         await channel.send(embed=embed)
@@ -637,17 +646,17 @@ class HelpChannels(Scheduler, commands.Cog):
         if self.is_in_category(channel, constants.Categories.help_in_use):
             log.trace(f"Checking if #{channel} ({channel.id}) has been answered.")
 
-            # Check if there is an entry in unanswered (does not persist across restarts)
-            if channel.id in self.unanswered:
-                claimant = self.help_channel_claimants.get(channel)
-                if not claimant:
-                    # The mapping for this channel was lost, we can't do anything.
+            # Check if there is an entry in unanswered
+            if await self.unanswered.contains(channel.id):
+                claimant_id = await self.help_channel_claimants.get(channel.id)
+                if not claimant_id:
+                    # The mapping for this channel doesn't exist, we can't do anything.
                     return
 
                 # Check the message did not come from the claimant
-                if claimant.id != message.author.id:
+                if claimant_id != message.author.id:
                     # Mark the channel as answered
-                    self.unanswered[channel.id] = False
+                    await self.unanswered.set(channel.id, False)
 
     @commands.Cog.listener()
     async def on_message(self, message: discord.Message) -> None:
@@ -680,12 +689,15 @@ class HelpChannels(Scheduler, commands.Cog):
             await self.move_to_in_use(channel)
             await self.revoke_send_permissions(message.author)
             # Add user with channel for dormant check.
-            self.help_channel_claimants[channel] = message.author
+            await self.help_channel_claimants.set(channel.id, message.author.id)
 
             self.bot.stats.incr("help.claimed")
 
-            self.claim_times[channel.id] = datetime.now()
-            self.unanswered[channel.id] = True
+            # Must use a timezone-aware datetime to ensure a correct POSIX timestamp.
+            timestamp = datetime.now(timezone.utc).timestamp()
+            await self.claim_times.set(channel.id, timestamp)
+
+            await self.unanswered.set(channel.id, True)
 
             log.trace(f"Releasing on_message lock for {message.id}.")
 
@@ -720,15 +732,28 @@ class HelpChannels(Scheduler, commands.Cog):
         msg = await self.get_last_message(channel)
         return self.match_bot_embed(msg, AVAILABLE_MSG)
 
-    async def reset_send_permissions(self) -> None:
-        """Reset send permissions in the Available category for claimants."""
-        log.trace("Resetting send permissions in the Available category.")
+    async def check_cooldowns(self) -> None:
+        """Remove expired cooldowns and re-schedule active ones."""
+        log.trace("Checking all cooldowns to remove or re-schedule them.")
         guild = self.bot.get_guild(constants.Guild.id)
+        cooldown = constants.HelpChannels.claim_minutes * 60
+
+        for channel_id, member_id in await self.help_channel_claimants.items():
+            member = guild.get_member(member_id)
+            if not member:
+                continue  # Member probably left the guild.
+
+            in_use_time = await self.get_in_use_time(channel_id)
 
-        # TODO: replace with a persistent cache cause checking every member is quite slow
-        for member in guild.members:
-            if self.is_claimant(member):
+            if not in_use_time or in_use_time.seconds > cooldown:
+                # Remove the role if no claim time could be retrieved or if the cooldown expired.
+                # Since the channel is in the claimants cache, it is definitely strange for a time
+                # to not exist. However, it isn't a reason to keep the user stuck with a cooldown.
                 await self.remove_cooldown_role(member)
+            else:
+                # The member is still on a cooldown; re-schedule it for the remaining time.
+                remaining = cooldown - in_use_time.seconds
+                await self.schedule_cooldown_expiration(member, remaining)
 
     async def add_cooldown_role(self, member: discord.Member) -> None:
         """Add the help cooldown role to `member`."""
@@ -781,11 +806,14 @@ class HelpChannels(Scheduler, commands.Cog):
         # Would mean the user somehow bypassed the lack of permissions (e.g. user is guild owner).
         self.cancel_task(member.id, ignore_missing=True)
 
-        timeout = constants.HelpChannels.claim_minutes * 60
-        callback = self.remove_cooldown_role(member)
+        await self.schedule_cooldown_expiration(member, constants.HelpChannels.claim_minutes * 60)
 
-        log.trace(f"Scheduling {member}'s ({member.id}) send message permissions to be reinstated.")
-        self.schedule_task(member.id, TaskData(timeout, callback))
+    async def schedule_cooldown_expiration(self, member: discord.Member, seconds: int) -> None:
+        """Schedule the cooldown role for `member` to be removed after a duration of `seconds`."""
+        log.trace(f"Scheduling removal of {member}'s ({member.id}) cooldown.")
+
+        callback = self.remove_cooldown_role(member)
+        self.schedule_task(member.id, TaskData(seconds, callback))
 
     async def send_available_message(self, channel: discord.TextChannel) -> None:
         """Send the available message by editing a dormant message or sending a new message."""
diff --git a/bot/utils/redis_cache.py b/bot/utils/redis_cache.py
index 354e987b9..58cfe1df5 100644
--- a/bot/utils/redis_cache.py
+++ b/bot/utils/redis_cache.py
@@ -11,7 +11,7 @@ log = logging.getLogger(__name__)
 
 # Type aliases
 RedisKeyType = Union[str, int]
-RedisValueType = Union[str, int, float]
+RedisValueType = Union[str, int, float, bool]
 RedisKeyOrValue = Union[RedisKeyType, RedisValueType]
 
 # Prefix tuples
@@ -20,6 +20,7 @@ _VALUE_PREFIXES = (
     ("f|", float),
     ("i|", int),
     ("s|", str),
+    ("b|", bool),
 )
 _KEY_PREFIXES = (
     ("i|", int),
@@ -47,8 +48,8 @@ class RedisCache:
     behaves, and should be familiar to Python users. The biggest difference is that
     all the public methods in this class are coroutines, and must be awaited.
 
-    Because of limitations in Redis, this cache will only accept strings, integers and
-    floats both for keys and values.
+    Because of limitations in Redis, this cache will only accept strings and integers for keys,
+    and strings, integers, floats and booleans for values.
 
     Please note that this class MUST be created as a class attribute, and that that class
     must also contain an attribute with an instance of our Bot. See `__get__` and `__set_name__`
@@ -108,8 +109,15 @@ class RedisCache:
     def _to_typestring(key_or_value: RedisKeyOrValue, prefixes: _PrefixTuple) -> str:
         """Turn a valid Redis type into a typestring."""
         for prefix, _type in prefixes:
-            if isinstance(key_or_value, _type):
+            # Convert bools into integers before storing them.
+            if type(key_or_value) is bool:
+                bool_int = int(key_or_value)
+                return f"{prefix}{bool_int}"
+
+            # isinstance is a bad idea here, because isintance(False, int) == True.
+            if type(key_or_value) is _type:
                 return f"{prefix}{key_or_value}"
+
         raise TypeError(f"RedisCache._to_typestring only supports the following: {prefixes}.")
 
     @staticmethod
@@ -122,6 +130,13 @@ class RedisCache:
         # Now we convert our unicode string back into the type it originally was.
         for prefix, _type in prefixes:
             if key_or_value.startswith(prefix):
+
+                # For booleans, we need special handling because bool("False") is True.
+                if prefix == "b|":
+                    value = key_or_value[len(prefix):]
+                    return bool(int(value))
+
+                # Otherwise we can just convert normally.
                 return _type(key_or_value[len(prefix):])
         raise TypeError(f"RedisCache._from_typestring only supports the following: {prefixes}.")
 
diff --git a/tests/bot/utils/test_redis_cache.py b/tests/bot/utils/test_redis_cache.py
index e5d6e4078..a2f0fe55d 100644
--- a/tests/bot/utils/test_redis_cache.py
+++ b/tests/bot/utils/test_redis_cache.py
@@ -49,7 +49,9 @@ class RedisCacheTests(unittest.IsolatedAsyncioTestCase):
         test_cases = (
             ('favorite_fruit', 'melon'),
             ('favorite_number', 86),
-            ('favorite_fraction', 86.54)
+            ('favorite_fraction', 86.54),
+            ('favorite_boolean', False),
+            ('other_boolean', True),
         )
 
         # Test that we can get and set different types.