smarm/examples/graceful_shutdown.rs

//! Graceful shutdown, end to end: a supervised app tree, a server that
//! drains before it exits, and the two ways the whole thing winds down.
//!
//! Stopping an actor comes in two strengths, as in OTP:
//!   - `request_stop(pid)`      = `exit(Pid, kill)`: cooperative hard stop,
//!                                unwinds at the next observation point.
//!   - `request_shutdown(pid)`  = `exit(Pid, shutdown)`: a trapping target gets
//!                                an `ExitSignal { reason: Shutdown }` and winds
//!                                down on its own terms; a non-trapping one is
//!                                stopped outright.
//!
//! A supervisor traps exits. `request_shutdown(sup)` runs its ordered
//! shutdown — children in reverse start order, each per its `ChildSpec`
//! `Shutdown` policy (`Timeout(d)` default 5s, `Infinity`, `BrutalKill`) —
//! and the supervisor then returns normally.
//!
//! Two triggers are shown:
//!   1. **Root exit.** The run's root actor returning means "the program is
//!      done": the runtime delivers `request_shutdown` to every top-level actor
//!      (here: the supervisor). Trapping actors may keep running to drain and
//!      end the run when they stop themselves; non-trapping ones are stopped.
//!   2. **An outside thread** (e.g. a signal handler) driving it via
//!      `RuntimeHandle::request_shutdown` on the supervisor — the root then
//!      just waits for the tree to come down.

use smarm::gen_server::{
    GenServer, GenServerBuilder, GenServerCtx, ShutdownAction, StopHandle, TimerHandle,
};
use smarm::supervisor::{ChildSpec, OneForOne, Restart, Shutdown};
use smarm::{monitor, request_shutdown, self_pid, sleep, spawn, trap_exit, DownReason};
use std::thread;
use std::time::Duration;

/// A server with in-flight work: on shutdown it stops accepting, finishes what
/// it has (simulated with a ticking timer), then ends itself.
struct Drainer {
    pending: u32,
    stop: Option<StopHandle<Drainer>>,
    timer: Option<TimerHandle<Drainer>>,
}

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

    fn init(&mut self, ctx: &GenServerCtx<Self>) {
        ctx.trap_exit(); // opt in: shutdown arrives as handle_shutdown
        self.stop = Some(ctx.stop_handle());
        self.timer = Some(ctx.timer());
    }
    fn handle_call(&mut self, _: ()) {}
    fn handle_cast(&mut self, _: ()) {}
    fn handle_shutdown(&mut self) -> ShutdownAction {
        println!(
            "drainer: shutdown requested, {} items pending",
            self.pending
        );
        self.timer
            .as_ref()
            .unwrap()
            .tick_every(Duration::from_millis(20), ());
        ShutdownAction::Continue // keep serving until drained
    }
    fn handle_timer(&mut self, _: ()) {
        self.pending -= 1;
        if self.pending == 0 {
            println!("drainer: drained, stopping");
            self.stop.as_ref().unwrap().stop(); // normal exit
        }
    }
    fn terminate(&mut self) {
        // Graceful path: this runs on the normal path and may block.
        println!("drainer: terminate");
    }
}

/// A supervised child wrapping the server. (A gen_server is not yet directly
/// usable as a `ChildSpec` start fn; the wrapper traps, forwards the shutdown,
/// and waits for the server to finish. See ROADMAP "open items".)
fn drainer_child() {
    let inbox = trap_exit();
    let srv = GenServerBuilder::new(Drainer {
        pending: 3,
        stop: None,
        timer: None,
    })
    .under(self_pid())
    .start();
    let mon = monitor(srv.pid());
    // Wait for our shutdown, forward it, wait for the server.
    while let Ok(sig) = inbox.recv() {
        if sig.reason == DownReason::Shutdown {
            request_shutdown(srv.pid());
            let _ = mon.rx.recv();
            return;
        }
    }
}

fn app_tree() -> OneForOne {
    OneForOne::new()
        .child(
            ChildSpec::new(Restart::Permanent, || {
                // A plain worker that does not trap: stopped outright on shutdown.
                loop {
                    sleep(Duration::from_millis(10));
                }
            })
            .shutdown(Shutdown::Timeout(Duration::from_millis(100))),
        )
        .child(ChildSpec::new(Restart::Permanent, drainer_child).shutdown(Shutdown::Infinity))
}

fn main() {
    println!("--- 1. root exit drives the shutdown ---");
    smarm::run(|| {
        spawn(|| app_tree().run());
        sleep(Duration::from_millis(50)); // the app "runs" for a while
                                          // Returning here asks the supervisor to shut down; the run ends when
                                          // the tree — drainer included — is gone.
    });

    println!("--- 2. an outside thread drives the shutdown ---");
    let rt = smarm::init(smarm::Config::default());
    let handle = rt.handle(); // Send + Sync; grab it before run
    rt.run(move || {
        let sup = spawn(|| app_tree().run());
        let sup_pid = sup.pid();
        // Stand-in for a SIGTERM handler thread.
        thread::spawn(move || {
            thread::sleep(Duration::from_millis(50));
            println!("signal thread: requesting shutdown");
            handle.request_shutdown(sup_pid);
        });
        sup.join()
            .expect("supervisor returns normally after ordered shutdown");
        println!("supervisor down; root returns");
    });
}