fix missing heartbeat

websocket-tcp-chat/src/main.rs

@@ -21,6 +21,7 @@ use std::time::Instant;
 use actix::*;
 use actix_web::server::HttpServer;
 use actix_web::{fs, http, ws, App, Error, HttpRequest, HttpResponse};
+use std::time::Duration;
 
 mod codec;
 mod server;
@@ -68,6 +69,10 @@ impl Actor for WsChatSession {
         // before processing any other events.
         // HttpContext::state() is instance of WsChatSessionState, state is shared
         // across all routes within application
+
+        // we'll start heartbeat process on session start.
+        self.hb(ctx);
+
         let addr = ctx.address();
         ctx.state()
             .addr
@@ -183,6 +188,34 @@ impl StreamHandler<ws::Message, ws::ProtocolError> for WsChatSession {
     }
 }
 
+impl WsChatSession {
+    /// helper method that sends ping to client every second.
+    ///
+    /// also this method check heartbeats from client
+    fn hb(&self, ctx: &mut ws::WebsocketContext<Self, WsChatSessionState>) {
+        ctx.run_later(Duration::new(1, 0), |act, ctx| {
+            // check client heartbeats
+            if Instant::now().duration_since(act.hb) > Duration::new(10, 0) {
+                // heartbeat timed out
+                println!("Websocket Client heartbeat failed, disconnecting!");
+
+                // notify chat server
+                ctx.state()
+                    .addr
+                    .do_send(server::Disconnect { id: act.id });
+
+                // stop actor
+                ctx.stop();
+            }
+
+            ctx.ping("");
+
+            // if we can not send message to sink, sink is closed (disconnected)
+            act.hb(ctx);
+        });
+    }
+}
+
 fn main() {
     let _ = env_logger::init();
     let sys = actix::System::new("websocket-example");