fix blocking mesasges

server/src/ws.rs

@@ -7,11 +7,13 @@ use uuid::Uuid;
 
 use cookie::Cookie;
 
-use tungstenite::server::accept_hdr;
 use tungstenite::Message::Binary;
 use tungstenite::handshake::server::{Request, ErrorResponse};
+use tungstenite::handshake::HandshakeRole;
 use tungstenite::http::StatusCode;
 use tungstenite::protocol::WebSocket;
+use tungstenite::util::NonBlockingResult;
+use tungstenite::{accept_hdr};
 
 use crossbeam_channel::{unbounded, Receiver};
 
@@ -136,9 +138,11 @@ pub fn start(pool: PgPool, psr: Receiver<Message>) {
         let ws_pool = pool.clone();
         let ws_psr = psr.clone();
         spawn(move || {
-
             let (acc_s, acc_r) = unbounded();
 
+            let nb_stream = stream.unwrap();
+            nb_stream.set_nonblocking(true).unwrap();
+
             // search through the ws request for the auth cookie
             let cb = |req: &Request| {
                 let err = || ErrorResponse {
@@ -163,7 +167,7 @@ pub fn start(pool: PgPool, psr: Receiver<Message>) {
                 Ok(None)
             };
 
-            let mut websocket = accept_hdr(stream.unwrap(), cb).unwrap();
+            let mut websocket = accept_hdr(nb_stream, cb).unwrap();
 
             // get a copy of the account
             let account = match acc_r.recv().unwrap() {
@@ -206,8 +210,9 @@ pub fn start(pool: PgPool, psr: Receiver<Message>) {
             };
 
             loop {
-                match websocket.read_message() {
+                match websocket.read_message().no_block() {
                     Ok(msg) => {
+                        if let Some(msg) = msg {
                             match msg {
                                 Binary(data) => {
                                     let begin = Instant::now();
@@ -221,7 +226,7 @@ pub fn start(pool: PgPool, psr: Receiver<Message>) {
 
                                             if let Err(e) = websocket.write_message(Binary(response)) {
                                                 // connection closed
-                                            debug!("{:?}", e);
+                                                warn!("{:?}", e);
                                                 return;
                                             };
 
@@ -234,7 +239,7 @@ pub fn start(pool: PgPool, psr: Receiver<Message>) {
 
                                             if let Err(e) =  websocket.write_message(Binary(response)) {
                                                 // connection closed
-                                            debug!("{:?}", e);
+                                                warn!("{:?}", e);
                                                 return;
                                             };
                                         }
@@ -242,18 +247,19 @@ pub fn start(pool: PgPool, psr: Receiver<Message>) {
                                 },
                                 _ => (),
                             }
-                    },
-                    // connection is closed
-                    Err(e) => {
-                        debug!("{:?}", e);
-                        return;
-                    }
                         };
 
                         match ws_psr.try_recv() {
                             Ok(n) => handle_message(&subs, n, &mut websocket),
                             Err(_) => (),
                         };
+                    },
+                    // connection is closed
+                    Err(e) => {
+                        warn!("{:?}", e);
+                        return;
+                    }
+                };
             }
         });
     }