prefork 0.6.0

A library for forking processes
Documentation 
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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

use std::{collections::HashSet, process};

use crate::{error::Result, parent};
use log::debug;
use nix::unistd::{fork, ForkResult, Pid};
use signal_hook::{
    consts::{SIGCHLD, SIGINT, SIGQUIT, SIGTERM},
    iterator::Signals,
};

/// A Prefork Server.
/// A server is used to `fork()` child processes.
pub struct Server<T> {
    num_processes: u32,
    resource: Option<T>,
    child_init: Box<dyn Fn(u32, T)>,
    child_num: u32,
}

impl<T> Server<T> {
    /// Create a server.
    /// The `resource` is passed to child processes.
    /// The `child_init()` function is called in each child process.
    /// There will be `num_processes` forked in a call to `prefork()`.
    pub fn from_resource(resource: T, child_init: Box<dyn Fn(u32, T)>, num_processes: u32) -> Self {
        Self {
            num_processes,
            resource: Some(resource),
            child_init,
            child_num: 0,
        }
    }

    /// Fork multiple processes and return their process ids if this is the parent process.
    /// # Errors
    /// Errors returned by `fork()` are propagated.
    pub fn prefork(&mut self) -> Result<Option<HashSet<Pid>>> {
        let mut pids = HashSet::new();
        let mut is_parent = true;
        for _ in 0..self.num_processes {
            let pid = self.fork()?;
            if let Some(pid) = pid {
                pids.insert(pid);
            } else {
                is_parent = false;
                break;
            }
        }
        if is_parent {
            Ok(Some(pids))
        } else {
            Ok(None)
        }
    }

    /// Fork a single process and return the process id if this is the parent process.
    /// # Errors
    /// Errors returned by `fork()` are propagated.
    pub fn fork(&mut self) -> Result<Option<Pid>> {
        match unsafe { fork()? } {
            ForkResult::Parent { child } => {
                self.child_num += 1;
                Ok(Some(child))
            }
            ForkResult::Child => {
                if let Some(resource) = self.resource.take() {
                    (self.child_init)(self.child_num, resource);
                } else {
                    unreachable!("fork resource is empty");
                }
                Ok(None)
            }
        }
    }
}

pub fn run_parent(mut pids: HashSet<Pid>) -> Result<()> {
    let parent_pid = process::id();
    debug!("Handling signals in parent {parent_pid}");
    let mut signals = Signals::new([SIGINT, SIGTERM, SIGQUIT, SIGCHLD])?;
    for signal in signals.forever() {
        match signal {
            SIGCHLD => {
                parent::wait_process(&mut pids);
                if pids.is_empty() {
                    break;
                }
            }
            _ => break,
        }
    }
    parent::kill_all(pids)
}