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, GenServerName, ShutdownAction, StopHandle,
    TimerHandle,
};
use smarm::supervisor::{ChildSpec, OneForOne, Restart, Shutdown};
use smarm::{sleep, spawn};
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");
    }
}

/// The server's name: how the rest of the app reaches it (and the only handle
/// that survives a restart).
const DRAINER: GenServerName<Drainer> = GenServerName::new("drainer");

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))),
        )
        // A gen_server is a direct child: `named(N).run()` runs the loop as
        // the child actor itself, so the supervisor's shutdown arrives as
        // `handle_shutdown` and a restart re-binds the name.
        .child(
            ChildSpec::new(Restart::Permanent, || {
                GenServerBuilder::new(Drainer {
                    pending: 3,
                    stop: None,
                    timer: None,
                })
                .named(DRAINER)
                .run()
                .expect("drainer name is free");
            })
            .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");
    });
}