refactor

src/hbbr/main.rs

@@ -1,12 +1,10 @@
 use clap::App;
 mod relay_server;
-use hbb_common::{env_logger::*, tokio, ResultType};
-use relay_server::start;
+use hbb_common::{env_logger::*, ResultType};
+use relay_server::*;
+use std::sync::{Arc, Mutex};
 
-const DEFAULT_PORT: &'static str = "21117";
-
-#[tokio::main]
-async fn main() -> ResultType<()> {
+fn main() -> ResultType<()> {
     init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "info"));
     let args = format!(
         "-p, --port=[NUMBER(default={})] 'Sets the listening port'",
@@ -18,6 +16,7 @@ async fn main() -> ResultType<()> {
         .about("RustDesk Relay Server")
         .args_from_usage(&args)
         .get_matches();
-    start(matches.value_of("port").unwrap_or(DEFAULT_PORT)).await?;
+    let stop: Arc<Mutex<bool>> = Default::default();
+    start(matches.value_of("port").unwrap_or(DEFAULT_PORT), stop)?;
     Ok(())
 }

src/hbbr/relay_server.rs

@@ -4,7 +4,11 @@ use hbb_common::{
     rendezvous_proto::*,
     sleep,
     tcp::{new_listener, FramedStream},
-    tokio, ResultType,
+    tokio::{
+        self,
+        time::{interval, Duration},
+    },
+    ResultType,
 };
 use std::{
     collections::HashMap,
@@ -16,9 +20,13 @@ lazy_static::lazy_static! {
     static ref PEERS: Arc<Mutex<HashMap<String, FramedStream>>> = Arc::new(Mutex::new(HashMap::new()));
 }
 
-pub async fn start(port: &str) -> ResultType<()> {
+pub const DEFAULT_PORT: &'static str = "21117";
+
+#[tokio::main(basic_scheduler)]
+pub async fn start(port: &str, stop: Arc<Mutex<bool>>) -> ResultType<()> {
     let addr = format!("0.0.0.0:{}", port);
     log::info!("Listening on {}", addr);
+    let mut timer = interval(Duration::from_millis(300));
     let mut listener = new_listener(addr, false).await?;
     loop {
         tokio::select! {
@@ -27,8 +35,15 @@ pub async fn start(port: &str) -> ResultType<()> {
                     make_pair(FramedStream::from(stream), addr).await.ok();
                 });
             }
+            _ = timer.tick() => {
+                if *stop.lock().unwrap() {
+                    log::info!("Stopped");
+                    break;
+                }
+            }
         }
     }
+    Ok(())
 }
 
 async fn make_pair(stream: FramedStream, addr: SocketAddr) -> ResultType<()> {