Skip to main content

linuxutils_misc/
waitpid.rs

1use linuxutils_common::man::ManContent;
2
3pub const MAN: ManContent = ManContent::empty();
4
5use clap::Parser;
6use rustix::{
7    event::{PollFd, PollFlags, Timespec, poll},
8    process::{Pid, PidfdFlags, pidfd_open},
9};
10use std::{process::ExitCode, time::Instant};
11
12#[derive(Parser)]
13#[command(name = "waitpid", about = "Wait for arbitrary processes to exit")]
14pub struct Args {
15    /// Be more verbose
16    #[arg(short, long)]
17    verbose: bool,
18
19    /// Maximum number of seconds to wait
20    #[arg(short, long)]
21    timeout: Option<u64>,
22
23    /// Don't error on already-exited PIDs
24    #[arg(short, long)]
25    exited: bool,
26
27    /// Number of process exits to wait for
28    #[arg(short, long)]
29    count: Option<usize>,
30
31    /// PIDs to wait for
32    #[arg(required = true)]
33    pids: Vec<i32>,
34}
35
36pub fn run(args: Args) -> ExitCode {
37    let target_count = args.count.unwrap_or(args.pids.len());
38    let deadline = args
39        .timeout
40        .map(|t| Instant::now() + std::time::Duration::from_secs(t));
41
42    let mut pidfds = Vec::new();
43    for &raw_pid in &args.pids {
44        let pid = match Pid::from_raw(raw_pid) {
45            Some(p) => p,
46            None => {
47                eprintln!("waitpid: invalid PID: {raw_pid}");
48                return ExitCode::FAILURE;
49            }
50        };
51
52        match pidfd_open(pid, PidfdFlags::NONBLOCK) {
53            Ok(fd) => {
54                if args.verbose {
55                    eprintln!("waitpid: waiting for PID {raw_pid}");
56                }
57                pidfds.push((raw_pid, fd));
58            }
59            Err(e) => {
60                if e == rustix::io::Errno::SRCH && args.exited {
61                    if args.verbose {
62                        eprintln!("waitpid: PID {raw_pid} already exited");
63                    }
64                } else {
65                    eprintln!("waitpid: failed to open PID {raw_pid}: {e}");
66                    if e == rustix::io::Errno::NOSYS {
67                        return ExitCode::from(2);
68                    }
69                    return ExitCode::FAILURE;
70                }
71            }
72        }
73    }
74
75    let mut exited = args.pids.len() - pidfds.len(); // already-exited count
76
77    while exited < target_count && !pidfds.is_empty() {
78        let timeout = deadline.map(|d| {
79            let remaining = d.saturating_duration_since(Instant::now());
80            Timespec {
81                tv_sec: remaining.as_secs() as i64,
82                tv_nsec: remaining.subsec_nanos() as i64,
83            }
84        });
85
86        let mut fds: Vec<PollFd<'_>> = pidfds
87            .iter()
88            .map(|(_, fd)| PollFd::new(fd, PollFlags::IN))
89            .collect();
90
91        match poll(&mut fds, timeout.as_ref()) {
92            Ok(0) => {
93                // Timeout.
94                if args.verbose {
95                    eprintln!("waitpid: timeout");
96                }
97                return ExitCode::from(3);
98            }
99            Ok(_) => {
100                // Check which pidfds are ready (process exited).
101                let mut i = 0;
102                while i < pidfds.len() {
103                    let mut pfd = [PollFd::new(&pidfds[i].1, PollFlags::IN)];
104                    let ready = poll(
105                        &mut pfd,
106                        Some(&Timespec {
107                            tv_sec: 0,
108                            tv_nsec: 0,
109                        }),
110                    )
111                    .unwrap_or(0);
112                    if ready > 0 {
113                        let pid = pidfds[i].0;
114                        if args.verbose {
115                            eprintln!("waitpid: PID {pid} exited");
116                        }
117                        pidfds.swap_remove(i);
118                        exited += 1;
119                        if exited >= target_count {
120                            break;
121                        }
122                    } else {
123                        i += 1;
124                    }
125                }
126            }
127            Err(e) => {
128                if e != rustix::io::Errno::INTR {
129                    eprintln!("waitpid: poll failed: {e}");
130                    return ExitCode::FAILURE;
131                }
132            }
133        }
134
135        if let Some(d) = deadline
136            && Instant::now() >= d
137        {
138            if args.verbose {
139                eprintln!("waitpid: timeout");
140            }
141            return ExitCode::from(3);
142        }
143    }
144
145    ExitCode::SUCCESS
146}