From 252d3cb797311e73b2c552f3f49851fad8d5b930 Mon Sep 17 00:00:00 2001
From: rustdesk <info@rustdesk.com>
Date: Thu, 23 Dec 2021 11:11:48 +0800
Subject: [PATCH] refactor

---
 src/server/connection.rs    | 21 ++++++++++++---------
 src/server/input_service.rs |  2 --
 src/server/service.rs       |  9 ++++-----
 src/server/video_service.rs |  7 +++----
 4 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/src/server/connection.rs b/src/server/connection.rs
index 0782338e6..0cebac601 100644
--- a/src/server/connection.rs
+++ b/src/server/connection.rs
@@ -62,15 +62,18 @@ impl Subscriber for ConnInner {
 
     #[inline]
     fn send(&mut self, msg: Arc<Message>) {
-        self.tx.as_mut().map(|tx| {
-            allow_err!(tx.send((Instant::now(), msg)));
-        });
-    }
-
-    fn send_video_frame(&mut self, tm: std::time::Instant, msg: Arc<Message>) {
-        self.tx_video.as_mut().map(|tx| {
-            allow_err!(tx.send((tm.into(), msg)));
-        });
+        match &msg.union {
+            Some(message::Union::video_frame(_)) => {
+                self.tx_video.as_mut().map(|tx| {
+                    allow_err!(tx.send((Instant::now(), msg)));
+                });
+            }
+            _ => {
+                self.tx.as_mut().map(|tx| {
+                    allow_err!(tx.send((Instant::now(), msg)));
+                });
+            }
+        }
     }
 }
 
diff --git a/src/server/input_service.rs b/src/server/input_service.rs
index aee13304b..056267c5b 100644
--- a/src/server/input_service.rs
+++ b/src/server/input_service.rs
@@ -79,8 +79,6 @@ impl Subscriber for MouseCursorSub {
             self.inner.send(msg);
         }
     }
-
-    fn send_video_frame(&mut self, _: std::time::Instant, _: Arc<Message>) {}
 }
 
 pub const NAME_CURSOR: &'static str = "mouse_cursor";
diff --git a/src/server/service.rs b/src/server/service.rs
index 7f280130f..9c146b047 100644
--- a/src/server/service.rs
+++ b/src/server/service.rs
@@ -16,7 +16,6 @@ pub trait Service: Send + Sync {
 pub trait Subscriber: Default + Send + Sync + 'static {
     fn id(&self) -> i32;
     fn send(&mut self, msg: Arc<Message>);
-    fn send_video_frame(&mut self, tm: time::Instant, msg: Arc<Message>);
 }
 
 #[derive(Default)]
@@ -145,15 +144,15 @@ impl<T: Subscriber + From<ConnInner>> ServiceTmpl<T> {
         }
     }
 
-    pub fn send_video_frame(&self, tm: time::Instant, msg: Message) -> HashSet<i32> {
-        self.send_video_frame_shared(tm, Arc::new(msg))
+    pub fn send_video_frame(&self, msg: Message) -> HashSet<i32> {
+        self.send_video_frame_shared(Arc::new(msg))
     }
 
-    pub fn send_video_frame_shared(&self, tm: time::Instant, msg: Arc<Message>) -> HashSet<i32> {
+    pub fn send_video_frame_shared(&self, msg: Arc<Message>) -> HashSet<i32> {
         let mut conn_ids = HashSet::new();
         let mut lock = self.0.write().unwrap();
         for s in lock.subscribes.values_mut() {
-            s.send_video_frame(tm, msg.clone());
+            s.send(msg.clone());
             conn_ids.insert(s.id());
         }
         conn_ids
diff --git a/src/server/video_service.rs b/src/server/video_service.rs
index f087cc30b..9ec5f9849 100644
--- a/src/server/video_service.rs
+++ b/src/server/video_service.rs
@@ -108,7 +108,7 @@ impl VideoFrameController {
                         fetched_conn_ids.insert(id);
 
                         // break if all connections have received current frame
-                        if fetched_conn_ids.is_superset(&send_conn_ids) {
+                        if fetched_conn_ids.len() >= send_conn_ids.len() {
                             break;
                         }
                     }
@@ -257,7 +257,7 @@ fn run(sp: GenericService) -> ResultType<()> {
             Ok(frame) => {
                 let time = now - start;
                 let ms = (time.as_secs() * 1000 + time.subsec_millis() as u64) as i64;
-                let send_conn_ids = handle_one_frame(&sp, now, &frame, ms, &mut crc, &mut vpx)?;
+                let send_conn_ids = handle_one_frame(&sp, &frame, ms, &mut crc, &mut vpx)?;
                 frame_controller.set_send(now, send_conn_ids);
                 #[cfg(windows)]
                 {
@@ -327,7 +327,6 @@ fn create_frame(frame: &EncodeFrame) -> VP9 {
 #[inline]
 fn handle_one_frame(
     sp: &GenericService,
-    now: Instant,
     frame: &[u8],
     ms: i64,
     crc: &mut (u32, u32),
@@ -365,7 +364,7 @@ fn handle_one_frame(
 
         // to-do: flush periodically, e.g. 1 second
         if frames.len() > 0 {
-            send_conn_ids = sp.send_video_frame(now, create_msg(frames));
+            send_conn_ids = sp.send_video_frame(create_msg(frames));
         }
     }
     Ok(send_conn_ids)