//! An example of an unmanaged fork where the caller is responsible for looking
//! after the children.

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

use log::{debug, info};

use nix::sys::signal::{kill, Signal};
use prefork::{Pid, Prefork, Result};
use signal_hook::{
    consts::{SIGCHLD, SIGINT, SIGQUIT, SIGTERM, SIGUSR1, SIGUSR2},
    iterator::Signals,
};

// Children wait forever and print a message when they receive SIGUSR1 or SIGUSR2.
fn child(child_num: u32, _: ()) {
    info!("Starting child {child_num}");
    let mut signals = Signals::new([SIGUSR1, SIGUSR2]).expect("catch signals");
    for signal in signals.forever() {
        info!("Child {child_num} got signal {signal}");
    }
    info!("Stopping child {child_num}");
}

// The parent forwards SIGUSR1 and SIGUSR2 to the children.
// Other signals cause the children to be reaped.
fn manage_children(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, SIGUSR1, SIGUSR2, SIGCHLD])?;
    for signal in signals.forever() {
        match signal {
            SIGCHLD => {
                prefork::wait_process(&mut pids);
                if pids.is_empty() {
                    break;
                }
            }
            SIGUSR1 | SIGUSR2 => {
                // Pass signals to children.
                let signal = Signal::try_from(signal).expect("convert i32 signal for kill");
                for child in &pids {
                    kill(child.clone(), signal).expect("forward signal to child");
                }
            }
            _ => break,
        }
    }
    prefork::kill_all(pids)
}

fn main() {
    // Send logs to tracing crate.
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::DEBUG)
        .without_time()
        .init();

    // Create a prefork server.
    let mut server = Prefork::from_resource(())
        .with_num_processes(10)
        .with_init(child)
        .server();

    let Some(pids) = server.prefork().expect("fork") else {
        // We're the child process.
        return;
    };

    // Manage signals and reap children
    manage_children(pids).expect("well behaved children");
}