use std::os::unix::io::RawFd;
fn raw_fork() -> std::io::Result<i32> {
#[cfg(target_arch = "x86_64")]
const NR_FORK: i64 = 57;
#[cfg(target_arch = "aarch64")]
const NR_FORK: i64 = -1;
#[cfg(target_arch = "x86_64")]
{
let pid = unsafe { libc::syscall(NR_FORK) };
if pid < 0 {
Err(std::io::Error::last_os_error())
} else {
Ok(pid as i32)
}
}
#[cfg(target_arch = "aarch64")]
{
let pid = unsafe { libc::fork() };
if pid < 0 {
Err(std::io::Error::last_os_error())
} else {
Ok(pid)
}
}
}
pub(crate) fn fork_ready_loop_fn(
ctrl_fd: RawFd,
n: u32,
work_fn: &dyn Fn(u32),
stdout_write_fds: &[RawFd],
) {
let _ = unsafe { libc::fflush(std::ptr::null_mut()) };
let mut pids = Vec::with_capacity(n as usize);
for i in 0..n {
match raw_fork() {
Ok(0) => {
unsafe { libc::close(ctrl_fd) };
if (i as usize) < stdout_write_fds.len() && stdout_write_fds[i as usize] >= 0 {
unsafe { libc::dup2(stdout_write_fds[i as usize], 1) };
}
for &wfd in stdout_write_fds {
if wfd >= 0 { unsafe { libc::close(wfd) }; }
}
unsafe { libc::setpgid(0, 0) };
std::env::set_var("CLONE_ID", i.to_string());
work_fn(i);
unsafe { libc::fflush(std::ptr::null_mut()) };
unsafe { libc::_exit(0) };
}
Ok(pid) => {
pids.push(pid as u32);
}
Err(_) => {
pids.push(0);
}
}
}
for &wfd in stdout_write_fds {
if wfd >= 0 { unsafe { libc::close(wfd) }; }
}
let pid_bytes: Vec<u8> = pids.iter().flat_map(|p| p.to_be_bytes()).collect();
unsafe { libc::write(ctrl_fd, pid_bytes.as_ptr() as *const _, pid_bytes.len()) };
let mut exit_codes = Vec::with_capacity(pids.len());
for &pid in &pids {
if pid > 0 {
let mut status: i32 = 0;
unsafe { libc::waitpid(pid as i32, &mut status, 0) };
let code = if libc::WIFEXITED(status) { libc::WEXITSTATUS(status) } else { -1 };
exit_codes.push(code as i32);
} else {
exit_codes.push(-1);
}
}
let code_bytes: Vec<u8> = exit_codes.iter().flat_map(|c| c.to_be_bytes()).collect();
unsafe { libc::write(ctrl_fd, code_bytes.as_ptr() as *const _, code_bytes.len()) };
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_raw_fork() {
let pid = raw_fork().unwrap();
if pid == 0 {
unsafe { libc::_exit(42) };
}
let mut status: i32 = 0;
unsafe { libc::waitpid(pid, &mut status, 0) };
assert!(libc::WIFEXITED(status));
assert_eq!(libc::WEXITSTATUS(status), 42);
}
}