From 58d3e5f8c5400bc5417f827697db59665500b4c9 Mon Sep 17 00:00:00 2001
From: bott <mhpoin@gmail.com>
Date: Tue, 28 Aug 2018 09:53:51 +0200
Subject: [PATCH] fix missing heartbeat

---
 websocket-tcp-chat/src/main.rs | 39 +++++++++++++++++++++++++++++++---
 1 file changed, 36 insertions(+), 3 deletions(-)

diff --git a/websocket-tcp-chat/src/main.rs b/websocket-tcp-chat/src/main.rs
index a4874536..0c8f1c47 100644
--- a/websocket-tcp-chat/src/main.rs
+++ b/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");
@@ -205,15 +238,15 @@ fn main() {
         };
 
         App::with_state(state)
-        // redirect to websocket.html
+            // redirect to websocket.html
             .resource("/", |r| r.method(http::Method::GET).f(|_| {
                 HttpResponse::Found()
                     .header("LOCATION", "/static/websocket.html")
                     .finish()
             }))
-        // websocket
+            // websocket
             .resource("/ws/", |r| r.route().f(chat_route))
-        // static resources
+            // static resources
             .handler("/static/", fs::StaticFiles::new("static/").unwrap())
     }).bind("127.0.0.1:8080")
         .unwrap()