Resume when restarting WS sender/receiver

3 files changed, 37 insertions, 30 deletions

diff --git a/src/client/gateway/error.rs b/src/client/gateway/error.rs
index 30e87f0..121de71 100644
--- a/src/client/gateway/error.rs
+++ b/src/client/gateway/error.rs
@@ -12,6 +12,8 @@ pub enum Error {
     ExpectedHello,
     /// Expected a Ready or an InvalidateSession
     InvalidHandshake,
+    /// When a session Id was expected (for resuming), but was not present.
+    NoSessionId,
     /// Failed to reconnect after a number of attempts.
     ReconnectFailure,
 }
@@ -22,15 +24,10 @@ impl Display for Error {
             Error::Closed(s, ref v) => {
                 f.write_str(&format!("Connection closed {:?}: {:?}", s, v))
             },
-            Error::ExpectedHello => {
-                f.write_str("Expected Hello during handshake")
-            },
-            Error::InvalidHandshake => {
-                f.write_str("Expected Ready or InvalidateSession")
-            },
-            Error::ReconnectFailure => {
-                f.write_str("Failed to Reconnect")
-            },
+            Error::ExpectedHello => f.write_str("Expected Hello during handshake"),
+            Error::InvalidHandshake => f.write_str("Expected Ready or InvalidateSession"),
+            Error::NoSessionId => f.write_str("No Session Id present"),
+            Error::ReconnectFailure => f.write_str("Failed to Reconnect"),
         }
     }
 }
diff --git a/src/client/gateway/shard.rs b/src/client/gateway/shard.rs
index 3dfe0ee..283c7cd 100644
--- a/src/client/gateway/shard.rs
+++ b/src/client/gateway/shard.rs
@@ -302,12 +302,10 @@ impl Shard {
                           s,
                           self.seq);
 
-                    return if let Some(session_id) = self.session_id.clone() {
-                        self.resume(session_id, receiver)
-                            .map(|(ev, rec)| Some((ev, Some(rec))))
+                    return if self.session_id.is_some() {
+                        self.resume(receiver).map(|(ev, rec)| Some((ev, Some(rec))))
                     } else {
-                        self.reconnect(receiver)
-                            .map(|(ev, rec)| Some((ev, Some(rec))))
+                        self.reconnect(receiver).map(|(ev, rec)| Some((ev, Some(rec))))
                     };
                 }
 
@@ -331,12 +329,10 @@ impl Shard {
                     let _ = self.keepalive_channel.send(status);
                 }
 
-                if let Some(session_id) = self.session_id.clone() {
-                    self.resume(session_id, receiver)
-                        .map(|(ev, rec)| Some((ev, Some(rec))))
+                if self.session_id.is_some() {
+                    self.resume(receiver).map(|(ev, rec)| Some((ev, Some(rec))))
                 } else {
-                    self.reconnect(receiver)
-                        .map(|(ev, rec)| Some((ev, Some(rec))))
+                    self.reconnect(receiver).map(|(ev, rec)| Some((ev, Some(rec))))
                 }
             },
             Ok(GatewayEvent::InvalidateSession) => {
@@ -393,8 +389,8 @@ impl Shard {
                 if resume {
                     info!("Attempting to resume");
 
-                    if let Some(session_id) = self.session_id.clone() {
-                        match self.resume(session_id, receiver) {
+                    if self.session_id.is_some() {
+                        match self.resume(receiver) {
                             Ok((ev, rec)) => {
                                 info!("Resumed");
 
@@ -427,10 +423,10 @@ impl Shard {
                 // - InvalidateSession.
                 //
                 // Otherwise, fallback to reconnecting.
-                if let Some(session_id) = self.session_id.clone() {
+                if self.session_id.is_some() {
                     info!("Attempting to resume");
 
-                    match self.resume(session_id, &mut receiver) {
+                    match self.resume(&mut receiver) {
                         Ok((ev, rec)) => {
                             info!("Resumed");
 
@@ -596,9 +592,15 @@ impl Shard {
         Err(Error::Gateway(GatewayError::ReconnectFailure))
     }
 
-    fn resume(&mut self, session_id: String, receiver: &mut Receiver<WebSocketStream>)
+    #[doc(hidden)]
+    pub fn resume(&mut self, receiver: &mut Receiver<WebSocketStream>)
         -> Result<(Event, Receiver<WebSocketStream>)> {
-        receiver.get_mut().get_mut().shutdown(Shutdown::Both)?;
+        let session_id = match self.session_id.clone() {
+            Some(session_id) => session_id,
+            None => return Err(Error::Gateway(GatewayError::NoSessionId)),
+        };
+
+        let _ = receiver.shutdown_all();
         let url = prep::build_gateway_url(&self.ws_url)?;
 
         let response = WsClient::connect(url)?.send()?;
diff --git a/src/client/mod.rs b/src/client/mod.rs
index 073ffb5..5ef770a 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -1308,12 +1308,20 @@ fn handle_shard(info: &mut MonitorInfo) {
             Err(Error::WebSocket(WebSocketError::NoDataAvailable)) => {
                 debug!("Attempting to shutdown receiver/sender");
 
-                match info.receiver.shutdown_all() {
-                    Ok(_) => debug!("Successfully shutdown receiver/sender"),
-                    Err(why) => warn!("Err shutting down receiver/sender: {:?}", why),
-                }
+                match info.shard.lock().unwrap().resume(&mut info.receiver) {
+                    Ok((_, receiver)) => {
+                        debug!("Successfully resumed shard");
+
+                        info.receiver = receiver;
 
-                return;
+                        continue;
+                    },
+                    Err(why) => {
+                        warn!("Err resuming shard: {:?}", why);
+
+                        return;
+                    },
+                }
             },
             other => other,
         };