prefork/
server.rs

1use std::{collections::HashSet, process};
2
3use crate::{error::Result, parent};
4use log::debug;
5use nix::unistd::{fork, ForkResult, Pid};
6use signal_hook::{
7    consts::{SIGCHLD, SIGINT, SIGQUIT, SIGTERM},
8    iterator::Signals,
9};
10
11/// A Prefork Server.
12/// A server is used to `fork()` child processes.
13pub struct Server<T> {
14    num_processes: u32,
15    resource: Option<T>,
16    child_init: Box<dyn Fn(u32, T)>,
17    child_num: u32,
18}
19
20impl<T> Server<T> {
21    /// Create a server.
22    /// The `resource` is passed to child processes.
23    /// The `child_init()` function is called in each child process.
24    /// There will be `num_processes` forked in a call to `prefork()`.
25    pub fn from_resource(resource: T, child_init: Box<dyn Fn(u32, T)>, num_processes: u32) -> Self {
26        Self {
27            num_processes,
28            resource: Some(resource),
29            child_init,
30            child_num: 0,
31        }
32    }
33
34    /// Fork multiple processes and return their process ids if this is the parent process.
35    /// # Errors
36    /// Errors returned by `fork()` are propagated.
37    pub fn prefork(&mut self) -> Result<Option<HashSet<Pid>>> {
38        let mut pids = HashSet::new();
39        let mut is_parent = true;
40        for _ in 0..self.num_processes {
41            let pid = self.fork()?;
42            if let Some(pid) = pid {
43                pids.insert(pid);
44            } else {
45                is_parent = false;
46                break;
47            }
48        }
49        if is_parent {
50            Ok(Some(pids))
51        } else {
52            Ok(None)
53        }
54    }
55
56    /// Fork a single process and return the process id if this is the parent process.
57    /// # Errors
58    /// Errors returned by `fork()` are propagated.
59    pub fn fork(&mut self) -> Result<Option<Pid>> {
60        match unsafe { fork()? } {
61            ForkResult::Parent { child } => {
62                self.child_num += 1;
63                Ok(Some(child))
64            }
65            ForkResult::Child => {
66                if let Some(resource) = self.resource.take() {
67                    (self.child_init)(self.child_num, resource);
68                } else {
69                    unreachable!("fork resource is empty");
70                }
71                Ok(None)
72            }
73        }
74    }
75}
76
77pub fn run_parent(mut pids: HashSet<Pid>) -> Result<()> {
78    let parent_pid = process::id();
79    debug!("Handling signals in parent {parent_pid}");
80    let mut signals = Signals::new([SIGINT, SIGTERM, SIGQUIT, SIGCHLD])?;
81    for signal in signals.forever() {
82        match signal {
83            SIGCHLD => {
84                parent::wait_process(&mut pids);
85                if pids.is_empty() {
86                    break;
87                }
88            }
89            _ => break,
90        }
91    }
92    parent::kill_all(pids)
93}
prefork/server.rs

prefork/
server.rs
