safe_fork/
lib.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
use std::io::{Error, Result};
use std::os::unix::process::ExitStatusExt;
use std::process::ExitStatus;

/// Ensures the current process is single-threaded.
pub fn ensure_single_threaded() -> Result<()> {
    // SAFETY: `unshare` does not have special safety requirements.
    if unsafe { libc::unshare(libc::CLONE_VM) } == 0 {
        Ok(())
    } else {
        Err(Error::last_os_error())
    }
}

/// Check if the current process is single-threaded.
pub fn is_single_threaded() -> bool {
    ensure_single_threaded().is_ok()
}

pub struct Child {
    pid: libc::pid_t,
}

impl Child {
    /// Returns the OS-assigned process identifier associated with this child.
    pub fn pid(&self) -> u32 {
        self.pid as _
    }

    pub fn join(self) -> Result<ExitStatus> {
        // SAFETY: `waitpid` does not have special safety requirements.
        let mut status = 0;
        let ret = unsafe { libc::waitpid(self.pid, &mut status, 0) };
        if ret < 0 {
            return Err(std::io::Error::last_os_error());
        }
        Ok(ExitStatus::from_raw(status))
    }
}

/// Fork the current process.
///
/// The forking process must be single-threaded. Otherwise, this call will fail.
pub fn fork() -> Result<Option<Child>> {
    ensure_single_threaded()?;

    // SAFETY: fork is safe for single-threaded process.
    match unsafe { libc::fork() } {
        -1 => Err(std::io::Error::last_os_error()),
        0 => Ok(None),
        pid => Ok(Some(Child { pid })),
    }
}

/// Fork the current process, and execute the provided closure within child process.
pub fn fork_spawn(f: impl FnOnce() -> i32) -> Result<Child> {
    Ok(match fork()? {
        Some(c) => c,
        None => {
            std::process::exit(f());
        }
    })
}

/// Fork the current process, and execute the provided closure within child process, and wait for it to complete.
pub fn fork_join(f: impl FnOnce() -> i32) -> Result<i32> {
    let exit = fork_spawn(f)?.join()?;
    Ok(exit
        .code()
        .or_else(|| exit.signal().map(|x| x + 128))
        .unwrap_or(1))
}