diff options
| -rw-r--r-- | bot/bot.py | 23 | ||||
| -rw-r--r-- | bot/exts/info/reddit.py | 2 | ||||
| -rw-r--r-- | bot/exts/moderation/infraction/_scheduler.py | 9 | ||||
| -rw-r--r-- | bot/exts/moderation/infraction/infractions.py | 16 | ||||
| -rw-r--r-- | bot/exts/moderation/watchchannels/_watchchannel.py | 17 |
5 files changed, 52 insertions, 15 deletions
diff --git a/bot/bot.py b/bot/bot.py index b2e5237fe..9a60474b3 100644 --- a/bot/bot.py +++ b/bot/bot.py @@ -3,7 +3,7 @@ import logging import socket import warnings from collections import defaultdict -from typing import Dict, Optional +from typing import Dict, List, Optional import aiohttp import discord @@ -48,6 +48,9 @@ class Bot(commands.Bot): self.stats = AsyncStatsClient(self.loop, statsd_url, 8125, prefix="bot") + # All tasks that need to block closing until finished + self.closing_tasks: List[asyncio.Task] = [] + async def cache_filter_list_data(self) -> None: """Cache all the data in the FilterList on the site.""" full_cache = await self.api_client.get('bot/filter-lists') @@ -131,8 +134,26 @@ class Bot(commands.Bot): self._recreate() super().clear() + def _remove_extensions(self) -> None: + """Remove all extensions to trigger cog unloads.""" + extensions = list(self.extensions.keys()) + + for ext in extensions: + try: + self.unload_extension(ext) + except Exception: + pass + async def close(self) -> None: """Close the Discord connection and the aiohttp session, connector, statsd client, and resolver.""" + # Done before super().close() to allow tasks finish before the HTTP session closes. + self._remove_extensions() + + # Wait until all tasks that have to be completed before bot is closing is done + log.trace("Waiting for tasks before closing.") + await asyncio.gather(*self.closing_tasks) + + # Now actually do full close of bot await super().close() await self.api_client.close() diff --git a/bot/exts/info/reddit.py b/bot/exts/info/reddit.py index bad4c504d..6790be762 100644 --- a/bot/exts/info/reddit.py +++ b/bot/exts/info/reddit.py @@ -45,7 +45,7 @@ class Reddit(Cog): """Stop the loop task and revoke the access token when the cog is unloaded.""" self.auto_poster_loop.cancel() if self.access_token and self.access_token.expires_at > datetime.utcnow(): - asyncio.create_task(self.revoke_access_token()) + self.bot.closing_tasks.append(asyncio.create_task(self.revoke_access_token())) async def init_reddit_ready(self) -> None: """Sets the reddit webhook when the cog is loaded.""" diff --git a/bot/exts/moderation/infraction/_scheduler.py b/bot/exts/moderation/infraction/_scheduler.py index bebade0ae..ed67e3b26 100644 --- a/bot/exts/moderation/infraction/_scheduler.py +++ b/bot/exts/moderation/infraction/_scheduler.py @@ -74,8 +74,13 @@ class InfractionScheduler: return # Allowing mod log since this is a passive action that should be logged. - await apply_coro - log.info(f"Re-applied {infraction['type']} to user {infraction['user']} upon rejoining.") + try: + await apply_coro + except discord.NotFound: + # When user joined and then right after this left again before action completed, this can't add roles + log.info(f"Can't reapply {infraction['type']} to user {infraction['user']} because user left again.") + else: + log.info(f"Re-applied {infraction['type']} to user {infraction['user']} upon rejoining.") async def apply_infraction( self, diff --git a/bot/exts/moderation/infraction/infractions.py b/bot/exts/moderation/infraction/infractions.py index 746d4e154..8abb199db 100644 --- a/bot/exts/moderation/infraction/infractions.py +++ b/bot/exts/moderation/infraction/infractions.py @@ -277,10 +277,18 @@ class Infractions(InfractionScheduler, commands.Cog): self.mod_log.ignore(Event.member_update, user.id) async def action() -> None: - await user.add_roles(self._muted_role, reason=reason) - - log.trace(f"Attempting to kick {user} from voice because they've been muted.") - await user.move_to(None, reason=reason) + try: + await user.add_roles(self._muted_role, reason=reason) + except discord.HTTPException as e: + if e.code == 10007: + log.info(f"User {user} ({user.id}) left from guild. Can't give Muted role.") + else: + log.warning( + f"Got response {e.code} (HTTP {e.status}) while giving muted role to {user} ({user.id})." + ) + else: + log.trace(f"Attempting to kick {user} from voice because they've been muted.") + await user.move_to(None, reason=reason) await self.apply_infraction(ctx, infraction, user, action()) diff --git a/bot/exts/moderation/watchchannels/_watchchannel.py b/bot/exts/moderation/watchchannels/_watchchannel.py index 7118dee02..b576f2888 100644 --- a/bot/exts/moderation/watchchannels/_watchchannel.py +++ b/bot/exts/moderation/watchchannels/_watchchannel.py @@ -342,11 +342,14 @@ class WatchChannel(metaclass=CogABCMeta): """Takes care of unloading the cog and canceling the consumption task.""" self.log.trace("Unloading the cog") if self._consume_task and not self._consume_task.done(): + def done_callback(task: asyncio.Task) -> None: + """Send exception when consuming task have been cancelled.""" + try: + task.exception() + except asyncio.CancelledError: + self.log.error( + f"The consume task of {type(self).__name__} was canceled. Messages may be lost." + ) + + self._consume_task.add_done_callback(done_callback) self._consume_task.cancel() - try: - self._consume_task.result() - except asyncio.CancelledError as e: - self.log.exception( - "The consume task was canceled. Messages may be lost.", - exc_info=e - ) |