smarm/tests/gen_server_shutdown.rs

//! gen_server graceful shutdown.
//!
//!  - A server that does not opt in (`ctx.trap_exit()` in `init`) is stopped
//!    outright by `request_shutdown`, exactly as by `request_stop`.
//!  - A trapping server receives the request as `handle_shutdown`. The default
//!    returns `ShutdownAction::Exit`: the loop breaks and `terminate` runs on
//!    the normal (non-unwind) path, so it may block. `Continue` keeps the loop
//!    dispatching; the state later ends itself with a `StopHandle` — the only
//!    way for a gen_server to exit *normally* on its own (`request_stop` on
//!    self is an abnormal `Stopped`, which `Transient` restarts).
//!  - Other exit signals (linked peers dying) reach a trapping server via
//!    `handle_exit`.

use smarm::gen_server::{
    start, GenServer, GenServerBuilder, GenServerCtx, GenServerRef, ShutdownAction, StopHandle,
};
use smarm::supervisor::{ChildSpec, OneForOne, Restart};
use smarm::{link, monitor, request_shutdown, run, self_pid, sleep, spawn, DownReason, ExitSignal};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

#[derive(Default, Clone)]
struct Log {
    events: Arc<Mutex<Vec<&'static str>>>,
}
impl Log {
    fn push(&self, e: &'static str) {
        self.events.lock().unwrap().push(e);
    }
    fn get(&self) -> Vec<&'static str> {
        self.events.lock().unwrap().clone()
    }
}

/// A server with configurable shutdown behaviour.
struct Srv {
    log: Log,
    trap: bool,
    action: ShutdownAction,
    stop: Option<StopHandle<Srv>>,
    exits: Arc<Mutex<Vec<ExitSignal>>>,
}

impl Srv {
    fn new(log: &Log, trap: bool, action: ShutdownAction) -> Self {
        Srv {
            log: log.clone(),
            trap,
            action,
            stop: None,
            exits: Default::default(),
        }
    }
}

enum Cast {
    Note(&'static str),
    StopNow,
}

impl GenServer for Srv {
    type Call = ();
    type Reply = ();
    type Cast = Cast;
    type Info = ();
    type Timer = ();

    fn init(&mut self, ctx: &GenServerCtx<Self>) {
        if self.trap {
            ctx.trap_exit();
        }
        self.stop = Some(ctx.stop_handle());
    }
    fn handle_call(&mut self, _: ()) {}
    fn handle_cast(&mut self, c: Cast) {
        match c {
            Cast::Note(s) => self.log.push(s),
            Cast::StopNow => self.stop.as_ref().unwrap().stop(),
        }
    }
    fn handle_shutdown(&mut self) -> ShutdownAction {
        self.log.push("handle_shutdown");
        self.action
    }
    fn handle_exit(&mut self, sig: ExitSignal) {
        self.log.push("handle_exit");
        self.exits.lock().unwrap().push(sig);
    }
    fn terminate(&mut self) {
        // Allowed to block on the graceful path.
        if self.trap {
            sleep(Duration::from_millis(10));
        }
        self.log.push("terminate");
    }
}

fn spawn_settled<G: GenServer>(state: G) -> GenServerRef<G> {
    let r = start(state);
    sleep(Duration::from_millis(20)); // let init (trap_exit) run
    r
}

#[test]
fn non_trapping_server_is_stopped_outright() {
    let log = Log::default();
    let l = log.clone();
    run(move || {
        let r = spawn_settled(Srv::new(&l, false, ShutdownAction::Exit));
        let mon = monitor(r.pid());
        request_shutdown(r.pid());
        let d = mon.rx.recv().unwrap();
        assert_eq!(d.reason, DownReason::Stopped);
    });
    assert_eq!(log.get(), vec!["terminate"]);
}

#[test]
fn trapping_server_exits_normally_via_handle_shutdown() {
    let log = Log::default();
    let l = log.clone();
    run(move || {
        let r = spawn_settled(Srv::new(&l, true, ShutdownAction::Exit));
        let mon = monitor(r.pid());
        request_shutdown(r.pid());
        let d = mon.rx.recv().unwrap();
        assert_eq!(d.reason, DownReason::Exit);
    });
    assert_eq!(log.get(), vec!["handle_shutdown", "terminate"]);
}

#[test]
fn continue_keeps_dispatching_until_stop_handle() {
    let log = Log::default();
    let l = log.clone();
    run(move || {
        let r = spawn_settled(Srv::new(&l, true, ShutdownAction::Continue));
        let mon = monitor(r.pid());
        request_shutdown(r.pid());
        sleep(Duration::from_millis(20));
        r.cast(Cast::Note("after-shutdown-request")).unwrap();
        r.cast(Cast::StopNow).unwrap();
        let d = mon.rx.recv().unwrap();
        assert_eq!(d.reason, DownReason::Exit);
    });
    assert_eq!(
        log.get(),
        vec!["handle_shutdown", "after-shutdown-request", "terminate"]
    );
}

#[test]
fn stop_handle_is_a_normal_exit_that_transient_does_not_restart() {
    let starts = Arc::new(AtomicUsize::new(0));
    let s = starts.clone();
    run(move || {
        let s2 = s.clone();
        let sup = spawn(move || {
            let s3 = s2.clone();
            OneForOne::new()
                .child(ChildSpec::new(Restart::Transient, move || {
                    s3.fetch_add(1, Ordering::SeqCst);
                    let log = Log::default();
                    let r = GenServerBuilder::new(Srv::new(&log, false, ShutdownAction::Exit))
                        .under(self_pid())
                        .start();
                    r.cast(Cast::StopNow).unwrap();
                    // Block until the server is gone; a bare spawn parent
                    // returning would not itself end the server.
                    let mon = monitor(r.pid());
                    let _ = mon.rx.recv();
                }))
                .run();
        });
        sup.join().unwrap(); // returns only if the child was not restarted forever
    });
    assert_eq!(starts.load(Ordering::SeqCst), 1);
}

#[test]
fn linked_peer_death_reaches_handle_exit() {
    let log = Log::default();
    let l = log.clone();
    let alive = Arc::new(AtomicBool::new(false));
    let a = alive.clone();
    run(move || {
        let r = spawn_settled(Srv::new(&l, true, ShutdownAction::Exit));
        let pid = r.pid();
        let peer = spawn(move || {
            link(pid);
            panic!("peer dies");
        });
        let _ = peer.join();
        sleep(Duration::from_millis(20));
        r.cast(Cast::Note("still-serving")).unwrap();
        sleep(Duration::from_millis(20));
        a.store(true, Ordering::SeqCst);
        let mon = monitor(r.pid());
        drop(r); // inbox closes → clean exit
        let _ = mon.rx.recv(); // don't let the root-exit sweep race terminate
    });
    assert!(alive.load(Ordering::SeqCst));
    assert_eq!(log.get(), vec!["handle_exit", "still-serving", "terminate"]);
}

#[test]
fn gen_server_ref_shutdown_is_graceful_for_a_trapping_server() {
    let log = Log::default();
    let l = log.clone();
    run(move || {
        let r = spawn_settled(Srv::new(&l, true, ShutdownAction::Exit));
        r.shutdown();
    });
    assert_eq!(log.get(), vec!["handle_shutdown", "terminate"]);
}