From b8d3b68e75207a584c421f961d6bb21ce92a0a88 Mon Sep 17 00:00:00 2001
From: Frederick Zhang <frederick888@tsundere.moe>
Date: Fri, 22 Jan 2021 12:53:55 +1100
Subject: [PATCH] feature: Use `ps` as fallback to query CPU usage under macOS
 (#390)

When running without elevated permissions under macOS, sysinfo cannot
query states of processes by root user, which results in 0.0% CPU usage
for all this kind of processes (and state = Unknown).

Here we use `ps`, which has SUID, as a fallback to query CPU usages.
This can be potentially applied to other properties if needed in the
future (we'll need a proper struct and parser).
---
 src/app/data_harvester/processes.rs | 51 +++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/src/app/data_harvester/processes.rs b/src/app/data_harvester/processes.rs
index 81ccc7f8..c9233337 100644
--- a/src/app/data_harvester/processes.rs
+++ b/src/app/data_harvester/processes.rs
@@ -229,6 +229,37 @@ fn get_linux_cpu_usage(
     }
 }
 
+#[cfg(target_os = "macos")]
+fn get_macos_cpu_usage(pids: &[i32]) -> std::io::Result<std::collections::HashMap<i32, f64>> {
+    use itertools::Itertools;
+    let output = std::process::Command::new("ps")
+        .args(&["-o", "pid=,pcpu=", "-p"])
+        .arg(
+            pids.iter()
+                .map(i32::to_string)
+                .intersperse(",".to_string())
+                .collect::<String>(),
+        )
+        .output()?;
+    let mut result = std::collections::HashMap::new();
+    String::from_utf8_lossy(&output.stdout)
+        .split_whitespace()
+        .chunks(2)
+        .into_iter()
+        .for_each(|chunk| {
+            let chunk: Vec<&str> = chunk.collect();
+            if chunk.len() != 2 {
+                panic!("Unexpected `ps` output");
+            }
+            let pid = chunk[0].parse();
+            let usage = chunk[1].parse();
+            if let (Ok(pid), Ok(usage)) = (pid, usage) {
+                result.insert(pid, usage);
+            }
+        });
+    Ok(result)
+}
+
 #[cfg(target_os = "linux")]
 fn get_uid_and_gid(path: &PathBuf) -> (Option<u32>, Option<u32>) {
     // FIXME: [OPT] - can we merge our /stat and /status calls?
@@ -566,6 +597,26 @@ pub fn get_process_data(
         }
     }
 
+    #[cfg(target_os = "macos")]
+    {
+        let unknown_state = ProcessStatus::Unknown(0).to_string().to_string();
+        let cpu_usage_unknown_pids: Vec<i32> = process_vector
+            .iter()
+            .filter(|process| process.process_state == unknown_state)
+            .map(|process| process.pid)
+            .collect();
+        let cpu_usages = get_macos_cpu_usage(&cpu_usage_unknown_pids)?;
+        for process in &mut process_vector {
+            if cpu_usages.contains_key(&process.pid) {
+                process.cpu_usage_percent = if num_cpus == 0.0 {
+                    *cpu_usages.get(&process.pid).unwrap()
+                } else {
+                    *cpu_usages.get(&process.pid).unwrap() / num_cpus
+                };
+            }
+        }
+    }
+
     Ok(process_vector)
 }