Modernize cogs by using async cog_load and cog_unload

docs/framework_modlog.rst

@@ -35,8 +35,7 @@ Basic Usage
 Registering Case types
 **********************
 
-To register case types, use an asynchronous ``initialize()`` method and call
-it from your setup function:
+To register case types, use a special ``cog_load()`` method which is called when you add a cog:
 
 .. code-block:: python
 
@@ -46,7 +45,7 @@ it from your setup function:
 
     class MyCog(commands.Cog):
 
-        async def initialize(self):
+        async def cog_load(self):
             await self.register_casetypes()
 
         @staticmethod
@@ -87,7 +86,6 @@ it from your setup function:
 
     async def setup(bot):
         cog = MyCog()
-        await cog.initialize()
         await bot.add_cog(cog)
 
 .. important::

redbot/cogs/admin/admin.py

@@ -84,19 +84,15 @@ def __init__(self, bot):
         )
 
         self.__current_announcer = None
-        self._ready = asyncio.Event()
-        asyncio.create_task(self.handle_migrations())
-        # As this is a data migration, don't store this for cancelation.
 
-    async def cog_before_invoke(self, ctx: commands.Context):
-        await self._ready.wait()
+    async def cog_load(self) -> None:
+        await self.handle_migrations()
 
     async def red_delete_data_for_user(self, **kwargs):
         """ Nothing to delete """
         return
 
     async def handle_migrations(self):
-
         lock = self.config.get_guilds_lock()
         async with lock:
             # This prevents the edge case of someone loading admin,
@@ -107,10 +103,7 @@ async def handle_migrations(self):
                 await self.migrate_config_from_0_to_1()
                 await self.config.schema_version.set(1)
 
-        self._ready.set()
-
-    async def migrate_config_from_0_to_1(self):
-
+    async def migrate_config_from_0_to_1(self) -> None:
         all_guilds = await self.config.all_guilds()
 
         for guild_id, guild_data in all_guilds.items():

redbot/cogs/alias/__init__.py

@@ -3,6 +3,4 @@
 
 
 async def setup(bot: Red) -> None:
-    cog = Alias(bot)
-    await bot.add_cog(cog)
-    cog.sync_init()
+    await bot.add_cog(Alias(bot))

redbot/cogs/alias/alias.py

@@ -51,7 +51,12 @@ def __init__(self, bot: Red):
         self.config.register_global(entries=[], handled_string_creator=False)
         self.config.register_guild(entries=[])
         self._aliases: AliasCache = AliasCache(config=self.config, cache_enabled=True)
-        self._ready_event = asyncio.Event()
+
+    async def cog_load(self) -> None:
+        await self._maybe_handle_string_keys()
+
+        if not self._aliases._loaded:
+            await self._aliases.load_aliases()
 
     async def red_delete_data_for_user(
         self,
@@ -62,12 +67,8 @@ async def red_delete_data_for_user(
         if requester != "discord_deleted_user":
             return
 
-        await self._ready_event.wait()
         await self._aliases.anonymize_aliases(user_id)
 
-    async def cog_before_invoke(self, ctx):
-        await self._ready_event.wait()
-
     async def _maybe_handle_string_keys(self):
         # This isn't a normal schema migration because it's being added
         # after the fact for GH-3788
@@ -95,12 +96,10 @@ async def _maybe_handle_string_keys(self):
         all_guild_aliases = await self.config.all_guilds()
 
         for guild_id, guild_data in all_guild_aliases.items():
-
             to_set = []
             modified = False
 
             for a in guild_data.get("entries", []):
-
                 for keyname in ("creator", "guild"):
                     if isinstance((val := a.get(keyname)), str):
                         try:
@@ -122,28 +121,6 @@ async def _maybe_handle_string_keys(self):
 
         await self.config.handled_string_creator.set(True)
 
-    def sync_init(self):
-        t = asyncio.create_task(self._initialize())
-
-        def done_callback(fut: asyncio.Future):
-            try:
-                t.result()
-            except Exception as exc:
-                log.exception("Failed to load alias cog", exc_info=exc)
-                # Maybe schedule extension unloading with message to owner in future
-
-        t.add_done_callback(done_callback)
-
-    async def _initialize(self):
-        """ Should only ever be a task """
-
-        await self._maybe_handle_string_keys()
-
-        if not self._aliases._loaded:
-            await self._aliases.load_aliases()
-
-        self._ready_event.set()
-
     def is_command(self, alias_name: str) -> bool:
         """
         The logic here is that if this returns true, the name should not be used for an alias
@@ -483,9 +460,6 @@ async def _list_global_alias(self, ctx: commands.Context):
 
     @commands.Cog.listener()
     async def on_message_without_command(self, message: discord.Message):
-
-        await self._ready_event.wait()
-
         if message.guild is not None:
             if await self.bot.cog_disabled_in_guild(self, message.guild):
                 return

redbot/cogs/audio/core/events/dpy.py

@@ -245,13 +245,11 @@ async def cog_command_error(self, ctx: commands.Context, error: Exception) -> No
         if not handled:
             await self.bot.on_command_error(ctx, error, unhandled_by_cog=True)
 
-    def cog_unload(self) -> None:
+    async def cog_unload(self) -> None:
         if not self.cog_cleaned_up:
             self.bot.dispatch("red_audio_unload", self)
-            self.session.detach()
-            self.bot.loop.create_task(self._close_database()).add_done_callback(
-                task_callback_trace
-            )
+            await self.session.close()
+            await self._close_database()
             if self.player_automated_timer_task:
                 self.player_automated_timer_task.cancel()
 
@@ -266,13 +264,9 @@ def cog_unload(self) -> None:
 
             lavalink.unregister_event_listener(self.lavalink_event_handler)
             lavalink.unregister_update_listener(self.lavalink_update_handler)
-            self.bot.loop.create_task(lavalink.close(self.bot)).add_done_callback(
-                task_callback_trace
-            )
+            await lavalink.close(self.bot)
             if self.player_manager is not None:
-                self.bot.loop.create_task(self.player_manager.shutdown()).add_done_callback(
-                    task_callback_trace
-                )
+                await self.player_manager.shutdown()
 
             self.cog_cleaned_up = True
 

redbot/cogs/filter/__init__.py

@@ -3,6 +3,4 @@
 
 
 async def setup(bot: Red) -> None:
-    cog = Filter(bot)
-    await cog.initialize()
-    await bot.add_cog(cog)
+    await bot.add_cog(Filter(bot))

redbot/cogs/filter/filter.py

@@ -56,7 +56,7 @@ async def red_delete_data_for_user(
             if user_id in guild_data:
                 await self.config.member_from_ids(guild_id, user_id).clear()
 
-    async def initialize(self) -> None:
+    async def cog_load(self) -> None:
         await self.register_casetypes()
 
     @staticmethod

redbot/cogs/image/__init__.py

@@ -4,6 +4,4 @@
 
 
 async def setup(bot: Red) -> None:
-    cog = Image(bot)
-    await cog.initialize()
-    await bot.add_cog(cog)
+    await bot.add_cog(Image(bot))

redbot/cogs/image/image.py

@@ -24,21 +24,21 @@ def __init__(self, bot):
         self.session = aiohttp.ClientSession()
         self.imgur_base_url = "https://api.imgur.com/3/"
 
-    def cog_unload(self):
-        self.session.detach()
-
-    async def red_delete_data_for_user(self, **kwargs):
-        """ Nothing to delete """
-        return
-
-    async def initialize(self) -> None:
+    async def cog_load(self) -> None:
         """Move the API keys from cog stored config to core bot config if they exist."""
         imgur_token = await self.config.imgur_client_id()
         if imgur_token is not None:
             if not await self.bot.get_shared_api_tokens("imgur"):
                 await self.bot.set_shared_api_tokens("imgur", client_id=imgur_token)
             await self.config.imgur_client_id.clear()
 
+    async def cog_unload(self):
+        await self.session.close()
+
+    async def red_delete_data_for_user(self, **kwargs):
+        """ Nothing to delete """
+        return
+
     @commands.group(name="imgur")
     async def _imgur(self, ctx):
         """Retrieve pictures from Imgur.

redbot/cogs/mod/__init__.py

@@ -3,6 +3,4 @@
 
 
 async def setup(bot: Red) -> None:
-    cog = Mod(bot)
-    await bot.add_cog(cog)
-    await cog.initialize()
+    await bot.add_cog(Mod(bot))

redbot/cogs/mod/mod.py

@@ -84,8 +84,6 @@ def __init__(self, bot: Red):
         self.tban_expiry_task = asyncio.create_task(self.tempban_expirations_task())
         self.last_case: dict = defaultdict(dict)
 
-        self._ready = asyncio.Event()
-
     async def red_delete_data_for_user(
         self,
         *,
@@ -114,12 +112,8 @@ async def red_delete_data_for_user(
                         pass
                     # possible with a context switch between here and getting all guilds
 
-    async def initialize(self):
+    async def cog_load(self) -> None:
         await self._maybe_update_config()
-        self._ready.set()
-
-    async def cog_before_invoke(self, ctx: commands.Context) -> None:
-        await self._ready.wait()
 
     def cog_unload(self):
         self.tban_expiry_task.cancel()

redbot/cogs/mutes/__init__.py

@@ -5,3 +5,4 @@
 async def setup(bot: Red) -> None:
     cog = Mutes(bot)
     await bot.add_cog(cog)
+    cog.create_init_task()

redbot/cogs/mutes/mutes.py

@@ -88,28 +88,42 @@ def __init__(self, bot: Red):
             "dm": False,
             "show_mod": False,
         }
-        self.config.register_global(force_role_mutes=True, schema_version=0)
         # Tbh I would rather force everyone to use role mutes.
         # I also honestly think everyone would agree they're the
         # way to go. If for whatever reason someone wants to
         # enable channel overwrite mutes for their bot they can.
         # Channel overwrite logic still needs to be in place
         # for channel mutes methods.
+        self.config.register_global(force_role_mutes=True, schema_version=0)
         self.config.register_guild(**default_guild)
         self.config.register_member(perms_cache={})
         self.config.register_channel(muted_users={})
         self._server_mutes: Dict[int, Dict[int, dict]] = {}
         self._channel_mutes: Dict[int, Dict[int, dict]] = {}
-        self._ready = asyncio.Event()
         self._unmute_tasks: Dict[str, asyncio.Task] = {}
-        self._unmute_task = None
+        self._unmute_task: Optional[asyncio.Task] = None
         self.mute_role_cache: Dict[int, int] = {}
-        self._channel_mute_events: Dict[int, asyncio.Event] = {}
         # this is a dict of guild ID's and asyncio.Events
         # to wait for a guild to finish channel unmutes before
         # checking for manual overwrites
+        self._channel_mute_events: Dict[int, asyncio.Event] = {}
+        self._ready = asyncio.Event()
+        self._init_task: Optional[asyncio.Task] = None
+        self._ready_raised = False
+
+    def create_init_task(self) -> None:
+        def _done_callback(task: asyncio.Task) -> None:
+            exc = task.exception()
+            if exc is not None:
+                log.error(
+                    "An unexpected error occurred during Mutes's initialization.",
+                    exc_info=exc,
+                )
+                self._ready_raised = True
+                self._ready.set()
 
-        self._init_task = self.bot.loop.create_task(self._initialize())
+        self._init_task = asyncio.create_task(self.initialize())
+        self._init_task.add_done_callback(_done_callback)
 
     async def red_delete_data_for_user(
         self,
@@ -126,13 +140,17 @@ async def red_delete_data_for_user(
             return
 
         await self._ready.wait()
+        if self._ready_raised:
+            raise RuntimeError(
+                "Mutes cog is in a bad state, can't proceed with data deletion request."
+            )
         all_members = await self.config.all_members()
         for g_id, data in all_members.items():
             for m_id, mutes in data.items():
                 if m_id == user_id:
                     await self.config.member_from_ids(g_id, m_id).clear()
 
-    async def _initialize(self):
+    async def initialize(self):
         await self.bot.wait_until_red_ready()
         await self._maybe_update_config()
 
@@ -185,11 +203,21 @@ async def _schema_0_to_1(self):
         )
 
     async def cog_before_invoke(self, ctx: commands.Context):
-        await self._ready.wait()
+        if not self._ready.is_set():
+            async with ctx.typing():
+                await self._ready.wait()
+        if self._ready_raised:
+            await ctx.send(
+                "There was an error during Mutes's initialization."
+                " Check logs for more information."
+            )
+            raise commands.CheckFailure()
 
     def cog_unload(self):
-        self._init_task.cancel()
-        self._unmute_task.cancel()
+        if self._init_task is not None:
+            self._init_task.cancel()
+        if self._unmute_task is not None:
+            self._unmute_task.cancel()
         for task in self._unmute_tasks.values():
             task.cancel()
 
@@ -209,6 +237,10 @@ async def _handle_automatic_unmute(self):
         """
         await self.bot.wait_until_red_ready()
         await self._ready.wait()
+        if self._ready_raised:
+            raise RuntimeError(
+                "Mutes cog is in a bad state, cancelling automatic unmute task."
+            )
         while True:
             await self._clean_tasks()
             try:

redbot/cogs/permissions/permissions.py

@@ -819,8 +819,8 @@ def _load_rules_for(
                 elif rule is False:
                     cog_or_command.deny_to(model_id, guild_id=guild_id)
 
-    def cog_unload(self) -> None:
-        self.bot.loop.create_task(self._unload_all_rules())
+    async def cog_unload(self) -> None:
+        await self._unload_all_rules()
 
     async def _unload_all_rules(self) -> None:
         """Unload all rules set by this cog.