fallow-process 3.34.1

//! Process-wide registry of live spawned-child PIDs.
//!
//! Keyed by a monotonic `AtomicU64` counter rather than `Child::id()`
//! because POSIX recycles PIDs aggressively on long-running runners; a
//! recycled PID would collide with a previously-deregistered entry.
//!
//! Stores termination targets (not `Child` handles): the `ScopedChild` wrapper
//! owns the `Child` outright so it can call `wait_with_output` / `wait`
//! normally. Most children register one PID; subprocess wrappers can register
//! a dedicated process tree so signal cleanup also terminates descendants.

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock, TryLockError};
use std::time::{Duration, Instant};

use rustc_hash::FxHashMap;

use crate::process_tree::ProcessTree;

static NEXT_ID: AtomicU64 = AtomicU64::new(1);
static REGISTRY: OnceLock<Mutex<FxHashMap<u64, KillTarget>>> = OnceLock::new();

#[derive(Clone)]
enum KillTarget {
    Process(u32),
    ProcessTree(ProcessTreeHandle),
}

pub type ProcessTreeHandle = Arc<ProcessTree>;

/// One-shot guard: repeated signals during drain (signal storm) no-op
/// the second-and-onwards entries.
static DRAINING: AtomicU64 = AtomicU64::new(0);

fn registry() -> &'static Mutex<FxHashMap<u64, KillTarget>> {
    REGISTRY.get_or_init(|| Mutex::new(FxHashMap::default()))
}

/// Register `pid`. Returns a monotonic key the caller stores in their
/// `ScopedChild` for deregister at wait/drop time.
pub fn register(pid: u32) -> u64 {
    register_target(KillTarget::Process(pid))
}

/// Register an owned process tree for process-wide signal cleanup.
pub fn register_process_tree(process_tree: ProcessTreeHandle) -> u64 {
    register_target(KillTarget::ProcessTree(process_tree))
}

fn register_target(target: KillTarget) -> u64 {
    let id = NEXT_ID.fetch_add(1, Ordering::SeqCst);
    registry()
        .lock()
        .unwrap_or_else(|e| e.into_inner())
        .insert(id, target);
    id
}

/// Remove the registry entry for `id`. Idempotent.
pub fn deregister(id: u64) {
    registry()
        .lock()
        .unwrap_or_else(|e| e.into_inner())
        .remove(&id);
}

#[cfg(test)]
pub fn is_registered(id: u64) -> bool {
    registry()
        .lock()
        .unwrap_or_else(|error| error.into_inner())
        .contains_key(&id)
}

/// Snapshot every registered PID and kill each. Polls for liveness
/// with a bounded budget. Caller is the platform signal handler thread.
///
/// First-call-wins via the `DRAINING` guard: subsequent invocations
/// during the same shutdown skip the body to avoid re-entering the
/// lock under signal storm.
pub fn drain_and_kill() {
    if !begin_drain() {
        return;
    }
    let targets: Vec<KillTarget> = registry()
        .lock()
        .unwrap_or_else(|e| e.into_inner())
        .drain()
        .map(|(_id, target)| target)
        .collect();
    kill_and_wait(&targets);
}

/// Like [`drain_and_kill`], but returns at once when another thread holds
/// the registry lock. A panic hook calls this: a blocking lock would
/// deadlock when the panicking thread itself holds the lock.
pub fn try_drain_and_kill() {
    let targets: Vec<KillTarget> = match registry().try_lock() {
        Ok(mut guard) => drain_targets(&mut guard),
        Err(TryLockError::Poisoned(error)) => drain_targets(&mut error.into_inner()),
        Err(TryLockError::WouldBlock) => return,
    };
    kill_and_wait(&targets);
}

fn drain_targets(registry: &mut FxHashMap<u64, KillTarget>) -> Vec<KillTarget> {
    if !begin_drain() {
        return Vec::new();
    }
    registry.drain().map(|(_id, target)| target).collect()
}

fn begin_drain() -> bool {
    DRAINING
        .compare_exchange(0, 1, Ordering::SeqCst, Ordering::SeqCst)
        .is_ok()
}

fn kill_and_wait(targets: &[KillTarget]) {
    for target in targets {
        kill_target(target);
    }

    let deadline = Instant::now() + drain_budget();
    while Instant::now() < deadline {
        if !targets.iter().any(target_is_alive) {
            return;
        }
        std::thread::sleep(Duration::from_millis(50));
    }
}

fn kill_target(target: &KillTarget) {
    match target {
        KillTarget::Process(pid) => kill_pid(*pid),
        KillTarget::ProcessTree(process_tree) => {
            let _ = process_tree.terminate();
        }
    }
}

fn target_is_alive(target: &KillTarget) -> bool {
    match target {
        KillTarget::Process(pid) => pid_is_alive(*pid),
        KillTarget::ProcessTree(process_tree) => process_tree.is_alive(),
    }
}

#[cfg(unix)]
pub fn kill_pid(pid: u32) {
    let _ = std::process::Command::new("kill")
        .args(["-9", &pid.to_string()])
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .status();
}

#[cfg(windows)]
#[expect(
    unsafe_code,
    reason = "FFI to Win32 OpenProcess/TerminateProcess/CloseHandle; preconditions documented inline"
)]
pub fn kill_pid(pid: u32) {
    use windows_sys::Win32::Foundation::{CloseHandle, FALSE, HANDLE};
    use windows_sys::Win32::System::Threading::{OpenProcess, PROCESS_TERMINATE, TerminateProcess};
    // SAFETY: OpenProcess takes plain values and returns null on failure.
    let handle: HANDLE = unsafe { OpenProcess(PROCESS_TERMINATE, FALSE, pid) };
    if handle.is_null() {
        return;
    }
    // SAFETY: `handle` is a valid process handle opened with PROCESS_TERMINATE.
    let _ = unsafe { TerminateProcess(handle, 1) };
    // SAFETY: `handle` is valid and this is its only close.
    let _ = unsafe { CloseHandle(handle) };
}

#[cfg(not(any(unix, windows)))]
pub fn kill_pid(_pid: u32) {}

/// Whether `pid` is still running. On Linux a zombie counts as exited.
#[cfg(unix)]
pub fn pid_is_alive(pid: u32) -> bool {
    #[cfg(target_os = "linux")]
    if let Some(running) = crate::proc_state::process_is_running(pid) {
        return running;
    }
    std::process::Command::new("kill")
        .args(["-0", &pid.to_string()])
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .status()
        .is_ok_and(|s| s.success())
}

#[cfg(windows)]
#[expect(
    unsafe_code,
    reason = "FFI to Win32 OpenProcess/WaitForSingleObject/CloseHandle; preconditions documented inline"
)]
pub fn pid_is_alive(pid: u32) -> bool {
    use windows_sys::Win32::Foundation::{CloseHandle, FALSE, HANDLE, WAIT_OBJECT_0};
    use windows_sys::Win32::System::Threading::{
        OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION, WaitForSingleObject,
    };
    // SAFETY: OpenProcess takes plain values and returns null on failure.
    let handle: HANDLE = unsafe { OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid) };
    if handle.is_null() {
        return false;
    }
    // SAFETY: `handle` is a valid process handle; a zero timeout only polls.
    let result = unsafe { WaitForSingleObject(handle, 0) };
    // SAFETY: `handle` is valid and this is its only close.
    let _ = unsafe { CloseHandle(handle) };
    result != WAIT_OBJECT_0
}

#[cfg(not(any(unix, windows)))]
pub fn pid_is_alive(_pid: u32) -> bool {
    false
}

#[cfg(unix)]
const fn drain_budget() -> Duration {
    Duration::from_millis(500)
}

#[cfg(windows)]
const fn drain_budget() -> Duration {
    Duration::from_millis(1500)
}

#[cfg(not(any(unix, windows)))]
const fn drain_budget() -> Duration {
    Duration::from_millis(500)
}

#[cfg(test)]
#[cfg_attr(
    unix,
    expect(
        clippy::expect_used,
        reason = "test setup failures should fail at the exact setup operation"
    )
)]
mod tests {
    use super::*;

    #[test]
    fn register_deregister_roundtrip() {
        let id = register(42);
        assert!(id > 0);
        assert!(is_registered(id));
        deregister(id);
        assert!(!is_registered(id));
        deregister(id);
        assert!(!is_registered(id));
    }

    /// A panic on a thread that holds the registry lock must not deadlock
    /// the panic hook. The hook then skips the cleanup.
    #[test]
    fn try_drain_returns_while_the_lock_is_held() {
        let id = register(42);
        let guard = registry().lock().unwrap_or_else(|error| error.into_inner());
        try_drain_and_kill();
        let still_registered = guard.contains_key(&id);
        drop(guard);
        deregister(id);
        assert!(
            still_registered,
            "a held lock must leave the registry as is"
        );
    }

    #[cfg(unix)]
    #[test]
    fn register_process_tree_roundtrip() {
        let process_tree = Arc::new(ProcessTree::for_pid(42).expect("test process tree"));
        let id = register_process_tree(Arc::clone(&process_tree));
        assert!(id > 0);
        assert!(matches!(
            registry()
                .lock()
                .unwrap_or_else(|error| error.into_inner())
                .get(&id),
            Some(KillTarget::ProcessTree(registered))
                if Arc::ptr_eq(registered, &process_tree)
        ));
        deregister(id);
    }

    #[cfg(unix)]
    #[test]
    fn process_tree_target_terminates_descendants() {
        use std::fs;

        let root = tempfile::tempdir().expect("temporary process-tree root");
        let mut command = std::process::Command::new("sh");
        command
            .args([
                "-c",
                "sleep 600 & child=$!; printf '%s' \"$child\" > child.pid; wait \"$child\"",
            ])
            .current_dir(root.path());
        crate::process_tree::configure_std_command(&mut command);
        let mut leader = command.spawn().expect("spawn process tree");
        let process_tree = Arc::new(ProcessTree::for_std_child(&leader).expect("own process tree"));
        let pid_path = root.path().join("child.pid");
        // Loaded CI runners can take multi-second pauses before the shell
        // writes the pid file; poll with generous headroom and require
        // non-empty content so a mid-write read cannot race the parse
        // (issue #2112 family, Linux leg).
        let deadline = Instant::now() + Duration::from_secs(30);
        let mut pid_text = String::new();
        while Instant::now() < deadline {
            if let Ok(contents) = fs::read_to_string(&pid_path)
                && !contents.trim().is_empty()
            {
                pid_text = contents;
                break;
            }
            std::thread::sleep(Duration::from_millis(20));
        }
        let child_pid = pid_text
            .trim()
            .parse::<u32>()
            .expect("numeric descendant pid");

        kill_target(&KillTarget::ProcessTree(process_tree));
        let _ = leader.wait();
        let deadline = Instant::now() + Duration::from_secs(10);
        while pid_is_alive(child_pid) && Instant::now() < deadline {
            std::thread::sleep(Duration::from_millis(20));
        }
        assert!(!pid_is_alive(child_pid), "descendant survived tree cleanup");
    }

    /// Poll `alive` for a bounded time and return its last answer.
    #[cfg(target_os = "linux")]
    fn still_alive_after_grace(alive: impl Fn() -> bool) -> bool {
        let deadline = Instant::now() + Duration::from_secs(2);
        while alive() && Instant::now() < deadline {
            std::thread::sleep(Duration::from_millis(20));
        }
        alive()
    }

    // The drain loop must not wait for a killed child that nobody has reaped
    // yet: a zombie has already exited.
    #[cfg(target_os = "linux")]
    #[test]
    fn killed_unreaped_process_target_is_not_alive() {
        let mut child = std::process::Command::new("sleep")
            .arg("30")
            .spawn()
            .expect("spawn sleep");
        let target = KillTarget::Process(child.id());
        assert!(target_is_alive(&target), "running target counts as exited");

        kill_target(&target);
        let alive = still_alive_after_grace(|| target_is_alive(&target));
        let _ = child.wait();

        assert!(!alive, "killed and unreaped child still counts as alive");
    }

    #[cfg(target_os = "linux")]
    #[test]
    fn killed_unreaped_process_tree_target_is_not_alive() {
        let mut command = std::process::Command::new("sh");
        command.args(["-c", "sleep 30 & wait"]);
        crate::process_tree::configure_std_command(&mut command);
        let mut leader = command.spawn().expect("spawn process tree");
        let target = KillTarget::ProcessTree(Arc::new(
            ProcessTree::for_std_child(&leader).expect("own process tree"),
        ));
        assert!(target_is_alive(&target), "running target counts as exited");

        kill_target(&target);
        let alive = still_alive_after_grace(|| target_is_alive(&target));
        let _ = leader.wait();

        assert!(
            !alive,
            "killed and unreaped process tree still counts as alive"
        );
    }

    #[test]
    fn ids_are_monotonic() {
        let a = register(100);
        let b = register(200);
        assert!(b > a);
        deregister(a);
        deregister(b);
    }
}