Scheduler Redesign

bot/cogs/filtering.py

@@ -19,7 +19,6 @@
 )
 from bot.utils.redis_cache import RedisCache
 from bot.utils.scheduling import Scheduler
-from bot.utils.time import wait_until
 
 log = logging.getLogger(__name__)
 
@@ -60,16 +59,15 @@ def expand_spoilers(text: str) -> str:
 OFFENSIVE_MSG_DELETE_TIME = timedelta(days=Filter.offensive_msg_delete_days)
 
 
-class Filtering(Cog, Scheduler):
+class Filtering(Cog):
     """Filtering out invites, blacklisting domains, and warning us of certain regular expressions."""
 
     # Redis cache mapping a user ID to the last timestamp a bad nickname alert was sent
     name_alerts = RedisCache()
 
     def __init__(self, bot: Bot):
         self.bot = bot
-        super().__init__()
-
+        self.scheduler = Scheduler(self.__class__.__name__)
         self.name_lock = asyncio.Lock()
 
         staff_mistake_str = "If you believe this was a mistake, please let staff know!"
@@ -268,7 +266,7 @@ async def _filter_message(self, msg: Message, delta: Optional[int] = None) -> No
                             }
 
                             await self.bot.api_client.post('bot/offensive-messages', json=data)
-                            self.schedule_task(msg.id, data)
+                            self.schedule_msg_delete(data)
                             log.trace(f"Offensive message {msg.id} will be deleted on {delete_date}")
 
                         if is_private:
@@ -457,12 +455,10 @@ async def notify_member(self, filtered_member: Member, reason: str, channel: Tex
         except discord.errors.Forbidden:
             await channel.send(f"{filtered_member.mention} {reason}")
 
-    async def _scheduled_task(self, msg: dict) -> None:
+    def schedule_msg_delete(self, msg: dict) -> None:
         """Delete an offensive message once its deletion date is reached."""
         delete_at = dateutil.parser.isoparse(msg['delete_date']).replace(tzinfo=None)
-
-        await wait_until(delete_at)
-        await self.delete_offensive_msg(msg)
+        self.scheduler.schedule_at(delete_at, msg['id'], self.delete_offensive_msg(msg))
 
     async def reschedule_offensive_msg_deletion(self) -> None:
         """Get all the pending message deletion from the API and reschedule them."""
@@ -477,7 +473,7 @@ async def reschedule_offensive_msg_deletion(self) -> None:
             if delete_at < now:
                 await self.delete_offensive_msg(msg)
             else:
-                self.schedule_task(msg['id'], msg)
+                self.schedule_msg_delete(msg)
 
     async def delete_offensive_msg(self, msg: Mapping[str, str]) -> None:
         """Delete an offensive message, and then delete it from the db."""

bot/cogs/help_channels.py

@@ -1,5 +1,4 @@
 import asyncio
-import inspect
 import json
 import logging
 import random
@@ -57,14 +56,7 @@
 CoroutineFunc = t.Callable[..., t.Coroutine]
 
 
-class TaskData(t.NamedTuple):
-    """Data for a scheduled task."""
-
-    wait_time: int
-    callback: t.Awaitable
-
-
-class HelpChannels(Scheduler, commands.Cog):
+class HelpChannels(commands.Cog):
     """
     Manage the help channel system of the guild.
 
@@ -114,9 +106,8 @@ class HelpChannels(Scheduler, commands.Cog):
     claim_times = RedisCache()
 
     def __init__(self, bot: Bot):
-        super().__init__()
-
         self.bot = bot
+        self.scheduler = Scheduler(self.__class__.__name__)
 
         # Categories
         self.available_category: discord.CategoryChannel = None
@@ -145,7 +136,7 @@ def cog_unload(self) -> None:
         for task in self.queue_tasks:
             task.cancel()
 
-        self.cancel_all()
+        self.scheduler.cancel_all()
 
     def create_channel_queue(self) -> asyncio.Queue:
         """
@@ -229,10 +220,11 @@ async def close_command(self, ctx: commands.Context) -> None:
                 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)
+                if ctx.author.id in self.scheduler:
+                    self.scheduler.cancel(ctx.author.id)
 
                 await self.move_to_dormant(ctx.channel, "command")
-                self.cancel_task(ctx.channel.id)
+                self.scheduler.cancel(ctx.channel.id)
         else:
             log.debug(f"{ctx.author} invoked command 'dormant' outside an in-use help channel")
 
@@ -474,16 +466,15 @@ async def move_idle_channel(self, channel: discord.TextChannel, has_task: bool =
         else:
             # Cancel the existing task, if any.
             if has_task:
-                self.cancel_task(channel.id)
-
-            data = TaskData(idle_seconds - time_elapsed, self.move_idle_channel(channel))
+                self.scheduler.cancel(channel.id)
 
+            delay = idle_seconds - time_elapsed
             log.info(
                 f"#{channel} ({channel.id}) is still active; "
-                f"scheduling it to be moved after {data.wait_time} seconds."
+                f"scheduling it to be moved after {delay} seconds."
             )
 
-            self.schedule_task(channel.id, data)
+            self.scheduler.schedule_later(delay, channel.id, self.move_idle_channel(channel))
 
     async def move_to_bottom_position(self, channel: discord.TextChannel, category_id: int, **options) -> None:
         """
@@ -588,8 +579,7 @@ async def move_to_in_use(self, channel: discord.TextChannel) -> None:
         timeout = constants.HelpChannels.idle_minutes * 60
 
         log.trace(f"Scheduling #{channel} ({channel.id}) to become dormant in {timeout} sec.")
-        data = TaskData(timeout, self.move_idle_channel(channel))
-        self.schedule_task(channel.id, data)
+        self.scheduler.schedule_later(timeout, channel.id, self.move_idle_channel(channel))
         self.report_stats()
 
     async def notify(self) -> None:
@@ -722,10 +712,10 @@ async def on_message_delete(self, msg: discord.Message) -> None:
         log.info(f"Claimant of #{msg.channel} ({msg.author}) deleted message, channel is empty now. Rescheduling task.")
 
         # Cancel existing dormant task before scheduling new.
-        self.cancel_task(msg.channel.id)
+        self.scheduler.cancel(msg.channel.id)
 
-        task = TaskData(constants.HelpChannels.deleted_idle_minutes * 60, self.move_idle_channel(msg.channel))
-        self.schedule_task(msg.channel.id, task)
+        delay = constants.HelpChannels.deleted_idle_minutes * 60
+        self.scheduler.schedule_later(delay, msg.channel.id, self.move_idle_channel(msg.channel))
 
     async def is_empty(self, channel: discord.TextChannel) -> bool:
         """Return True if the most recent message in `channel` is the bot's `AVAILABLE_MSG`."""
@@ -752,8 +742,8 @@ async def check_cooldowns(self) -> None:
                 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)
+                delay = cooldown - in_use_time.seconds
+                self.scheduler.schedule_later(delay, member.id, self.remove_cooldown_role(member))
 
     async def add_cooldown_role(self, member: discord.Member) -> None:
         """Add the help cooldown role to `member`."""
@@ -804,16 +794,11 @@ async def revoke_send_permissions(self, member: discord.Member) -> None:
 
         # Cancel the existing task, if any.
         # Would mean the user somehow bypassed the lack of permissions (e.g. user is guild owner).
-        self.cancel_task(member.id, ignore_missing=True)
+        if member.id in self.scheduler:
+            self.scheduler.cancel(member.id)
 
-        await self.schedule_cooldown_expiration(member, constants.HelpChannels.claim_minutes * 60)
-
-    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))
+        delay = constants.HelpChannels.claim_minutes * 60
+        self.scheduler.schedule_later(delay, member.id, self.remove_cooldown_role(member))
 
     async def send_available_message(self, channel: discord.TextChannel) -> None:
         """Send the available message by editing a dormant message or sending a new message."""
@@ -855,21 +840,6 @@ async def wait_for_dormant_channel(self) -> discord.TextChannel:
 
         return channel
 
-    async def _scheduled_task(self, data: TaskData) -> None:
-        """Await the `data.callback` coroutine after waiting for `data.wait_time` seconds."""
-        try:
-            log.trace(f"Waiting {data.wait_time} seconds before awaiting callback.")
-            await asyncio.sleep(data.wait_time)
-
-            # Use asyncio.shield to prevent callback from cancelling itself.
-            # The parent task (_scheduled_task) will still get cancelled.
-            log.trace("Done waiting; now awaiting the callback.")
-            await asyncio.shield(data.callback)
-        finally:
-            if inspect.iscoroutine(data.callback):
-                log.trace("Explicitly closing coroutine.")
-                data.callback.close()
-
 
 def validate_config() -> None:
     """Raise a ValueError if the cog's config is invalid."""

bot/cogs/moderation/management.py

@@ -135,11 +135,11 @@ async def infraction_edit(
         if 'expires_at' in request_data:
             # A scheduled task should only exist if the old infraction wasn't permanent
             if old_infraction['expires_at']:
-                self.infractions_cog.cancel_task(new_infraction['id'])
+                self.infractions_cog.scheduler.cancel(new_infraction['id'])
 
             # If the infraction was not marked as permanent, schedule a new expiration task
             if request_data['expires_at']:
-                self.infractions_cog.schedule_task(new_infraction['id'], new_infraction)
+                self.infractions_cog.scheduler.schedule(new_infraction['id'], new_infraction)
 
             log_text += f"""
                 Previous expiry: {old_infraction['expires_at'] or "Permanent"}

bot/cogs/moderation/scheduler.py

@@ -1,4 +1,3 @@
-import asyncio
 import logging
 import textwrap
 import typing as t
@@ -23,13 +22,13 @@
 log = logging.getLogger(__name__)
 
 
-class InfractionScheduler(Scheduler):
+class InfractionScheduler:
     """Handles the application, pardoning, and expiration of infractions."""
 
     def __init__(self, bot: Bot, supported_infractions: t.Container[str]):
-        super().__init__()
-
         self.bot = bot
+        self.scheduler = Scheduler(self.__class__.__name__)
+
         self.bot.loop.create_task(self.reschedule_infractions(supported_infractions))
 
     @property
@@ -49,7 +48,7 @@ async def reschedule_infractions(self, supported_infractions: t.Container[str])
         )
         for infraction in infractions:
             if infraction["expires_at"] is not None and infraction["type"] in supported_infractions:
-                self.schedule_task(infraction["id"], infraction)
+                self.schedule_expiration(infraction)
 
     async def reapply_infraction(
         self,
@@ -155,7 +154,7 @@ async def apply_infraction(
                 await action_coro
                 if expiry:
                     # Schedule the expiration of the infraction.
-                    self.schedule_task(infraction["id"], infraction)
+                    self.schedule_expiration(infraction)
             except discord.HTTPException as e:
                 # Accordingly display that applying the infraction failed.
                 confirm_msg = ":x: failed to apply"
@@ -278,7 +277,7 @@ async def pardon_infraction(
 
                 # Cancel pending expiration task.
                 if infraction["expires_at"] is not None:
-                    self.cancel_task(infraction["id"])
+                    self.scheduler.cancel(infraction["id"])
 
         # Accordingly display whether the user was successfully notified via DM.
         dm_emoji = ""
@@ -415,7 +414,7 @@ async def deactivate_infraction(
 
         # Cancel the expiration task.
         if infraction["expires_at"] is not None:
-            self.cancel_task(infraction["id"])
+            self.scheduler.cancel(infraction["id"])
 
         # Send a log message to the mod log.
         if send_log:
@@ -449,16 +448,12 @@ async def _pardon_action(self, infraction: utils.Infraction) -> t.Optional[t.Dic
         """
         raise NotImplementedError
 
-    async def _scheduled_task(self, infraction: utils.Infraction) -> None:
+    def schedule_expiration(self, infraction: utils.Infraction) -> None:
         """
         Marks an infraction expired after the delay from time of scheduling to time of expiration.
 
         At the time of expiration, the infraction is marked as inactive on the website and the
         expiration task is cancelled.
         """
         expiry = dateutil.parser.isoparse(infraction["expires_at"]).replace(tzinfo=None)
-        await time.wait_until(expiry)
-
-        # Because deactivate_infraction() explicitly cancels this scheduled task, it is shielded
-        # to avoid prematurely cancelling itself.
-        await asyncio.shield(self.deactivate_infraction(infraction))
+        self.scheduler.schedule_at(expiry, infraction["id"], self.deactivate_infraction(infraction))

bot/cogs/moderation/silence.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from contextlib import suppress
-from typing import NamedTuple, Optional
+from typing import Optional
 
 from discord import TextChannel
 from discord.ext import commands, tasks
@@ -16,13 +16,6 @@
 log = logging.getLogger(__name__)
 
 
-class TaskData(NamedTuple):
-    """Data for a scheduled task."""
-
-    delay: int
-    ctx: Context
-
-
 class SilenceNotifier(tasks.Loop):
     """Loop notifier for posting notices to `alert_channel` containing added channels."""
 
@@ -61,25 +54,17 @@ async def _notifier(self) -> None:
             await self._alert_channel.send(f"<@&{Roles.moderators}> currently silenced channels: {channels_text}")
 
 
-class Silence(Scheduler, commands.Cog):
+class Silence(commands.Cog):
     """Commands for stopping channel messages for `verified` role in a channel."""
 
     def __init__(self, bot: Bot):
-        super().__init__()
         self.bot = bot
+        self.scheduler = Scheduler(self.__class__.__name__)
         self.muted_channels = set()
+
         self._get_instance_vars_task = self.bot.loop.create_task(self._get_instance_vars())
         self._get_instance_vars_event = asyncio.Event()
 
-    async def _scheduled_task(self, task: TaskData) -> None:
-        """Calls `self.unsilence` on expired silenced channel to unsilence it."""
-        await asyncio.sleep(task.delay)
-        log.info("Unsilencing channel after set delay.")
-
-        # Because `self.unsilence` explicitly cancels this scheduled task, it is shielded
-        # to avoid prematurely cancelling itself
-        await asyncio.shield(task.ctx.invoke(self.unsilence))
-
     async def _get_instance_vars(self) -> None:
         """Get instance variables after they're available to get from the guild."""
         await self.bot.wait_until_guild_available()
@@ -109,12 +94,7 @@ async def silence(self, ctx: Context, duration: HushDurationConverter = 10) -> N
 
         await ctx.send(f"{Emojis.check_mark} silenced current channel for {duration} minute(s).")
 
-        task_data = TaskData(
-            delay=duration*60,
-            ctx=ctx
-        )
-
-        self.schedule_task(ctx.channel.id, task_data)
+        self.scheduler.schedule_later(duration * 60, ctx.channel.id, ctx.invoke(self.unsilence))
 
     @commands.command(aliases=("unhush",))
     async def unsilence(self, ctx: Context) -> None:
@@ -164,7 +144,7 @@ async def _unsilence(self, channel: TextChannel) -> bool:
         if current_overwrite.send_messages is False:
             await channel.set_permissions(self._verified_role, **dict(current_overwrite, send_messages=None))
             log.info(f"Unsilenced channel #{channel} ({channel.id}).")
-            self.cancel_task(channel.id)
+            self.scheduler.cancel(channel.id)
             self.notifier.remove_channel(channel)
             self.muted_channels.discard(channel)
             return True

bot/cogs/moderation/superstarify.py

@@ -146,7 +146,7 @@ async def superstarify(
         log.debug(f"Changing nickname of {member} to {forced_nick}.")
         self.mod_log.ignore(constants.Event.member_update, member.id)
         await member.edit(nick=forced_nick, reason=reason)
-        self.schedule_task(id_, infraction)
+        self.schedule_expiration(infraction)
 
         # Send a DM to the user to notify them of their new infraction.
         await utils.notify_infraction(