Only poll known tracees with wait(2) (#102)

CHANGELOG.md

@@ -9,10 +9,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add accessors for `wait()` poll delay
+
 ### Changed
 
 ### Fixed
 
+- Only poll known tracees for `wait(2)` status changes
+
 ## [v0.11.0] - 2023-09-04
 
 ### Added

src/ptracer.rs

@@ -6,6 +6,7 @@ use std::io;
 use std::marker::PhantomData;
 use std::os::unix::process::CommandExt;
 use std::process::{Child, Command};
+use std::time::Duration;
 
 use nix::{
     errno::Errno,
@@ -46,8 +47,6 @@ pub type Registers = libc::user_regs_struct;
 /// Extra signal info, such as its cause.
 pub type Siginfo = libc::siginfo_t;
 
-const WALL: Option<WaitPidFlag> = Some(WaitPidFlag::__WALL);
-
 /// Linux constant defined in `include/uapi/linux/elf.h`.
 #[cfg(target_arch = "aarch64")]
 const NT_PRSTATUS: i32 = 0x1;
@@ -343,16 +342,22 @@ pub struct Ptracer {
     /// Ptrace options that will be applied to tracees, by default.
     options: Options,
 
+    /// Time to sleep for before polling tracees for new events.
+    poll_delay: Duration,
+
     /// Known tracees, and their state.
     tracees: BTreeMap<i32, State>,
 }
 
+const DEFAULT_POLL_DELAY: Duration = Duration::from_millis(100);
+
 impl Ptracer {
     pub fn new() -> Self {
         let options = Options::all();
+        let poll_delay = DEFAULT_POLL_DELAY;
         let tracees = BTreeMap::new();
 
-        Self { options, tracees }
+        Self { options, poll_delay, tracees }
     }
 
     /// Returns a reference to the default ptrace options applied to newly-spawned tracees.
@@ -365,6 +370,16 @@ impl Ptracer {
         &mut self.options
     }
 
+    /// Returns a reference to the poll delay.
+    pub fn poll_delay(&self) -> &Duration {
+        &self.poll_delay
+    }
+
+    /// Returns a mutable reference to the poll delay.
+    pub fn poll_delay_mut(&mut self) -> &mut Duration {
+        &mut self.poll_delay
+    }
+
     /// Resume the stopped tracee, delivering any pending signal.
     pub fn restart(&mut self, tracee: Tracee, restart: Restart) -> Result<()> {
         let Tracee { pid, pending, .. } = tracee;
@@ -417,21 +432,58 @@ impl Ptracer {
         r
     }
 
+    // Poll tracees for a `wait(2)` status change.
+    fn poll_tracees(&self) -> Result<Option<WaitStatus>> {
+        let flag = WaitPidFlag::__WALL | WaitPidFlag::WNOHANG;
+
+        for tracee in self.tracees.keys().copied() {
+            let pid = Pid::from_raw(tracee);
+
+            match wait::waitpid(pid, Some(flag)) {
+                Ok(WaitStatus::StillAlive) => {
+                    // Alive, no state change. Check remaining tracees.
+                    continue;
+                },
+                Ok(status) => {
+                    // One of our tracees changed state.
+                    return Ok(Some(status));
+                },
+                Err(errno) if errno == Errno::ECHILD => {
+                    // No more children to wait on: we're done.
+                    return Ok(None)
+                },
+                Err(err) => {
+                    // Something else went wrong.
+                    return Err(err.into())
+                },
+            };
+        }
+
+        // No tracee changed state.
+        Ok(None)
+    }
+
     /// Wait for some running tracee process to stop.
     ///
     /// If there are no tracees to wait on, returns `None`.
     pub fn wait(&mut self) -> Result<Option<Tracee>> {
         use Signal::*;
 
-        let status = match wait::waitpid(None, WALL) {
-            Ok(status) =>
-                status,
-            Err(errno) if errno == Errno::ECHILD =>
-                // No more children to wait on: we're done.
-                return Ok(None),
-            Err(err) =>
-                return Err(err.into()),
-        };
+        let status;
+
+        loop {
+            if self.tracees.is_empty() {
+                return Ok(None);
+            }
+
+            if let Some(new_status) = self.poll_tracees()? {
+                // A tracee changed state, examine its `wait(2)` status.
+                status = new_status;
+                break;
+            } else {
+                std::thread::sleep(self.poll_delay);
+            }
+        }
 
         let tracee = match status {
             WaitStatus::Exited(pid, _exit_code) => {

tests/tracee_died.rs

@@ -17,7 +17,7 @@ macro_rules! assert_matches {
 
 #[cfg(target_arch = "x86_64")]
 #[test]
-#[timeout(100)]
+#[timeout(1000)]
 fn test_tracee_died() -> Result<()> {
     let cmd = Command::new("true");
 

tests/wait_untraced_child.rs

@@ -0,0 +1,41 @@
+use std::process::Command;
+
+use anyhow::Result;
+use ntest::timeout;
+use pete::{Ptracer, Restart};
+
+#[test]
+#[timeout(5000)]
+fn test_wait_untraced_child() -> Result<()> {
+    // Untraced, exits before tracee.
+    let mut fast = Command::new("sleep").arg("0.1").spawn()?;
+
+    // Untraced, exits after tracee.
+    let mut slow = Command::new("sleep").arg("2").spawn()?;
+
+    // Traced.
+    let mut traceme = Command::new("sleep");
+    traceme.arg("1");
+
+    let mut tracer = Ptracer::new();
+    let mut tracee = tracer.spawn(traceme)?;
+
+    while let Some(tracee) = tracer.wait()? {
+        eprintln!("{}: {:?}", tracee.pid, tracee.stop);
+
+        tracer.restart(tracee, Restart::Continue)?;
+    }
+
+    eprintln!("waiting on tracee: {}", tracee.id());
+    println!("tracee status: {}", tracee.wait()?);
+
+    eprintln!("waiting on fast: {}", fast.id());
+    println!("fast status: {}", fast.wait()?);
+
+    eprintln!("waiting on slow: {}", slow.id());
+    println!("slow status: {}", slow.wait()?);
+
+    eprintln!("ok!");
+
+    Ok(())
+}