rightkit_child_process/
power.rs

1//! Keep the machine awake while work runs.
2//!
3//! [`PowerAssertion::acquire`] returns a guard. While it is alive the system does not go to
4//! idle sleep (and, when asked, the display stays on); dropping it releases the request.
5//! Windows uses `SetThreadExecutionState` on a small holder thread (the state belongs to the
6//! calling thread, so a thread that lives exactly as long as the guard owns it). macOS uses an
7//! IOKit power assertion (`IOPMAssertionCreateWithName`), which is process-wide. Other hosts
8//! return `ErrorKind::Unsupported`: nothing is held and nothing is faked.
9//!
10//! The request is released if the process dies, by either mechanism. It is advisory for idle
11//! sleep: closing a laptop lid or a low battery still sleeps the machine.
12
13use std::io;
14
15/// What to keep awake.
16#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
17pub enum PowerKind {
18    /// Prevent idle system sleep; the display may turn off.
19    #[default]
20    System,
21    /// Prevent idle system sleep and keep the display on.
22    SystemAndDisplay,
23}
24
25/// A live keep-awake request. Released on drop.
26#[derive(Debug)]
27pub struct PowerAssertion {
28    // Held only for its Drop, which releases the request.
29    #[allow(dead_code)]
30    inner: imp::Inner,
31}
32
33impl PowerAssertion {
34    /// Take a keep-awake request. `reason` is shown by `pmset -g assertions` on macOS and is
35    /// otherwise informational; it must not contain a NUL byte.
36    pub fn acquire(reason: &str, kind: PowerKind) -> io::Result<Self> {
37        if reason.bytes().any(|byte| byte == 0) {
38            return Err(io::Error::new(
39                io::ErrorKind::InvalidInput,
40                "reason contains a NUL byte",
41            ));
42        }
43        Ok(Self {
44            inner: imp::Inner::acquire(reason, kind)?,
45        })
46    }
47}
48
49#[cfg(windows)]
50mod imp {
51    use super::PowerKind;
52    use std::{io, sync::mpsc, thread};
53    use windows::Win32::System::Power::{
54        SetThreadExecutionState, ES_CONTINUOUS, ES_DISPLAY_REQUIRED, ES_SYSTEM_REQUIRED,
55    };
56
57    #[derive(Debug)]
58    pub(super) struct Inner {
59        release: Option<mpsc::Sender<()>>,
60        holder: Option<thread::JoinHandle<()>>,
61    }
62
63    impl Inner {
64        pub(super) fn acquire(_reason: &str, kind: PowerKind) -> io::Result<Self> {
65            let (release, wait) = mpsc::channel::<()>();
66            let (ready, started) = mpsc::channel::<io::Result<()>>();
67            let holder = thread::Builder::new()
68                .name("rightkit-power-assertion".into())
69                .spawn(move || {
70                    let mut flags = ES_CONTINUOUS | ES_SYSTEM_REQUIRED;
71                    if kind == PowerKind::SystemAndDisplay {
72                        flags |= ES_DISPLAY_REQUIRED;
73                    }
74                    // SAFETY: SetThreadExecutionState takes a flag value only. A zero return
75                    // means the request was refused.
76                    if unsafe { SetThreadExecutionState(flags) }.0 == 0 {
77                        let _ = ready.send(Err(io::Error::other(
78                            "SetThreadExecutionState refused the request",
79                        )));
80                        return;
81                    }
82                    let _ = ready.send(Ok(()));
83                    // Returns on release or when the guard (the sender) is gone.
84                    let _ = wait.recv();
85                    // SAFETY: as above; ES_CONTINUOUS alone clears this thread's request.
86                    unsafe { SetThreadExecutionState(ES_CONTINUOUS) };
87                })?;
88            match started.recv() {
89                Ok(Ok(())) => Ok(Self {
90                    release: Some(release),
91                    holder: Some(holder),
92                }),
93                Ok(Err(error)) => {
94                    let _ = holder.join();
95                    Err(error)
96                }
97                Err(_) => Err(io::Error::other("power assertion thread ended early")),
98            }
99        }
100    }
101
102    impl Drop for Inner {
103        fn drop(&mut self) {
104            drop(self.release.take());
105            if let Some(holder) = self.holder.take() {
106                let _ = holder.join();
107            }
108        }
109    }
110}
111
112#[cfg(target_os = "macos")]
113mod imp {
114    use super::PowerKind;
115    use std::{
116        ffi::{c_char, c_void, CString},
117        io,
118    };
119
120    type CfStringRef = *const c_void;
121    const K_CF_STRING_ENCODING_UTF8: u32 = 0x0800_0100;
122    const K_IOPM_ASSERTION_LEVEL_ON: u32 = 255;
123
124    #[link(name = "CoreFoundation", kind = "framework")]
125    extern "C" {
126        fn CFStringCreateWithCString(
127            alloc: *const c_void,
128            text: *const c_char,
129            encoding: u32,
130        ) -> CfStringRef;
131        fn CFRelease(object: *const c_void);
132    }
133    #[link(name = "IOKit", kind = "framework")]
134    extern "C" {
135        fn IOPMAssertionCreateWithName(
136            assertion_type: CfStringRef,
137            level: u32,
138            name: CfStringRef,
139            assertion_id: *mut u32,
140        ) -> i32;
141        fn IOPMAssertionRelease(assertion_id: u32) -> i32;
142    }
143
144    /// Owns one retained CFString.
145    struct CfString(CfStringRef);
146    impl CfString {
147        fn new(text: &str) -> io::Result<Self> {
148            let text = CString::new(text)
149                .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "NUL byte"))?;
150            // SAFETY: a valid NUL-terminated string and a null (default) allocator.
151            let string = unsafe {
152                CFStringCreateWithCString(
153                    std::ptr::null(),
154                    text.as_ptr(),
155                    K_CF_STRING_ENCODING_UTF8,
156                )
157            };
158            if string.is_null() {
159                return Err(io::Error::other("CFStringCreateWithCString failed"));
160            }
161            Ok(Self(string))
162        }
163    }
164    impl Drop for CfString {
165        fn drop(&mut self) {
166            // SAFETY: created by a Create function, released exactly once.
167            unsafe { CFRelease(self.0) };
168        }
169    }
170
171    #[derive(Debug)]
172    pub(super) struct Inner {
173        id: u32,
174    }
175
176    impl Inner {
177        pub(super) fn acquire(reason: &str, kind: PowerKind) -> io::Result<Self> {
178            let assertion_type = match kind {
179                PowerKind::System => "PreventUserIdleSystemSleep",
180                PowerKind::SystemAndDisplay => "PreventUserIdleDisplaySleep",
181            };
182            let assertion_type = CfString::new(assertion_type)?;
183            let name = CfString::new(reason)?;
184            let mut id = 0u32;
185            // SAFETY: both strings outlive the call (IOKit copies them); `id` is a valid out
186            // pointer.
187            let status = unsafe {
188                IOPMAssertionCreateWithName(
189                    assertion_type.0,
190                    K_IOPM_ASSERTION_LEVEL_ON,
191                    name.0,
192                    &mut id,
193                )
194            };
195            if status != 0 {
196                return Err(io::Error::other(format!(
197                    "IOPMAssertionCreateWithName failed: {status:#x}"
198                )));
199            }
200            Ok(Self { id })
201        }
202    }
203
204    impl Drop for Inner {
205        fn drop(&mut self) {
206            // SAFETY: `id` came from a successful create and is released once.
207            unsafe { IOPMAssertionRelease(self.id) };
208        }
209    }
210}
211
212#[cfg(not(any(windows, target_os = "macos")))]
213mod imp {
214    use super::PowerKind;
215    use std::io;
216
217    #[derive(Debug)]
218    pub(super) struct Inner;
219
220    impl Inner {
221        pub(super) fn acquire(_reason: &str, _kind: PowerKind) -> io::Result<Self> {
222            Err(io::Error::new(
223                io::ErrorKind::Unsupported,
224                "keep-awake is supported on Windows and macOS only",
225            ))
226        }
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233
234    #[test]
235    fn a_reason_with_a_nul_byte_is_rejected() {
236        let error = PowerAssertion::acquire("a\u{0}b", PowerKind::System).unwrap_err();
237        assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
238    }
239
240    #[cfg(any(windows, target_os = "macos"))]
241    #[test]
242    fn acquire_and_release_round_trip() {
243        let guard = PowerAssertion::acquire("rightkit power test", PowerKind::System).unwrap();
244        drop(guard);
245        let guard =
246            PowerAssertion::acquire("rightkit power test", PowerKind::SystemAndDisplay).unwrap();
247        drop(guard);
248    }
249
250    #[cfg(not(any(windows, target_os = "macos")))]
251    #[test]
252    fn other_hosts_report_unsupported() {
253        let error = PowerAssertion::acquire("x", PowerKind::System).unwrap_err();
254        assert_eq!(error.kind(), io::ErrorKind::Unsupported);
255    }
256}
rightkit_child_process/power.rs

rightkit_child_process/
power.rs
