diff options
| -rw-r--r-- | bot/cogs/help_channels.py | 124 | ||||
| -rw-r--r-- | bot/utils/redis_cache.py | 23 | ||||
| -rw-r--r-- | tests/bot/utils/test_redis_cache.py | 4 |
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. |