Implement websocket heartbeats (#48)

websocket/src/client.rs

@@ -74,6 +74,9 @@ impl ChatClient {
         ctx.run_later(Duration::new(1, 0), |act, ctx| {
             act.0.ping("");
             act.hb(ctx);
+
+            // client should also check for a timeout here, similar to the
+            // server code
         });
     }
 }

websocket/src/main.rs

@@ -8,22 +8,38 @@ extern crate actix;
 extern crate actix_web;
 extern crate env_logger;
 
+use std::time::{Instant, Duration};
+
 use actix::prelude::*;
 use actix_web::{
     fs, http, middleware, server, ws, App, Error, HttpRequest, HttpResponse,
 };
 
+/// How often heartbeat pings are sent
+const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
+/// How long before lack of client response causes a timeout
+const CLIENT_TIMEOUT: Duration = Duration::from_secs(10);
+
 /// do websocket handshake and start `MyWebSocket` actor
 fn ws_index(r: &HttpRequest) -> Result<HttpResponse, Error> {
-    ws::start(r, MyWebSocket)
+    ws::start(r, MyWebSocket::new())
 }
 
 /// websocket connection is long running connection, it easier
 /// to handle with an actor
-struct MyWebSocket;
+struct MyWebSocket {
+    /// Client must send ping at least once per 10 seconds (CLIENT_TIMEOUT),
+    /// otherwise we drop connection.
+    hb: Instant,
+}
 
 impl Actor for MyWebSocket {
     type Context = ws::WebsocketContext<Self>;
+
+    /// Method is called on actor start. We start the heartbeat process here.
+    fn started(&mut self, ctx: &mut Self::Context) {
+        self.hb(ctx);
+    }
 }
 
 /// Handler for `ws::Message`
@@ -32,7 +48,10 @@ impl StreamHandler<ws::Message, ws::ProtocolError> for MyWebSocket {
         // process websocket messages
         println!("WS: {:?}", msg);
         match msg {
-            ws::Message::Ping(msg) => ctx.pong(&msg),
+            ws::Message::Ping(msg) => {
+                self.hb = Instant::now();
+                ctx.pong(&msg);
+            }
             ws::Message::Text(text) => ctx.text(text),
             ws::Message::Binary(bin) => ctx.binary(bin),
             ws::Message::Close(_) => {
@@ -43,6 +62,33 @@ impl StreamHandler<ws::Message, ws::ProtocolError> for MyWebSocket {
     }
 }
 
+impl MyWebSocket {
+    fn new() -> Self {
+        Self { hb: Instant::now() }
+    }
+
+    /// helper method that sends ping to client every second.
+    ///
+    /// also this method checks heartbeats from client
+    fn hb(&self, ctx: &mut <Self as Actor>::Context) {
+        ctx.run_interval(HEARTBEAT_INTERVAL, |act, ctx| {
+            // check client heartbeats
+            if Instant::now().duration_since(act.hb) > CLIENT_TIMEOUT {
+                // heartbeat timed out
+                println!("Websocket Client heartbeat failed, disconnecting!");
+
+                // stop actor
+                ctx.stop();
+
+                // don't try to send a ping
+                return;
+            }
+
+            ctx.ping("");
+        });
+    }
+}
+
 fn main() {
     ::std::env::set_var("RUST_LOG", "actix_web=info");
     env_logger::init();