perf: cache windows foreground process snapshots

refs #1158

1 files changed, 77 insertions, 2 deletions

diff --git a/src/platform/windows.rs b/src/platform/windows.rs
index ef41e662..317ac495 100644
--- a/src/platform/windows.rs
+++ b/src/platform/windows.rs
@@ -4,6 +4,8 @@ use std::{
     mem::{size_of, MaybeUninit},
     path::PathBuf,
     ptr::{copy_nonoverlapping, null_mut},
+    sync::{Arc, Mutex},
+    time::{Duration, Instant},
 };
 
 use windows_sys::{
@@ -39,6 +41,21 @@ use windows_sys::{
 use super::{ClipboardImage, ForegroundJob, Signal};
 
 const STILL_ACTIVE: u32 = 259;
+const FOREGROUND_PROCESS_SNAPSHOT_CACHE_TTL: Duration = Duration::from_millis(250);
+
+#[derive(Debug)]
+struct CachedProcessSnapshot {
+    built_at: Instant,
+    entries: Arc<Vec<WindowsProcessEntry>>,
+}
+
+#[derive(Debug)]
+struct ProcessSnapshotCache {
+    cached: Option<CachedProcessSnapshot>,
+}
+
+static FOREGROUND_PROCESS_SNAPSHOT_CACHE: Mutex<ProcessSnapshotCache> =
+    Mutex::new(ProcessSnapshotCache { cached: None });
 
 pub(crate) fn should_draw_host_cursor_by_default() -> bool {
     true
@@ -238,7 +255,7 @@ fn available_pane_shell_from_snapshot(
 }
 
 pub fn foreground_group_leader_job(process_group_id: u32) -> Option<ForegroundJob> {
-    let entries = snapshot_processes();
+    let entries = cached_foreground_processes();
     let entry = entries.iter().find(|entry| entry.pid == process_group_id)?;
     Some(ForegroundJob {
         process_group_id,
@@ -247,7 +264,8 @@ pub fn foreground_group_leader_job(process_group_id: u32) -> Option<ForegroundJo
 }
 
 pub fn foreground_process_group_id(child_pid: u32) -> Option<u32> {
-    foreground_job(child_pid).map(|job| job.process_group_id)
+    let entries = cached_foreground_processes();
+    select_pane_foreground_job(child_pid, &entries).map(|job| job.process_group_id)
 }
 
 pub fn process_cwd(pid: u32) -> Option<PathBuf> {
@@ -418,6 +436,34 @@ fn snapshot_processes() -> Vec<WindowsProcessEntry> {
     output
 }
 
+fn cached_foreground_processes() -> Arc<Vec<WindowsProcessEntry>> {
+    let mut cache = FOREGROUND_PROCESS_SNAPSHOT_CACHE
+        .lock()
+        .unwrap_or_else(|err| err.into_inner());
+    cache.snapshot(FOREGROUND_PROCESS_SNAPSHOT_CACHE_TTL, snapshot_processes)
+}
+
+impl ProcessSnapshotCache {
+    fn snapshot(
+        &mut self,
+        max_age: Duration,
+        build: impl FnOnce() -> Vec<WindowsProcessEntry>,
+    ) -> Arc<Vec<WindowsProcessEntry>> {
+        if let Some(cached) = &self.cached {
+            if cached.built_at.elapsed() < max_age {
+                return Arc::clone(&cached.entries);
+            }
+        }
+
+        let entries = Arc::new(build());
+        self.cached = Some(CachedProcessSnapshot {
+            built_at: Instant::now(),
+            entries: Arc::clone(&entries),
+        });
+        entries
+    }
+}
+
 fn process_command_line(pid: u32) -> Option<String> {
     let process = ProcessHandle::open(pid, PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_VM_READ)?;
     let parameters = read_process_parameters(process.0)?;
@@ -744,6 +790,7 @@ mod tests {
     use std::{
         fs,
         process::{Command, Stdio},
+        sync::Arc,
         thread,
         time::{Duration, Instant},
     };
@@ -1049,6 +1096,34 @@ mod tests {
     }
 
     #[test]
+    fn windows_foreground_process_snapshot_is_shared_within_ttl() {
+        let mut cache = super::ProcessSnapshotCache { cached: None };
+        let mut builds = 0;
+        let mut first_build_completed_at = None;
+
+        let first = cache.snapshot(Duration::from_secs(60), || {
+            builds += 1;
+            let entries = vec![test_entry(10, 1, "powershell.exe", &["powershell.exe"])];
+            first_build_completed_at = Some(Instant::now());
+            entries
+        });
+        assert!(cache.cached.as_ref().unwrap().built_at >= first_build_completed_at.unwrap());
+        let second = cache.snapshot(Duration::from_secs(60), || {
+            builds += 1;
+            Vec::new()
+        });
+        let refreshed = cache.snapshot(Duration::ZERO, || {
+            builds += 1;
+            vec![test_entry(20, 1, "pwsh.exe", &["pwsh.exe"])]
+        });
+
+        assert!(Arc::ptr_eq(&first, &second));
+        assert!(!Arc::ptr_eq(&second, &refreshed));
+        assert_eq!(builds, 2);
+        assert_eq!(refreshed[0].pid, 20);
+    }
+
+    #[test]
     fn windows_process_tree_selects_wrapped_agent_descendant() {
         let entries = vec![
             test_entry(10, 1, "cmd.exe", &["cmd.exe"]),