Fix disconnections on successful gateway resume

Only break of out the `ShardRunner` loop if the resume failed and we have to reidentify.
mkrasnitski · Feb 3, 2025 · 10a47d6 · 10a47d6
1 parent f0072d1
commit 10a47d6
Showing 1 changed file with 24 additions and 16 deletions.
diff --git a/src/gateway/sharding/shard_runner.rs b/src/gateway/sharding/shard_runner.rs
@@ -141,8 +141,9 @@ impl ShardRunner {
             if let Some(action) = action {
                 match action {
                     ShardAction::Reconnect => {
-                        self.reconnect().await;
-                        return Ok(());
+                        if !self.reconnect().await {
+                            return Ok(());
+                        }
                     },
                     ShardAction::Heartbeat => {
                         if let Err(e) = self.shard.heartbeat().await {
@@ -151,8 +152,9 @@ impl ShardRunner {
                             self.shard.shard_info(),
                             e
                         );
-                            self.reconnect().await;
-                            return Ok(());
+                            if !self.reconnect().await {
+                                return Ok(());
+                            }
                         }
                     },
                     ShardAction::Identify => {
@@ -162,8 +164,9 @@ impl ShardRunner {
                             self.shard.shard_info(),
                             e
                         );
-                            self.reconnect().await;
-                            return Ok(());
+                            if !self.reconnect().await {
+                                return Ok(());
+                            }
                         }
                     },
                     ShardAction::Dispatch(event) => {
@@ -442,22 +445,27 @@ impl ShardRunner {
     }
 
     #[cfg_attr(feature = "tracing_instrument", instrument(skip(self)))]
-    async fn reconnect(&mut self) {
+    async fn reconnect(&mut self) -> bool {
         if self.shard.session_id().is_some() {
-            if let Err(why) = self.shard.resume().await {
-                warn!(
-                    "[ShardRunner {:?}] Resume failed, reidentifying: {:?}",
-                    self.shard.shard_info(),
-                    why,
-                );
+            match self.shard.resume().await {
+                Ok(()) => true,
+                Err(why) => {
+                    warn!(
+                        "[ShardRunner {:?}] Resume failed, reidentifying: {:?}",
+                        self.shard.shard_info(),
+                        why,
+                    );
 
-                // Don't spam reattempts on internet connection loss
-                tokio::time::sleep(std::time::Duration::from_secs(1)).await;
+                    // Don't spam reattempts on internet connection loss
+                    tokio::time::sleep(std::time::Duration::from_secs(1)).await;
 
-                self.request_restart().await;
+                    self.request_restart().await;
+                    false
+                },
             }
         } else {
             self.request_restart().await;
+            false
         }
     }