diff options
| -rw-r--r-- | bot/cogs/reminders.py | 6 | ||||
| -rw-r--r-- | bot/utils/scheduling.py | 36 |
2 files changed, 30 insertions, 12 deletions
diff --git a/bot/cogs/reminders.py b/bot/cogs/reminders.py index d96dedd20..603b627fb 100644 --- a/bot/cogs/reminders.py +++ b/bot/cogs/reminders.py @@ -74,9 +74,6 @@ class Reminders(Scheduler, Cog): log.debug(f"Deleting reminder {reminder_id} (the user has been reminded).") await self._delete_reminder(reminder_id) - # Now we can begone with it from our schedule list. - self.cancel_task(reminder_id) - async def _delete_reminder(self, reminder_id: str) -> None: """Delete a reminder from the database, given its ID, and cancel the running task.""" await self.bot.api_client.delete('bot/reminders/' + str(reminder_id)) @@ -86,7 +83,10 @@ class Reminders(Scheduler, Cog): async def _reschedule_reminder(self, reminder: dict) -> None: """Reschedule a reminder object.""" + log.trace(f"Cancelling old task #{reminder['id']}") self.cancel_task(reminder["id"]) + + log.trace(f"Scheduling new task #{reminder['id']}") self.schedule_task(reminder["id"], reminder) async def send_reminder(self, reminder: dict, late: relativedelta = None) -> None: diff --git a/bot/utils/scheduling.py b/bot/utils/scheduling.py index df46ccdd9..40d26249f 100644 --- a/bot/utils/scheduling.py +++ b/bot/utils/scheduling.py @@ -2,6 +2,7 @@ import asyncio import contextlib import logging from abc import abstractmethod +from functools import partial from typing import Dict from bot.utils import CogABCMeta @@ -36,6 +37,8 @@ class Scheduler(metaclass=CogABCMeta): `task_data` is passed to the `Scheduler._scheduled_task()` coroutine. """ + log.trace(f"{self.cog_name}: scheduling task #{task_id}...") + if task_id in self.scheduled_tasks: log.debug( f"{self.cog_name}: did not schedule task #{task_id}; task was already scheduled." @@ -43,13 +46,15 @@ class Scheduler(metaclass=CogABCMeta): return task = asyncio.create_task(self._scheduled_task(task_data)) - task.add_done_callback(_handle_task_exception) + task.add_done_callback(partial(self._task_done_callback, task_id)) self.scheduled_tasks[task_id] = task - log.debug(f"{self.cog_name}: scheduled task #{task_id}.") + log.debug(f"{self.cog_name}: scheduled task #{task_id} {id(task)}.") def cancel_task(self, task_id: str) -> None: """Un-schedules a task.""" + log.trace(f"{self.cog_name}: cancelling task #{task_id}...") + task = self.scheduled_tasks.get(task_id) if task is None: @@ -57,14 +62,27 @@ class Scheduler(metaclass=CogABCMeta): return task.cancel() - log.debug(f"{self.cog_name}: unscheduled task #{task_id}.") + log.debug(f"{self.cog_name}: unscheduled task #{task_id} {id(task)}.") del self.scheduled_tasks[task_id] + def _task_done_callback(self, task_id: str, task: asyncio.Task) -> None: + """ + Unschedule the task and raise its exception if one exists. + + If the task was cancelled, the CancelledError is retrieved and suppressed. In this case, + the task is already assumed to have been unscheduled. + """ + log.trace(f"{self.cog_name}: performing done callback for task #{task_id} {id(task)}") + + if task.cancelled(): + with contextlib.suppress(asyncio.CancelledError): + task.exception() + else: + # Check if it exists to avoid logging a warning. + if task_id in self.scheduled_tasks: + # Only cancel if the task is not cancelled to avoid a race condition when a new + # task is scheduled using the same ID. Reminders do this when re-scheduling after + # editing. + self.cancel_task(task_id) -def _handle_task_exception(task: asyncio.Task) -> None: - """Raise the task's exception, if any, unless the task is cancelled and has a CancelledError.""" - if task.cancelled(): - with contextlib.suppress(asyncio.CancelledError): task.exception() - else: - task.exception() |