Silence: cancel init task when cog unloads

2 files changed, 12 insertions, 4 deletions

diff --git a/bot/cogs/moderation/silence.py b/bot/cogs/moderation/silence.py
index 5851be00a..c339fd4d0 100644
--- a/bot/cogs/moderation/silence.py
+++ b/bot/cogs/moderation/silence.py
@@ -231,8 +231,13 @@ class Silence(commands.Cog):
                 self.scheduler.schedule_later(delta, channel_id, self._unsilence_wrapper(channel))
 
     def cog_unload(self) -> None:
-        """Cancel scheduled tasks."""
-        self.scheduler.cancel_all()
+        """Cancel the init task and scheduled tasks."""
+        # It's important to wait for _init_task (specifically for _reschedule) to be cancelled
+        # before cancelling scheduled tasks. Otherwise, it's possible for _reschedule to schedule
+        # more tasks after cancel_all has finished, despite _init_task.cancel being called first.
+        # This is cause cancel() on its own doesn't block until the task is cancelled.
+        self._init_task.cancel()
+        self._init_task.add_done_callback(lambda _: self.scheduler.cancel_all)
 
     # This cannot be static (must have a __func__ attribute).
     def cog_check(self, ctx: Context) -> bool:
diff --git a/tests/bot/cogs/moderation/test_silence.py b/tests/bot/cogs/moderation/test_silence.py
index a66d27d08..d56a731b6 100644
--- a/tests/bot/cogs/moderation/test_silence.py
+++ b/tests/bot/cogs/moderation/test_silence.py
@@ -127,9 +127,12 @@ class SilenceCogTests(unittest.IsolatedAsyncioTestCase):
         self.cog._reschedule.assert_awaited_once_with()
 
     def test_cog_unload_cancelled_tasks(self):
-        """All scheduled tasks were cancelled."""
+        """The init task was cancelled."""
+        self.cog._init_task = asyncio.Future()
         self.cog.cog_unload()
-        self.cog.scheduler.cancel_all.assert_called_once_with()
+
+        # It's too annoying to test cancel_all since it's a done callback and wrapped in a lambda.
+        self.assertTrue(self.cog._init_task.cancelled())
 
     @autospec(silence, "with_role_check")
     @mock.patch.object(silence, "MODERATION_ROLES", new=(1, 2, 3))