1 files changed, 27 insertions, 1 deletions

diff --git a/src/client/bridge/gateway/shard_runner.rs b/src/client/bridge/gateway/shard_runner.rs
index cff8d7e..7c6cf83 100644
--- a/src/client/bridge/gateway/shard_runner.rs
+++ b/src/client/bridge/gateway/shard_runner.rs
@@ -141,8 +141,16 @@ impl<H: EventHandler + Send + Sync + 'static> ShardRunner<H> {
                 }
             }
 
+            let pre = self.shard.stage();
+
             let (event, action, successful) = self.recv_event();
 
+            let post = self.shard.stage();
+
+            if post != pre {
+                self.update_manager();
+            }
+
             match action {
                 Some(ShardAction::Reconnect(ReconnectType::Reidentify)) => {
                     return self.request_restart()
@@ -249,7 +257,8 @@ impl<H: EventHandler + Send + Sync + 'static> ShardRunner<H> {
 
                     true
                 },
-                ShardManagerMessage::ShutdownInitiated => {
+                ShardManagerMessage::ShardUpdate { .. }
+                    | ShardManagerMessage::ShutdownInitiated => {
                     // nb: not sent here
 
                     true
@@ -411,6 +420,13 @@ impl<H: EventHandler + Send + Sync + 'static> ShardRunner<H> {
             },
         };
 
+        match event {
+            Ok(GatewayEvent::HeartbeatAck) => {
+                self.update_manager();
+            },
+            _ => {},
+        }
+
         let event = match event {
             Ok(GatewayEvent::Dispatch(_, event)) => Some(event),
             _ => None,
@@ -420,6 +436,8 @@ impl<H: EventHandler + Send + Sync + 'static> ShardRunner<H> {
     }
 
     fn request_restart(&self) -> Result<()> {
+        self.update_manager();
+
         debug!(
             "[ShardRunner {:?}] Requesting restart",
             self.shard.shard_info(),
@@ -430,4 +448,12 @@ impl<H: EventHandler + Send + Sync + 'static> ShardRunner<H> {
 
         Ok(())
     }
+
+    fn update_manager(&self) {
+        let _ = self.manager_tx.send(ShardManagerMessage::ShardUpdate {
+            id: ShardId(self.shard.shard_info()[0]),
+            latency: self.shard.latency(),
+            stage: self.shard.stage(),
+        });
+    }
 }