Reconnect on failed heartbeats

2 files changed, 15 insertions, 7 deletions

diff --git a/src/client/mod.rs b/src/client/mod.rs
index 40859f2..f35a400 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -1328,7 +1328,11 @@ fn handle_shard(info: &mut MonitorInfo) {
         {
             let mut shard = info.shard.lock().unwrap();
 
-            shard.check_heartbeat();
+            if let Err(why) = shard.check_heartbeat() {
+                error!("Failed to heartbeat and reconnect: {:?}", why);
+
+                return;
+            }
         }
 
         #[cfg(feature="voice")]
diff --git a/src/gateway/shard.rs b/src/gateway/shard.rs
index 3477946..9c3403c 100644
--- a/src/gateway/shard.rs
+++ b/src/gateway/shard.rs
@@ -746,10 +746,10 @@ impl Shard {
         }
     }
 
-    pub(crate) fn check_heartbeat(&mut self) {
+    pub(crate) fn check_heartbeat(&mut self) -> Result<()> {
         let heartbeat_interval = match self.heartbeat_interval {
             Some(heartbeat_interval) => heartbeat_interval,
-            None => return,
+            None => return Ok(()),
         };
 
         let wait = StdDuration::from_secs(heartbeat_interval / 1000);
@@ -758,7 +758,7 @@ impl Shard {
         // then don't perform a keepalive or attempt to reconnect.
         if let Some(last_sent) = self.heartbeat_instants.0 {
             if last_sent.elapsed() <= wait {
-                return;
+                return Ok(());
             }
         }
 
@@ -773,15 +773,19 @@ impl Shard {
                       why);
             }
 
-            return;
+            return Ok(());
         }
 
         // Otherwise, we're good to heartbeat.
         if let Err(why) = self.heartbeat() {
             warn!("[Shard {:?}] Err heartbeating: {:?}", self.shard_info, why);
-        }
 
-        self.heartbeat_instants.0 = Some(Instant::now());
+            self.reconnect()
+        } else {
+            self.heartbeat_instants.0 = Some(Instant::now());
+
+            Ok(())
+        }
     }
 
     pub(crate) fn autoreconnect(&mut self) -> Result<()> {