From f21a180053f644269a27d6d31195d7387ef325da Mon Sep 17 00:00:00 2001
From: Jovi Hsu <jv.hsu@outlook.com>
Date: Wed, 17 Nov 2021 14:42:08 +0800
Subject: [PATCH] shutdown websocket if the server disconnected, level up the
 default log print

---
 backend/src/agent/agent.rs | 3 +++
 backend/src/agent/ws.rs    | 4 ++++
 backend/src/main.rs        | 2 +-
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/backend/src/agent/agent.rs b/backend/src/agent/agent.rs
index 7abcdee..2b34002 100644
--- a/backend/src/agent/agent.rs
+++ b/backend/src/agent/agent.rs
@@ -111,6 +111,9 @@ impl StreamHandler<Result<Bytes, io::Error>> for Agent {
             }
             Err(err) => {
                 error!("error: {:?}", err);
+                if self.ws_addr.is_some() {
+                    self.ws_addr.as_ref().unwrap().do_send(ws::WsMsg::Close);
+                }
                 ctx.address().do_send(AgentMsg::Shutdown);
             }
         }
diff --git a/backend/src/agent/ws.rs b/backend/src/agent/ws.rs
index f92e8cd..f979464 100644
--- a/backend/src/agent/ws.rs
+++ b/backend/src/agent/ws.rs
@@ -16,6 +16,7 @@ use super::agent::*;
 #[rtype(result = "()")]
 pub enum WsMsg {
     SendToClient(Bytes),
+    Close,
 }
 
 /// Define Websocket actor
@@ -41,6 +42,9 @@ impl Handler<WsMsg> for WsSession {
             WsMsg::SendToClient(data) => {
                 ctx.binary(data);
             }
+            WsMsg::Close => {
+                ctx.stop();
+            }
         };
     }
 }
diff --git a/backend/src/main.rs b/backend/src/main.rs
index cd024da..a2e95ee 100644
--- a/backend/src/main.rs
+++ b/backend/src/main.rs
@@ -44,7 +44,7 @@ impl Default for AppData {
 fn setup_logger() {
     let logger = femme::pretty::Logger::new();
     async_log::Logger::wrap(logger, || 12)
-        .start(log::LevelFilter::Trace)
+        .start(log::LevelFilter::Warn)
         .unwrap();
 }