Immediately return from run on shutdown

Workaround for serenity-rs/serenity#2507.

Cargo.toml

@@ -4,6 +4,6 @@ default-members = ["crate/serenity-utils"]
 resolver = "2"
 
 [workspace.package]
-version = "0.10.3"
+version = "0.10.4"
 authors = ["Fenhl <[email protected]>"]
 edition = "2021"

crate/serenity-utils/Cargo.toml

@@ -29,7 +29,7 @@ branch = "next"
 features = ["unstable_discord_api"]
 
 [dependencies.serenity-utils-derive]
-version = "=0.10.3"
+version = "=0.10.4"
 path = "../serenity-utils-derive"
 
 [dependencies.tokio]

crate/serenity-utils/src/builder.rs

@@ -28,7 +28,10 @@ use {
         model::prelude::*,
         prelude::*,
     },
-    tokio::time::sleep,
+    tokio::{
+        select,
+        time::sleep,
+    },
     crate::{
         RwFuture,
         handler::{
@@ -246,7 +249,8 @@ impl Builder {
     /// Convenience method wrapping `self` in [`Ok`] which can be used at the end of a method call chain.
     pub fn ok<E>(self) -> Result<Self, E> { Ok(self) }
 
-    #[doc(hidden)] pub async fn run(mut self) -> serenity::Result<()> { // used in `serenity_utils::main`
+    /// Runs the Discord bot until shutdown.
+    pub async fn run(mut self) -> serenity::Result<()> { // used in `serenity_utils::main`
         self.intents |= self.handler.intents;
         self.client = self.client.event_handler(self.handler);
         let mut client = self.client
@@ -256,8 +260,16 @@ impl Builder {
         {
             let mut data = client.data.write().await;
             data.insert::<crate::ShardManagerContainer>(Arc::clone(&client.shard_manager));
+            data.insert::<crate::ShutdownNotifier>(Arc::default());
+        }
+        let shutdown_notifier = Arc::clone(client.data.read().await.get::<crate::ShutdownNotifier>().expect("missing shutdown notifier"));
+        select! {
+            res = client.start_autosharded() => {
+                let () = res?;
+            }
+            // workaround for https://github.com/serenity-rs/serenity/issues/2507
+            () = shutdown_notifier.notified() => {}
         }
-        client.start_autosharded().await?;
         sleep(Duration::from_secs(1)).await; // wait to make sure websockets can be closed cleanly
         Ok(())
     }

crate/serenity-utils/src/lib.rs

@@ -125,12 +125,18 @@ impl<T: Send + Sync + Default> Default for RwFuture<T> {
 }
 
 /// A `typemap` key holding the [`ShardManager`]. Used in `shut_down`.
-pub struct ShardManagerContainer;
+pub enum ShardManagerContainer {}
 
 impl TypeMapKey for ShardManagerContainer {
     type Value = Arc<Mutex<ShardManager>>;
 }
 
+enum ShutdownNotifier {}
+
+impl TypeMapKey for ShutdownNotifier {
+    type Value = Arc<tokio::sync::Notify>;
+}
+
 /// Creates a builder for setting up and running a bot.
 ///
 /// An advantage of using this compared to constructing a [`Client`] manually is that the bot will automatically request the required intents.
@@ -142,6 +148,8 @@ pub async fn builder(token: String) -> serenity::Result<Builder> {
 pub async fn shut_down(ctx: &Context) {
     ctx.invisible(); // hack to prevent the bot showing as online when it's not
     let data = ctx.data.read().await;
+    // workaround for https://github.com/serenity-rs/serenity/issues/2507
+    data.get::<ShutdownNotifier>().expect("missing shutdown notifier").notify_waiters();
     let mut shard_manager = data.get::<ShardManagerContainer>().expect("missing shard manager").lock().await;
     shard_manager.shutdown_all().await;
     sleep(Duration::from_secs(1)).await; // wait to make sure websockets can be closed cleanly