Add a connection timeout

Add a timeout on new connections of 15 seconds. If 15 seconds comes and
no heartbeat interval has yet been received (indicating a lack of a
Hello message being received), then the shard will indicate that it is
in a failed state and needs to be restarted.

If a Hello is received but the handshake is not further completed to a
Ready stage, then a heartbeat acknowledgement check will clean up the
shard and indicate its failed status.

1 files changed, 9 insertions, 1 deletions

diff --git a/src/gateway/shard.rs b/src/gateway/shard.rs
index 8f6969d..aee2bd0 100644
--- a/src/gateway/shard.rs
+++ b/src/gateway/shard.rs
@@ -79,6 +79,10 @@ pub struct Shard {
     /// Whether the shard has permanently shutdown.
     shutdown: bool,
     stage: ConnectionStage,
+    /// Instant of when the shard was started.
+    // This acts as a timeout to determine if the shard has - for some reason -
+    // not started within a decent amount of time.
+    pub started: Instant,
     pub token: Arc<Mutex<String>>,
     ws_url: Arc<Mutex<String>>,
 }
@@ -147,6 +151,7 @@ impl Shard {
             last_heartbeat_acknowledged,
             seq,
             stage,
+            started: Instant::now(),
             token,
             session_id,
             shard_info,
@@ -567,7 +572,9 @@ impl Shard {
         let wait = {
             let heartbeat_interval = match self.heartbeat_interval {
                 Some(heartbeat_interval) => heartbeat_interval,
-                None => return true,
+                None => {
+                    return self.started.elapsed() < StdDuration::from_secs(15);
+                },
             };
 
             StdDuration::from_secs(heartbeat_interval / 1000)
@@ -769,6 +776,7 @@ impl Shard {
         // This is used to accurately assess whether the state of the shard is
         // accurate when a Hello is received.
         self.stage = ConnectionStage::Connecting;
+        self.started = Instant::now();
         let mut client = connect(&self.ws_url.lock())?;
         self.stage = ConnectionStage::Handshake;