Don't wait for websocket to close on Shard disconnect (#3369)

Author: GnomedDev

src/gateway/sharding/mod.rs

@@ -698,7 +698,8 @@ async fn connect(base_url: &str, compression: TransportCompression) -> Result<Ws
         Error::Gateway(GatewayError::BuildingUrl)
     })?;
 
-    WsClient::connect(url, compression).await
+    let client = WsClient::connect(url, compression).await?;
+    Ok(client)
 }
 
 struct ResumeMetadata {

src/gateway/sharding/shard_runner.rs

@@ -3,7 +3,6 @@ use std::sync::Arc;
 use dashmap::DashMap;
 use dashmap::try_result::TryResult;
 use futures::channel::mpsc::{self, UnboundedReceiver as Receiver, UnboundedSender as Sender};
-use tokio_tungstenite::tungstenite;
 use tokio_tungstenite::tungstenite::error::Error as TungsteniteError;
 use tokio_tungstenite::tungstenite::protocol::frame::CloseFrame;
 #[cfg(feature = "tracing_instrument")]
@@ -234,7 +233,13 @@ impl ShardRunner {
         }
     }
 
-    // Shuts down the WebSocket client.
+    /// Shuts down the WebSocket client.
+    ///
+    /// The Shard will be in an indeterminate state after this call, especially if called after
+    /// error.
+    ///
+    /// Therefore, the only correct code path is to exit out of the ShardRunner loop and discard the
+    /// WebSocket client entirely.
     #[cfg_attr(feature = "tracing_instrument", instrument(skip(self)))]
     async fn shutdown(&mut self, close_code: u16) {
         debug!("[ShardRunner {:?}] Shutting down.", self.shard.shard_info());
@@ -248,22 +253,6 @@ impl ShardRunner {
                 }))
                 .await,
         );
-
-        // In return, we wait for either a Close Frame response, or an error, after which this WS
-        // is deemed disconnected from Discord.
-        loop {
-            match self.shard.client.next().await {
-                Some(Ok(tungstenite::Message::Close(_))) => return,
-                Some(Err(_)) => {
-                    warn!(
-                        "[ShardRunner {:?}] Received an error awaiting close frame",
-                        self.shard.shard_info(),
-                    );
-                    return;
-                },
-                _ => {},
-            }
-        }
     }
 
     #[cfg(feature = "voice")]

src/gateway/ws.rs

@@ -20,6 +20,7 @@ use zstd_safe::{DStream as ZstdInflater, InBuffer, OutBuffer};
 
 use super::{ActivityData, ChunkGuildFilter, GatewayError, PresenceData, TransportCompression};
 use crate::constants::{self, Opcode};
+use crate::internal::prelude::*;
 use crate::model::event::GatewayEvent;
 use crate::model::gateway::{GatewayIntents, ShardInfo};
 #[cfg(feature = "voice")]
@@ -236,7 +237,10 @@ pub struct WsClient {
 const TIMEOUT: Duration = Duration::from_millis(500);
 
 impl WsClient {
-    pub(crate) async fn connect(url: Url, compression: TransportCompression) -> Result<Self> {
+    pub(crate) async fn connect(
+        url: Url,
+        compression: TransportCompression,
+    ) -> StdResult<Self, WsError> {
         let config = {
             let mut config = WebSocketConfig::default();
             config.max_message_size = None;
@@ -288,13 +292,8 @@ impl WsClient {
         Ok(())
     }
 
-    /// Delegate to `StreamExt::next`
-    pub(crate) async fn next(&mut self) -> Option<std::result::Result<Message, WsError>> {
-        self.stream.next().await
-    }
-
     /// Delegate to `WebSocketStream::close`
-    pub(crate) async fn close(&mut self, msg: Option<CloseFrame>) -> Result<()> {
+    pub(crate) async fn close(&mut self, msg: Option<CloseFrame>) -> StdResult<(), WsError> {
         self.stream.close(msg).await?;
         Ok(())
     }