smarm/tests/supervisor_shutdown.rs

//! Supervisor shutdown — the OTP child-spec `shutdown` policy.
//!
//! A supervisor traps exits. A `request_shutdown` reaching it (from its parent
//! supervisor, or from the app via `request_shutdown`/`RuntimeHandle`) runs
//! the ordered shutdown: children are stopped in reverse start order, each
//! per its `Shutdown` policy — `request_shutdown`, wait up to the timeout for
//! its termination signal, `request_stop` if it overstays — and then `run()`
//! returns normally. Every supervisor-initiated child stop (ordered shutdown,
//! OneForAll/RestForOne sibling cycling) goes through the same policy.
//!
//! A *hard* `request_stop` on a supervisor unwinds it; a drop guard then
//! hard-stops its live children so the subtree is never orphaned.

use smarm::supervisor::{ChildSpec, OneForOne, Restart, Shutdown, Strategy};
use smarm::{
    monitor, request_shutdown, request_stop, run, sleep, spawn, trap_exit, DownReason, JoinHandle,
};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

/// A child that traps exits, records the order it was shut down in, and exits
/// normally on the request (after `delay`). Ignores the request if `comply`
/// is false — a straggler that must be hard-stopped.
fn polite_child(
    tag: usize,
    log: &Arc<Mutex<Vec<usize>>>,
    delay: Duration,
    comply: bool,
) -> impl Fn() + Send + Sync + 'static {
    let log = log.clone();
    move || {
        let inbox = trap_exit();
        loop {
            let sig = match inbox.recv() {
                Ok(s) => s,
                Err(_) => return,
            };
            if sig.reason == DownReason::Shutdown {
                log.lock().unwrap().push(tag);
                if comply {
                    sleep(delay);
                    return;
                }
                // Not complying: keep running until hard-stopped.
                loop {
                    sleep(Duration::from_millis(5));
                }
            }
        }
    }
}

/// Spawn `sup`, let its children reach `trap_exit`, return the handle.
fn spawn_settled(sup: OneForOne) -> JoinHandle {
    let h = spawn(move || sup.run());
    sleep(Duration::from_millis(30));
    h
}

#[test]
fn shutdown_stops_children_in_reverse_order_and_returns_normally() {
    let log = Arc::new(Mutex::new(Vec::new()));
    let l = log.clone();
    run(move || {
        let sup = OneForOne::new()
            .child(ChildSpec::new(
                Restart::Permanent,
                polite_child(1, &l, Duration::ZERO, true),
            ))
            .child(ChildSpec::new(
                Restart::Permanent,
                polite_child(2, &l, Duration::ZERO, true),
            ))
            .child(ChildSpec::new(
                Restart::Permanent,
                polite_child(3, &l, Duration::ZERO, true),
            ));
        let h = spawn_settled(sup);
        let mon = monitor(h.pid());
        request_shutdown(h.pid());
        let down = mon.rx.recv().expect("down");
        assert_eq!(
            down.reason,
            DownReason::Exit,
            "supervisor exits normally after shutdown"
        );
    });
    assert_eq!(*log.lock().unwrap(), vec![3, 2, 1]);
}

#[test]
fn non_trapping_child_is_simply_stopped() {
    let dropped = Arc::new(AtomicBool::new(false));
    let d = dropped.clone();
    run(move || {
        struct G(Arc<AtomicBool>);
        impl Drop for G {
            fn drop(&mut self) {
                self.0.store(true, Ordering::SeqCst);
            }
        }
        let sup = OneForOne::new().child(ChildSpec::new(Restart::Permanent, move || {
            let _g = G(d.clone());
            loop {
                sleep(Duration::from_millis(5));
            }
        }));
        let h = spawn_settled(sup);
        request_shutdown(h.pid());
        h.join().expect("sup");
    });
    assert!(dropped.load(Ordering::SeqCst));
}

#[test]
fn straggler_is_hard_stopped_after_timeout() {
    let log = Arc::new(Mutex::new(Vec::new()));
    let l = log.clone();
    run(move || {
        let sup = OneForOne::new().child(
            ChildSpec::new(
                Restart::Permanent,
                polite_child(1, &l, Duration::ZERO, false),
            )
            .shutdown(Shutdown::Timeout(Duration::from_millis(50))),
        );
        let h = spawn_settled(sup);
        let t0 = Instant::now();
        request_shutdown(h.pid());
        h.join().expect("sup");
        let took = t0.elapsed();
        assert!(
            took >= Duration::from_millis(50),
            "returned before the grace period: {took:?}"
        );
        assert!(
            took < Duration::from_secs(2),
            "did not fall back to a hard stop: {took:?}"
        );
    });
    assert_eq!(
        *log.lock().unwrap(),
        vec![1],
        "the straggler did receive the request"
    );
}

#[test]
fn infinity_waits_for_a_slow_but_compliant_child() {
    let log = Arc::new(Mutex::new(Vec::new()));
    let l = log.clone();
    let finished = Arc::new(AtomicBool::new(false));
    let f = finished.clone();
    run(move || {
        let f2 = f.clone();
        let l2 = l.clone();
        let sup = OneForOne::new().child(
            ChildSpec::new(Restart::Permanent, move || {
                let inbox = trap_exit();
                let _ = inbox.recv();
                l2.lock().unwrap().push(1);
                sleep(Duration::from_millis(150));
                f2.store(true, Ordering::SeqCst); // only reached if not hard-stopped
            })
            .shutdown(Shutdown::Infinity),
        );
        let h = spawn_settled(sup);
        request_shutdown(h.pid());
        h.join().expect("sup");
    });
    assert!(
        finished.load(Ordering::SeqCst),
        "Infinity must not hard-stop a compliant child"
    );
}

#[test]
fn brutal_kill_skips_the_request() {
    let log = Arc::new(Mutex::new(Vec::new()));
    let l = log.clone();
    run(move || {
        let sup = OneForOne::new().child(
            ChildSpec::new(
                Restart::Permanent,
                polite_child(1, &l, Duration::ZERO, true),
            )
            .shutdown(Shutdown::BrutalKill),
        );
        let h = spawn_settled(sup);
        request_shutdown(h.pid());
        h.join().expect("sup");
    });
    assert!(
        log.lock().unwrap().is_empty(),
        "a BrutalKill child never sees the request"
    );
}

#[test]
fn hard_stop_of_supervisor_does_not_orphan_children() {
    let alive = Arc::new(AtomicUsize::new(0));
    let a = alive.clone();
    run(move || {
        struct Alive(Arc<AtomicUsize>);
        impl Drop for Alive {
            fn drop(&mut self) {
                self.0.fetch_sub(1, Ordering::SeqCst);
            }
        }
        let mk = |a: Arc<AtomicUsize>| {
            move || {
                a.fetch_add(1, Ordering::SeqCst);
                let _g = Alive(a.clone());
                loop {
                    sleep(Duration::from_millis(5));
                }
            }
        };
        let sup = OneForOne::new()
            .child(ChildSpec::new(Restart::Permanent, mk(a.clone())))
            .child(ChildSpec::new(Restart::Permanent, mk(a.clone())));
        let h = spawn_settled(sup);
        assert_eq!(a.load(Ordering::SeqCst), 2);
        let mon = monitor(h.pid());
        request_stop(h.pid());
        let _ = mon.rx.recv();
        sleep(Duration::from_millis(50));
        assert_eq!(
            a.load(Ordering::SeqCst),
            0,
            "children orphaned by a hard supervisor stop"
        );
    });
}

#[test]
fn nested_shutdown_reaches_grandchildren() {
    let log = Arc::new(Mutex::new(Vec::new()));
    let l = log.clone();
    run(move || {
        let l_inner = l.clone();
        let inner = move || {
            OneForOne::new()
                .child(ChildSpec::new(
                    Restart::Permanent,
                    polite_child(10, &l_inner, Duration::ZERO, true),
                ))
                .child(ChildSpec::new(
                    Restart::Permanent,
                    polite_child(11, &l_inner, Duration::ZERO, true),
                ))
                .run()
        };
        let sup = OneForOne::new()
            .child(ChildSpec::new(
                Restart::Permanent,
                polite_child(1, &l, Duration::ZERO, true),
            ))
            .child(ChildSpec::new(Restart::Permanent, inner).shutdown(Shutdown::Infinity));
        let h = spawn_settled(sup);
        request_shutdown(h.pid());
        h.join().expect("sup");
    });
    assert_eq!(*log.lock().unwrap(), vec![11, 10, 1]);
}

#[test]
fn sibling_cycling_uses_graceful_shutdown() {
    // OneForAll: when child A dies, sibling B (trapping) must receive a
    // Shutdown request rather than a bare stop.
    let log = Arc::new(Mutex::new(Vec::new()));
    let l = log.clone();
    let a_runs = Arc::new(AtomicUsize::new(0));
    let ar = a_runs.clone();
    run(move || {
        let ar2 = ar.clone();
        let sup = OneForOne::new()
            .strategy(Strategy::OneForAll)
            .intensity(5, Duration::from_secs(60))
            .child(ChildSpec::new(Restart::Transient, move || {
                let n = ar2.fetch_add(1, Ordering::SeqCst) + 1;
                sleep(Duration::from_millis(30));
                if n == 1 {
                    panic!("first run dies");
                }
                // Second run: park until shut down.
                let inbox = trap_exit();
                let _ = inbox.recv();
            }))
            .child(ChildSpec::new(
                Restart::Permanent,
                polite_child(2, &l, Duration::ZERO, true),
            ));
        let h = spawn(move || sup.run());
        sleep(Duration::from_millis(150));
        request_shutdown(h.pid());
        h.join().expect("sup");
    });
    // B was shut down once by the cycle and once by the final shutdown.
    assert_eq!(*log.lock().unwrap(), vec![2, 2]);
    assert_eq!(a_runs.load(Ordering::SeqCst), 2);
}