Skip to main content

handshake_core/
keep_awake.rs

1//! App-wide, temporary sleep prevention, independent of folder and agent lifetimes.
2use std::process::{Child, Command, Stdio};
3use std::sync::Mutex;
4
5use thiserror::Error;
6
7static CONTROLLER: Mutex<Controller> = Mutex::new(Controller { child: None });
8
9#[derive(Debug, Error)]
10pub enum KeepAwakeError {
11    #[error("Keep Awake is only available on macOS")]
12    Unsupported,
13    #[error("Keep Awake could not access its process: {0}")]
14    Process(#[from] std::io::Error),
15    #[error("Keep Awake state is unavailable")]
16    Poisoned,
17}
18
19/// Reads the live state, optionally enabling or disabling it for the whole application process.
20/// Call with `Some(false)` during normal app shutdown. The child also watches this process so
21/// a crash cannot leave sleep prevention running. State is never persisted.
22///
23/// # Errors
24/// Returns an error if process supervision fails or the platform does not support Keep Awake.
25pub fn keep_awake(enabled: Option<bool>) -> Result<bool, KeepAwakeError> {
26    let mut controller = CONTROLLER.lock().map_err(|_| KeepAwakeError::Poisoned)?;
27    if let Some(enabled) = enabled {
28        controller.set(enabled, || {
29            if !cfg!(target_os = "macos") {
30                return Err(KeepAwakeError::Unsupported);
31            }
32            Ok(Command::new("/usr/bin/caffeinate")
33                .args(["-disu", "-w", &std::process::id().to_string()])
34                .stdin(Stdio::null())
35                .stdout(Stdio::null())
36                .stderr(Stdio::null())
37                .spawn()?)
38        })?;
39    }
40    Ok(controller.active()?)
41}
42
43#[derive(Default)]
44struct Controller {
45    child: Option<Child>,
46}
47
48impl Controller {
49    fn active(&mut self) -> std::io::Result<bool> {
50        if let Some(child) = &mut self.child
51            && child.try_wait()?.is_some()
52        {
53            self.child = None;
54        }
55        Ok(self.child.is_some())
56    }
57
58    fn set(
59        &mut self,
60        enabled: bool,
61        launch: impl FnOnce() -> Result<Child, KeepAwakeError>,
62    ) -> Result<(), KeepAwakeError> {
63        let active = self.active()?;
64        if enabled && !active {
65            self.child = Some(launch()?);
66        } else if !enabled {
67            self.stop()?;
68        }
69        Ok(())
70    }
71
72    fn stop(&mut self) -> std::io::Result<()> {
73        if let Some(child) = &mut self.child {
74            child.kill()?;
75            child.wait()?;
76        }
77        self.child = None;
78        Ok(())
79    }
80}
81
82impl Drop for Controller {
83    fn drop(&mut self) {
84        if let Err(error) = self.stop() {
85            tracing::warn!(%error, "Keep Awake process cleanup failed");
86        }
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    fn sleeper() -> Result<Child, KeepAwakeError> {
95        Ok(Command::new("/bin/sleep").arg("30").spawn()?)
96    }
97
98    #[test]
99    fn repeated_enable_reuses_the_child_and_disable_reaps_it() {
100        let mut controller = Controller::default();
101        assert!(!controller.active().unwrap());
102        controller.set(true, sleeper).unwrap();
103        let pid = controller.child.as_ref().unwrap().id();
104        controller.set(true, || panic!("already running")).unwrap();
105        assert_eq!(controller.child.as_ref().unwrap().id(), pid);
106        controller
107            .set(false, || panic!("disable must not launch"))
108            .unwrap();
109        assert!(!controller.active().unwrap());
110        assert!(
111            !Command::new("/bin/kill")
112                .args(["-0", &pid.to_string()])
113                .stderr(Stdio::null())
114                .status()
115                .unwrap()
116                .success()
117        );
118        controller.set(false, sleeper).unwrap();
119    }
120
121    #[test]
122    fn launch_failure_stays_off_and_can_be_retried() {
123        let mut controller = Controller::default();
124        assert!(
125            controller
126                .set(true, || Err(std::io::Error::from(
127                    std::io::ErrorKind::NotFound
128                )
129                .into()))
130                .is_err()
131        );
132        assert!(!controller.active().unwrap());
133        controller.set(true, sleeper).unwrap();
134        assert!(controller.active().unwrap());
135    }
136
137    #[test]
138    fn external_exit_is_reaped_and_can_be_restarted() {
139        let mut controller = Controller::default();
140        controller.set(true, sleeper).unwrap();
141        controller.child.as_mut().unwrap().kill().unwrap();
142        controller.child.as_mut().unwrap().wait().unwrap();
143        assert!(!controller.active().unwrap());
144        controller.set(true, sleeper).unwrap();
145        assert!(controller.active().unwrap());
146    }
147}