use std::process::{Child, Command};
pub(super) fn configure(command: &mut Command) {
#[cfg(unix)]
{
use std::os::unix::process::CommandExt;
command.process_group(0);
}
#[cfg(not(unix))]
let _ = command;
}
pub(super) struct ProcessTree {
#[cfg(unix)]
group: rustix::process::Pid,
#[cfg(windows)]
job: windows_sys::Win32::Foundation::HANDLE,
}
impl ProcessTree {
#[cfg(unix)]
pub(super) fn attach(child: &Child) -> std::io::Result<Self> {
let raw = i32::try_from(child.id()).map_err(std::io::Error::other)?;
let group = rustix::process::Pid::from_raw(raw)
.ok_or_else(|| std::io::Error::other("owned process group must have a positive PID"))?;
Ok(Self { group })
}
#[cfg(windows)]
#[allow(unsafe_code)]
pub(super) fn attach(child: &Child) -> std::io::Result<Self> {
use std::os::windows::io::AsRawHandle;
use windows_sys::Win32::System::JobObjects::*;
let job = unsafe { CreateJobObjectW(std::ptr::null(), std::ptr::null()) };
if job.is_null() {
return Err(std::io::Error::last_os_error());
}
let tree = Self { job };
let mut limits: JOBOBJECT_EXTENDED_LIMIT_INFORMATION = unsafe { std::mem::zeroed() };
limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
let ok = unsafe {
SetInformationJobObject(
job,
JobObjectExtendedLimitInformation,
(&limits as *const JOBOBJECT_EXTENDED_LIMIT_INFORMATION).cast(),
std::mem::size_of_val(&limits) as u32,
) != 0
&& AssignProcessToJobObject(job, child.as_raw_handle()) != 0
};
if !ok {
return Err(std::io::Error::last_os_error());
}
Ok(tree)
}
#[cfg(not(any(unix, windows)))]
pub(super) fn attach(_child: &Child) -> std::io::Result<Self> {
Err(std::io::Error::other(
"native process containment is unsupported",
))
}
#[cfg(unix)]
pub(super) fn terminate(&self) -> std::io::Result<()> {
match rustix::process::kill_process_group(self.group, rustix::process::Signal::KILL) {
Ok(()) | Err(rustix::io::Errno::SRCH) => Ok(()),
Err(error) => Err(error.into()),
}
}
#[cfg(windows)]
#[allow(unsafe_code)]
pub(super) fn terminate(&self) -> std::io::Result<()> {
if unsafe { windows_sys::Win32::System::JobObjects::TerminateJobObject(self.job, 1) } == 0 {
return Err(std::io::Error::last_os_error());
}
Ok(())
}
#[cfg(not(any(unix, windows)))]
pub(super) fn terminate(&self) -> std::io::Result<()> {
Ok(())
}
}
#[cfg(windows)]
impl Drop for ProcessTree {
#[allow(unsafe_code)]
fn drop(&mut self) {
unsafe { windows_sys::Win32::Foundation::CloseHandle(self.job) };
}
}
#[cfg(all(test, unix))]
mod tests {
use super::*;
use std::{
io::{BufRead, BufReader},
process::Stdio,
};
#[test]
fn terminating_owned_group_closes_descendant_pipe_before_releasing_custody() {
let mut command = Command::new("/bin/sh");
command
.args(["-c", "sleep 60 & echo ready; wait"])
.stdout(Stdio::piped());
configure(&mut command);
let mut child = command.spawn().unwrap();
let tree = ProcessTree::attach(&child).unwrap();
assert_ne!(tree.group.as_raw_nonzero().get(), 0);
let mut output = BufReader::new(child.stdout.take().unwrap());
let mut line = String::new();
output.read_line(&mut line).unwrap();
assert_eq!(line.trim(), "ready");
tree.terminate().unwrap();
child.wait().unwrap();
let (send, receive) = std::sync::mpsc::channel();
std::thread::spawn(move || {
let _ = send.send(output.read_line(&mut String::new()));
});
assert_eq!(
receive
.recv_timeout(std::time::Duration::from_secs(3))
.unwrap()
.unwrap(),
0,
"descendant must not keep the provider pipe alive"
);
}
}