Merge pull request #965 from sharkdp/codex/wait4-resource-usage

Use `wait4` for per-command resource usage on Unix

src/quantity/mod.rs

@@ -3,6 +3,8 @@ use std::marker::PhantomData;
 pub use uom::num_traits::Zero;
 pub use uom::si::f64::{Information, Ratio, Time};
 pub use uom::si::information::byte;
+#[cfg(any(not(windows), test))]
+pub use uom::si::information::kibibyte;
 pub use uom::si::ratio::ratio;
 #[cfg(any(not(windows), test))]
 pub use uom::si::time::microsecond;
@@ -55,7 +57,7 @@ pub fn format_duration_value(duration: f64, unit: Option<TimeUnit>) -> (String,
 }
 
 #[cfg(test)]
-use uom::si::{information::kibibyte, time::millisecond};
+use uom::si::time::millisecond;
 
 #[test]
 fn test_time() {

src/timer/mod.rs

@@ -24,19 +24,6 @@ use std::process::{ChildStdout, Command};
 
 use anyhow::Result;
 
-#[cfg(not(windows))]
-#[derive(Debug, Copy, Clone)]
-struct CPUTimes {
-    /// Total amount of time spent executing in user mode
-    pub user_usec: i64,
-
-    /// Total amount of time spent executing in kernel mode
-    pub system_usec: i64,
-
-    /// Maximum amount of memory used by the process, in bytes
-    pub memory_usage_byte: u64,
-}
-
 /// Discard the output of a child process.
 fn discard(output: ChildStdout) {
     const CHUNK_SIZE: usize = 64 << 10;
@@ -95,9 +82,15 @@ pub fn execute_and_measure(mut command: Command) -> Result<Measurement> {
         discard(output);
     }
 
+    #[cfg(not(windows))]
+    let (time_user, time_system, peak_memory_usage, exit_status) = cpu_timer.stop(child)?;
+
+    #[cfg(windows)]
     let exit_status = child.wait()?;
 
     let time_wall_clock = wallclock_timer.stop();
+
+    #[cfg(windows)]
     let (time_user, time_system, peak_memory_usage) = cpu_timer.stop();
 
     Ok(Measurement {

src/timer/unix_timer.rs

@@ -1,108 +1,89 @@
 #![cfg(not(windows))]
 
-use std::convert::TryFrom;
-use std::mem;
+use std::io;
+use std::mem::MaybeUninit;
+use std::os::unix::process::ExitStatusExt;
+use std::process::{Child, ExitStatus};
 
-use crate::quantity::{byte, microsecond, Information, Time};
-use crate::timer::CPUTimes;
+use anyhow::Result;
+
+use crate::quantity::{byte, kibibyte, microsecond, second, Information, Time};
 
 #[derive(Debug, Copy, Clone)]
-pub struct CPUInterval {
+struct ResourceUsage {
     /// Total amount of time spent executing in user mode
-    pub user: Time,
+    pub time_user: Time,
 
     /// Total amount of time spent executing in kernel mode
-    pub system: Time,
+    pub time_system: Time,
+
+    /// Maximum amount of memory used by the process, in bytes
+    pub memory_usage: Information,
 }
 
-pub struct CPUTimer {
-    start_cpu: CPUTimes,
+#[allow(clippy::useless_conversion)]
+fn convert_timeval(tv: libc::timeval) -> Time {
+    let sec = tv.tv_sec as f64;
+    let usec = tv.tv_usec as f64;
+
+    Time::new::<second>(sec) + Time::new::<microsecond>(usec)
 }
 
-impl CPUTimer {
-    pub fn start() -> Self {
-        CPUTimer {
-            start_cpu: get_cpu_times(),
+#[allow(clippy::useless_conversion)]
+fn wait4(mut child: Child) -> io::Result<(ExitStatus, ResourceUsage)> {
+    drop(child.stdin.take());
+
+    let pid = child.id() as i32;
+    let mut status = 0;
+    let mut rusage = MaybeUninit::zeroed();
+
+    loop {
+        let result = unsafe { libc::wait4(pid, &mut status, 0, rusage.as_mut_ptr()) };
+        if result >= 0 {
+            break;
+        }
+        let error = io::Error::last_os_error();
+        if error.kind() != io::ErrorKind::Interrupted {
+            return Err(error);
         }
     }
 
-    pub fn stop(&self) -> (Time, Time, Information) {
-        let end_cpu = get_cpu_times();
-        let cpu_interval = cpu_time_interval(&self.start_cpu, &end_cpu);
-        (
-            cpu_interval.user,
-            cpu_interval.system,
-            Information::new::<byte>(end_cpu.memory_usage_byte as f64),
-        )
-    }
-}
+    let rusage = unsafe { rusage.assume_init() };
 
-/// Read CPU execution times ('user' and 'system')
-fn get_cpu_times() -> CPUTimes {
-    use libc::{getrusage, rusage, RUSAGE_CHILDREN};
-
-    let result: rusage = unsafe {
-        let mut buf = mem::zeroed();
-        let success = getrusage(RUSAGE_CHILDREN, &mut buf);
-        assert_eq!(0, success);
-        buf
-    };
-
-    const MICROSEC_PER_SEC: i64 = 1000 * 1000;
-
-    // Linux and *BSD return the value in KibiBytes, Darwin flavors in bytes
-    let max_rss_byte = if cfg!(target_os = "macos") || cfg!(target_os = "ios") {
-        result.ru_maxrss
+    let memory_usage_byte = if cfg!(target_os = "macos") || cfg!(target_os = "ios") {
+        // Linux and *BSD return the value in KibiBytes, Darwin flavors in bytes
+        // A f64 can represent integers up to 2^53 exactly, so we can represent
+        // all values up to 2^53 bytes = 8 PiB exactly. Beyond that, our precision
+        // is larger than 1 byte, but that's considered acceptable.
+        Information::new::<byte>(rusage.ru_maxrss as f64)
     } else {
-        result.ru_maxrss * 1024
+        Information::new::<kibibyte>(rusage.ru_maxrss as f64)
     };
 
-    #[allow(clippy::useless_conversion)]
-    CPUTimes {
-        user_usec: i64::from(result.ru_utime.tv_sec) * MICROSEC_PER_SEC
-            + i64::from(result.ru_utime.tv_usec),
-        system_usec: i64::from(result.ru_stime.tv_sec) * MICROSEC_PER_SEC
-            + i64::from(result.ru_stime.tv_usec),
-        memory_usage_byte: u64::try_from(max_rss_byte).unwrap_or(0),
+    Ok((
+        ExitStatus::from_raw(status),
+        ResourceUsage {
+            time_user: convert_timeval(rusage.ru_utime),
+            time_system: convert_timeval(rusage.ru_stime),
+            memory_usage: memory_usage_byte.into(),
+        },
+    ))
+}
+
+pub struct CPUTimer {}
+
+impl CPUTimer {
+    pub fn start() -> Self {
+        Self {}
+    }
+
+    pub fn stop(&self, child: Child) -> Result<(Time, Time, Information, ExitStatus)> {
+        let (status, usage) = wait4(child)?;
+        Ok((
+            usage.time_user,
+            usage.time_system,
+            usage.memory_usage,
+            status,
+        ))
     }
 }
-
-/// Compute the time intervals in between two `CPUTimes` snapshots
-fn cpu_time_interval(start: &CPUTimes, end: &CPUTimes) -> CPUInterval {
-    CPUInterval {
-        user: Time::new::<microsecond>((end.user_usec - start.user_usec) as f64),
-        system: Time::new::<microsecond>((end.system_usec - start.system_usec) as f64),
-    }
-}
-
-#[cfg(test)]
-use crate::quantity::second;
-#[cfg(test)]
-use approx::assert_relative_eq;
-
-#[test]
-fn test_cpu_time_interval() {
-    let t_a = CPUTimes {
-        user_usec: 12345,
-        system_usec: 54321,
-        memory_usage_byte: 0,
-    };
-
-    let t_b = CPUTimes {
-        user_usec: 20000,
-        system_usec: 70000,
-        memory_usage_byte: 0,
-    };
-
-    let t_zero = cpu_time_interval(&t_a, &t_a);
-    assert!(t_zero.user.get::<second>().abs() < f64::EPSILON);
-    assert!(t_zero.system.get::<second>().abs() < f64::EPSILON);
-
-    let t_ab = cpu_time_interval(&t_a, &t_b);
-    assert_relative_eq!(0.007655, t_ab.user.get::<second>());
-    assert_relative_eq!(0.015679, t_ab.system.get::<second>());
-
-    let t_ba = cpu_time_interval(&t_b, &t_a);
-    assert_relative_eq!(-0.007655, t_ba.user.get::<second>());
-    assert_relative_eq!(-0.015679, t_ba.system.get::<second>());
-}