fix windows wakelock, add set_display (#6623)

Signed-off-by: 21pages <pages21@163.com>

Cargo.lock

@@ -3333,7 +3333,7 @@ dependencies = [
 [[package]]
 name = "keepawake"
 version = "0.4.3"
-source = "git+https://github.com/rustdesk-org/keepawake-rs#ac395ef826b32a077bc5d2fe108cf71fde8fe2e6"
+source = "git+https://github.com/rustdesk-org/keepawake-rs#ad94454a75cf1ff9e95e217dee9dd6a378bf625e"
 dependencies = [
  "anyhow",
  "apple-sys",

src/platform/macos.rs

@@ -733,4 +733,11 @@ impl WakeLock {
                 .ok(),
         )
     }
+
+    pub fn set_display(&mut self, display: bool) -> ResultType<()> {
+        self.0
+            .as_mut()
+            .map(|h| h.set_display(display))
+            .ok_or(anyhow!("no AwakeHandle"))?
+    }
 }

src/platform/mod.rs

@@ -98,7 +98,7 @@ impl WakeLock {
     }
 }
 
-pub fn get_wake_lock(_display: bool) -> WakeLock {
+pub fn get_wakelock(_display: bool) -> WakeLock {
     hbb_common::log::info!("new wakelock, require display on: {_display}");
     #[cfg(target_os = "android")]
     return crate::platform::WakeLock::new("server");

src/platform/windows.rs

@@ -2109,7 +2109,7 @@ pub fn is_process_consent_running() -> ResultType<bool> {
         .output()?;
     Ok(output.status.success() && !output.stdout.is_empty())
 }
-pub struct WakeLock;
+pub struct WakeLock(u32);
 // Failed to compile keepawake-rs on i686
 impl WakeLock {
     pub fn new(display: bool, idle: bool, sleep: bool) -> Self {
@@ -2124,7 +2124,20 @@ impl WakeLock {
             flag |= ES_AWAYMODE_REQUIRED;
         }
         unsafe { SetThreadExecutionState(flag) };
-        WakeLock {}
+        WakeLock(flag)
+    }
+
+    pub fn set_display(&mut self, display: bool) -> ResultType<()> {
+        let flag = if display {
+            self.0 | ES_DISPLAY_REQUIRED
+        } else {
+            self.0 & !ES_DISPLAY_REQUIRED
+        };
+        if flag != self.0 {
+            unsafe { SetThreadExecutionState(flag) };
+            self.0 = flag;
+        }
+        Ok(())
     }
 }
 

src/server/connection.rs

@@ -70,7 +70,7 @@ lazy_static::lazy_static! {
     static ref ALIVE_CONNS: Arc::<Mutex<Vec<i32>>> = Default::default();
     static ref AUTHED_CONNS: Arc::<Mutex<Vec<(i32, AuthConnType)>>> = Default::default();
     static ref SWITCH_SIDES_UUID: Arc::<Mutex<HashMap<String, (Instant, uuid::Uuid)>>> = Default::default();
-    static ref WAKE_LOCK: Arc::<Mutex<Option<(crate::platform::WakeLock, bool)>>> = Default::default();
+    static ref WAKELOCK_SENDER: Arc::<Mutex<std::sync::mpsc::Sender<(usize, usize)>>> = Arc::new(Mutex::new(start_wakelock_thread()));
 }
 
 #[cfg(any(target_os = "windows", target_os = "linux"))]
@@ -3176,6 +3176,51 @@ impl FileRemoveLogControl {
     }
 }
 
+fn start_wakelock_thread() -> std::sync::mpsc::Sender<(usize, usize)> {
+    use crate::platform::{get_wakelock, WakeLock};
+    let (tx, rx) = std::sync::mpsc::channel::<(usize, usize)>();
+    std::thread::spawn(move || {
+        let mut wakelock: Option<WakeLock> = None;
+        let mut last_display = false;
+        loop {
+            match rx.recv() {
+                Ok((conn_count, remote_count)) => {
+                    if conn_count == 0 {
+                        wakelock = None;
+                        log::info!("drop wakelock");
+                    } else {
+                        let mut display = remote_count > 0;
+                        if let Some(w) = wakelock.as_mut() {
+                            if display != last_display {
+                                #[cfg(any(target_os = "windows", target_os = "macos"))]
+                                {
+                                    log::info!("set wakelock display to {display}");
+                                    if let Err(e) = w.set_display(display) {
+                                        log::error!(
+                                            "failed to set wakelock display to {display}: {e:?}"
+                                        );
+                                    }
+                                }
+                            }
+                        } else {
+                            if cfg!(target_os = "linux") {
+                                display = true;
+                            }
+                            wakelock = Some(get_wakelock(display));
+                        }
+                        last_display = display;
+                    }
+                }
+                Err(e) => {
+                    log::error!("wakelock receive error: {e:?}");
+                    break;
+                }
+            }
+        }
+    });
+    tx
+}
+
 #[cfg(windows)]
 pub struct PortableState {
     pub last_uac: bool,
@@ -3252,7 +3297,6 @@ impl LinuxHeadlessHandle {
 extern "C" fn connection_shutdown_hook() {
     // https://stackoverflow.com/questions/35980148/why-does-an-atexit-handler-panic-when-it-accesses-stdout
     // Please make sure there is no print in the call stack
-    *WAKE_LOCK.lock().unwrap() = None;
     #[cfg(any(target_os = "windows", target_os = "linux"))]
     {
         *WALLPAPER_REMOVER.lock().unwrap() = None;
@@ -3296,27 +3340,17 @@ mod raii {
         }
 
         fn check_wake_lock() {
-            let mut wake_lock = WAKE_LOCK.lock().unwrap();
+            let conn_count = AUTHED_CONNS.lock().unwrap().len();
             let remote_count = AUTHED_CONNS
                 .lock()
                 .unwrap()
                 .iter()
                 .filter(|c| c.1 == AuthConnType::Remote)
                 .count();
-            let display = remote_count > 0;
-            if let Some((_, last_display)) = *wake_lock {
-                if last_display != display {
-                    *wake_lock = None;
-                }
-            }
-            let empty = AUTHED_CONNS.lock().unwrap().is_empty();
-            if empty {
-                *wake_lock = None;
-            } else {
-                if wake_lock.is_none() {
-                    *wake_lock = Some((crate::platform::get_wake_lock(display), display));
-                }
-            }
+            allow_err!(WAKELOCK_SENDER
+                .lock()
+                .unwrap()
+                .send((conn_count, remote_count)));
         }
     }