Fix shard shutdown via Context

This fixes the shard shutdown via the Context. This works by setting a
"shutdown" field on the Shard struct as being true, indicating that the
Shard has shutdown. The Shard Runner detects this and requests a
shutdown from the Shard Manager.

The ideal solution here would be to add a "Shutdown" variant to
serenity::gateway::ConnectionStage, but that would be a breaking change,
so we instead need to opt for adding a new struct to gateway::Shard.

The Shard Manager has also been updated to only attempt the shutdown of
a shard if it doesn't already know for certain that it shut itself down,
which avoids an error logged saying that there was an error sending a
shutdown message to its Shard Runner.

When all shards have been shutdown (for most bots, this will only be
one), the Shard Manager will end and the Client will stop its
operations, returning thread control to the user.

4 files changed, 46 insertions, 8 deletions

diff --git a/src/client/bridge/gateway/shard_manager.rs b/src/client/bridge/gateway/shard_manager.rs
index 98ada1e..dfd4a16 100644
--- a/src/client/bridge/gateway/shard_manager.rs
+++ b/src/client/bridge/gateway/shard_manager.rs
@@ -174,13 +174,17 @@ impl ShardManager {
     }
 
     fn shutdown(&mut self, shard_id: ShardId) {
-        info!("Shutting down shard {}", shard_id);
-
         if let Some(runner) = self.runners.lock().get(&shard_id) {
-            let msg = ShardManagerMessage::Shutdown(shard_id);
+            let is_shutdown = runner.shard.lock().is_shutdown();
+
+            if !is_shutdown {
+                info!("Shutting down shard {}", shard_id);
+
+                let msg = ShardManagerMessage::Shutdown(shard_id);
 
-            if let Err(why) = runner.runner_tx.send(msg) {
-                warn!("Failed to cleanly shutdown shard {}: {:?}", shard_id, why);
+                if let Err(why) = runner.runner_tx.send(msg) {
+                    warn!("Failed to cleanly shutdown shard {}: {:?}", shard_id, why);
+                }
             }
         }
 
diff --git a/src/client/bridge/gateway/shard_runner.rs b/src/client/bridge/gateway/shard_runner.rs
index 7ee18c3..aa0e064 100644
--- a/src/client/bridge/gateway/shard_runner.rs
+++ b/src/client/bridge/gateway/shard_runner.rs
@@ -140,8 +140,16 @@ impl<H: EventHandler + Send + Sync + 'static> ShardRunner<H> {
                 }}
             }
 
-            if !successful && !self.shard.lock().stage().is_connecting() {
-                return self.request_restart();
+            {
+                let shard = self.shard.lock();
+
+                if !successful && !shard.stage().is_connecting() {
+                    return self.request_restart();
+                }
+
+                if shard.is_shutdown() {
+                    return self.request_shutdown();
+                }
             }
         }
     }
@@ -219,4 +227,11 @@ impl<H: EventHandler + Send + Sync + 'static> ShardRunner<H> {
 
         Ok(())
     }
+
+    fn request_shutdown(&self) -> Result<()> {
+        debug!("[ShardRunner {:?}] Requesting shutdown", self.shard_info);
+        let _ = self.manager_tx.send(ShardManagerMessage::ShutdownAll);
+
+        Ok(())
+    }
 }
diff --git a/src/gateway/mod.rs b/src/gateway/mod.rs
index b593f5c..ed54bca 100644
--- a/src/gateway/mod.rs
+++ b/src/gateway/mod.rs
@@ -132,7 +132,7 @@ impl ConnectionStage {
 
         match *self {
             Connecting | Handshake | Identifying | Resuming => true,
-            _ => false,
+            Connected | Disconnected => false,
         }
     }
 }
diff --git a/src/gateway/shard.rs b/src/gateway/shard.rs
index 3a8209d..e5fe65d 100644
--- a/src/gateway/shard.rs
+++ b/src/gateway/shard.rs
@@ -89,6 +89,8 @@ pub struct Shard {
     seq: u64,
     session_id: Option<String>,
     shard_info: [u64; 2],
+    /// Whether the shard has permanently shutdown.
+    shutdown: bool,
     stage: ConnectionStage,
     token: Arc<Mutex<String>>,
     ws_url: Arc<Mutex<String>>,
@@ -144,6 +146,7 @@ impl Shard {
                 let user = http::get_current_user()?;
 
                 Shard {
+                    shutdown: false,
                     client,
                     current_presence,
                     heartbeat_instants,
@@ -160,6 +163,7 @@ impl Shard {
                 }
             } else {
                 Shard {
+                    shutdown: false,
                     client,
                     current_presence,
                     heartbeat_instants,
@@ -176,6 +180,18 @@ impl Shard {
         })
     }
 
+    /// Whether the shard has permanently shutdown.
+    ///
+    /// This should normally happen due to manual calling of [`shutdown`] or
+    /// [`shutdown_clean`].
+    ///
+    /// [`shutdown`]: #method.shutdown
+    /// [`shutdown_clean`]: #method.shutdown_clean
+    #[inline]
+    pub fn is_shutdown(&self) -> bool {
+        self.shutdown
+    }
+
     /// Retrieves a copy of the current shard information.
     ///
     /// The first element is the _current_ shard - 0-indexed - while the second
@@ -604,6 +620,7 @@ impl Shard {
         stream.flush()?;
         stream.shutdown(Shutdown::Both)?;
 
+        self.shutdown = true;
         debug!("[Shard {:?}] Cleanly shutdown shard", self.shard_info);
 
         Ok(())
@@ -616,6 +633,8 @@ impl Shard {
         stream.flush()?;
         stream.shutdown(Shutdown::Both)?;
 
+        self.shutdown = true;
+
         Ok(())
     }