Initial public release of Warp.

Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b

crates/prevent_sleep/Cargo.toml

@@ -0,0 +1,29 @@
+[package]
+name = "prevent_sleep"
+edition = "2024"
+authors.workspace = true
+publish.workspace = true
+license.workspace = true
+
+[dependencies]
+cfg-if.workspace = true
+futures.workspace = true
+log.workspace = true
+pin-project.workspace = true
+
+[target.'cfg(windows)'.dependencies]
+itertools.workspace = true
+parking_lot.workspace = true
+windows = { workspace = true, features = [
+    "Win32_System_Power",
+] }
+
+[target.'cfg(target_os = "macos")'.dependencies]
+objc2.workspace = true
+objc2-foundation = { workspace = true, features = [
+    "NSProcessInfo",
+    "NSString",
+] }
+
+[build-dependencies]
+cfg_aliases.workspace = true

crates/prevent_sleep/build.rs

@@ -0,0 +1,8 @@
+use cfg_aliases::cfg_aliases;
+
+fn main() {
+    cfg_aliases! {
+        macos: { target_os = "macos" },
+        noop: { not(any(macos, windows)) },
+    }
+}

crates/prevent_sleep/src/lib.rs

@@ -0,0 +1,47 @@
+#[cfg_attr(macos, path = "mac.rs")]
+#[cfg_attr(windows, path = "windows.rs")]
+#[cfg_attr(noop, path = "noop.rs")]
+mod imp;
+
+use std::{
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use pin_project::pin_project;
+
+pub use imp::Guard;
+
+/// Returns a guard that prevents the system from going to sleep while the guard is held.
+///
+/// Callers should provide a description of the reason for preventing sleep. Depending on
+/// platform, this may appear in logs, so write it as though it may be user-visible, e.g.:
+/// "Agent Mode request in-progress".
+pub fn prevent_sleep(reason: &'static str) -> Guard {
+    imp::prevent_sleep(reason)
+}
+
+/// A simple wrapper around a stream that optionally prevents the system from going to sleep
+/// while the stream is being polled.
+#[pin_project]
+pub struct Stream<S> {
+    #[pin]
+    inner: S,
+    guard: Option<Guard>,
+}
+
+impl<S> Stream<S> {
+    /// Wraps the provided stream, maintaining the provided sleep guard as long as the stream is
+    /// being polled.
+    pub fn wrap(inner: S, guard: Option<Guard>) -> Self {
+        Self { inner, guard }
+    }
+}
+
+impl<S: futures::stream::Stream> futures::stream::Stream for Stream<S> {
+    type Item = S::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        self.project().inner.poll_next(cx)
+    }
+}

crates/prevent_sleep/src/mac.rs

@@ -0,0 +1,42 @@
+use objc2::{rc::Retained, runtime::ProtocolObject};
+use objc2_foundation::{NSActivityOptions, NSObjectProtocol, NSProcessInfo, NSString};
+
+/// A guard object that prevents system sleep while it remains in scope.
+pub struct Guard {
+    process_info: Retained<NSProcessInfo>,
+    activity_token: Retained<ProtocolObject<dyn NSObjectProtocol>>,
+    reason: Retained<NSString>,
+}
+
+// Mark the guard as safe for being sent across threads.  We don't need to worry about thread safety
+// here because the underlying process info and activity marker can be shared across threads, and
+// we are sure that there aren't synchronization issues because we only interact with the activity
+// during creation and drop.
+unsafe impl Send for Guard {}
+unsafe impl Sync for Guard {}
+
+impl Drop for Guard {
+    fn drop(&mut self) {
+        unsafe {
+            self.process_info.endActivity(&self.activity_token);
+        }
+        log::info!("No longer preventing sleep with reason: {}", self.reason);
+    }
+}
+
+/// Returns a guard that prevents system sleep while it remains in scope.
+pub fn prevent_sleep(reason: &'static str) -> Guard {
+    let reason = NSString::from_str(reason);
+
+    let process_info = NSProcessInfo::processInfo();
+    let activity_token =
+        process_info.beginActivityWithOptions_reason(NSActivityOptions::UserInitiated, &reason);
+
+    log::info!("Preventing sleep with reason: {reason}");
+
+    Guard {
+        process_info,
+        activity_token,
+        reason,
+    }
+}

crates/prevent_sleep/src/noop.rs

@@ -0,0 +1,5 @@
+pub struct Guard;
+
+pub fn prevent_sleep(_reason: &'static str) -> Guard {
+    Guard
+}

crates/prevent_sleep/src/windows.rs

@@ -0,0 +1,165 @@
+use std::{
+    sync::{LazyLock, Once, mpsc},
+    thread::JoinHandle,
+};
+
+use itertools::Itertools as _;
+use parking_lot::Mutex;
+use windows::Win32::System::Power::{self, SetThreadExecutionState};
+
+/// The global backing state for the sleep prevention logic.
+static STATE: LazyLock<State> = LazyLock::new(State::new);
+
+/// Ensures that we only log message send failures once.
+static SEND_FAILURE: Once = Once::new();
+
+enum StateUpdate {
+    AddTask { task_id: u64, reason: &'static str },
+    RemoveTask { task_id: u64 },
+}
+
+/// The underlying state for the sleep prevention logic.
+struct State {
+    inner: Mutex<StateInner>,
+}
+
+impl State {
+    /// Constructs a new state object.
+    fn new() -> Self {
+        let (update_tx, update_rx) = mpsc::channel::<StateUpdate>();
+
+        let join_handle = std::thread::Builder::new()
+            .name("prevent_sleep".to_string())
+            .spawn(move || {
+                Self::thread_main(update_rx);
+            })
+            .expect("should not fail to spawn thread");
+
+        State {
+            inner: Mutex::new(StateInner {
+                update_tx,
+                join_handle: Some(join_handle),
+                next_task_id: 0,
+            }),
+        }
+    }
+
+    /// The main function of the thread that handles changes to the set of sleep-preventing
+    /// tasks and updates the system state accordingly.
+    fn thread_main(update_rx: mpsc::Receiver<StateUpdate>) {
+        let mut active_tasks: Vec<(u64, &'static str)> = Default::default();
+
+        while let Ok(task) = update_rx.recv() {
+            match task {
+                StateUpdate::AddTask { task_id, reason } => {
+                    let was_empty = active_tasks.is_empty();
+                    active_tasks.push((task_id, reason));
+
+                    // If this is the first task, prevent sleep.
+                    if was_empty {
+                        unsafe {
+                            SetThreadExecutionState(
+                                Power::ES_CONTINUOUS
+                                    | Power::ES_AWAYMODE_REQUIRED
+                                    | Power::ES_SYSTEM_REQUIRED,
+                            );
+                        }
+                    }
+
+                    Self::log_active_tasks(&active_tasks);
+                }
+                StateUpdate::RemoveTask { task_id } => {
+                    // Remove the task with this ID.
+                    active_tasks.retain(|(id, _)| *id != task_id);
+
+                    if active_tasks.is_empty() {
+                        // Allow sleep again.
+                        unsafe {
+                            SetThreadExecutionState(Power::ES_CONTINUOUS);
+                        }
+                        log::info!("No longer preventing sleep");
+                    } else {
+                        // Log remaining active reasons.
+                        Self::log_active_tasks(&active_tasks);
+                    }
+                }
+            }
+        }
+
+        // The channel was closed, so allow sleep and terminate the thread.
+        unsafe {
+            SetThreadExecutionState(Power::ES_CONTINUOUS);
+        }
+        log::warn!("Sleep-prevention thread terminating...");
+    }
+
+    fn log_active_tasks(active_tasks: &[(u64, &'static str)]) {
+        let reasons = active_tasks.iter().map(|(_, reason)| reason).collect_vec();
+        log::info!("Preventing sleep with reasons: {reasons:?}");
+    }
+
+    fn new_guard(&self, reason: &'static str) -> Guard {
+        let (task_id, update_tx) = {
+            let mut inner = self.inner.lock();
+
+            let task_id = inner.next_task_id;
+            inner.next_task_id += 1;
+
+            let update_tx = inner.update_tx.clone();
+
+            (task_id, update_tx)
+        };
+
+        if let Err(err) = update_tx.send(StateUpdate::AddTask { task_id, reason }) {
+            SEND_FAILURE.call_once(|| {
+                log::warn!("Failed to send AddTask to sleep-prevention thread: {err}");
+            });
+        }
+
+        Guard { task_id, update_tx }
+    }
+}
+
+/// The internal state of for the sleep prevention logic.
+struct StateInner {
+    update_tx: mpsc::Sender<StateUpdate>,
+    join_handle: Option<JoinHandle<()>>,
+    next_task_id: u64,
+}
+
+impl Drop for StateInner {
+    fn drop(&mut self) {
+        // Close the channel to signal the thread to exit.  We replace the
+        // sender with a new one, then drop the original sender.
+        let old_sender = std::mem::replace(&mut self.update_tx, mpsc::channel().0);
+        std::mem::drop(old_sender);
+
+        // Wait for the thread to finish.
+        if let Some(handle) = self.join_handle.take() {
+            let _ = handle.join();
+        }
+    }
+}
+
+/// A guard that prevents system sleep while it continues to exist.
+pub struct Guard {
+    task_id: u64,
+    update_tx: mpsc::Sender<StateUpdate>,
+}
+
+impl Drop for Guard {
+    fn drop(&mut self) {
+        if let Err(err) = self.update_tx.send(StateUpdate::RemoveTask {
+            task_id: self.task_id,
+        }) {
+            SEND_FAILURE.call_once(|| {
+                log::warn!("Failed to send RemoveTask to sleep-prevention thread: {err}");
+            });
+        }
+    }
+}
+
+/// Returns a guard that prevents system sleep while it remains in scope.
+pub fn prevent_sleep(reason: &'static str) -> Guard {
+    STATE.new_guard(reason)
+}