handshake_core/
keep_awake.rs1use 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
19pub 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}