Skip to main content

malkuth/
hooks.rs

1//! Pluggable lifecycle hooks (requirement: "exit probes / heartbeats are
2//! optional, hookable — use the default, or supply your own").
3//!
4//! Every default facility in the `malkuth` crate (signal-based exit, HTTP
5//! probes, timed heartbeats) is wired through one of these traits. If the
6//! default does not fit — e.g. you want drain to be triggered when your server
7//! receives an application-level "stop" command, or you want heartbeats to
8//! carry your own payload — implement the trait and hand it to the supervisor.
9//! Nothing here depends on a runtime or a wire framework.
10
11use std::time::Duration;
12
13use async_trait::async_trait;
14
15use crate::{DrainController, HealthStatus, HeartbeatBeat, ReadyStatus, ShutdownKind};
16
17/// Why an [`ExitSource`] fired.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub struct ExitReason {
20    /// The kind of shutdown implied.
21    pub kind: ShutdownKind,
22    /// Whether the process should actually exit once drain completes
23    /// (`false` for a pure reload that keeps the process alive).
24    pub should_exit: bool,
25}
26
27impl ExitReason {
28    /// A graceful drain that should terminate the process.
29    pub const fn graceful() -> Self {
30        Self {
31            kind: ShutdownKind::Graceful,
32            should_exit: true,
33        }
34    }
35    /// An immediate (skip-drain) exit.
36    pub const fn immediate() -> Self {
37        Self {
38            kind: ShutdownKind::Immediate,
39            should_exit: true,
40        }
41    }
42    /// A reload: drain bit stays clear, process keeps running.
43    pub const fn reload() -> Self {
44        Self {
45            kind: ShutdownKind::Reload,
46            should_exit: false,
47        }
48    }
49}
50
51/// Source of process-exit / drain triggers.
52///
53/// The default implementation (in the `malkuth` crate) installs OS signal
54/// handlers (`SIGTERM`/`SIGINT`/`SIGHUP`/`SIGQUIT`). A custom implementation
55/// might instead trigger drain when the server receives an in-band "stop"
56/// command, when a parent supervisor sends a control message over IPC, or when
57/// an orchestrator flips a file bit.
58///
59/// Implementations **must** call `ctrl.begin_drain(reason.kind)` (and arrange
60/// process exit if `reason.should_exit`) when their condition fires.
61#[async_trait]
62pub trait ExitSource: Send + Sync {
63    /// Block until the source fires, then describe why.
64    ///
65    /// Returning is the signal that drain should begin; the caller then drives
66    /// [`DrainController::wait_for_drain`] and exits if asked.
67    async fn wait(&self, ctrl: DrainController) -> ExitReason;
68}
69
70/// Readiness/liveness state as observed by probes (HTTP `/readyz` or RPC
71/// `Lifecycle.Status`). Implementations decide *how* the state is computed
72/// (read atomics, ping a dependency, query a registry, …).
73#[async_trait]
74pub trait ProbeSink: Send + Sync {
75    /// Current readiness (drain bit + dependency checks).
76    async fn ready(&self) -> ReadyStatus;
77    /// Current liveness.
78    async fn health(&self) -> HealthStatus;
79}
80
81/// Heartbeat report produced by a [`Heartbeat`] source.
82#[derive(Debug, Clone, Default)]
83pub struct HeartbeatReport {
84    /// The beat to publish, if any.
85    pub beat: Option<HeartbeatBeat>,
86    /// Interval until the next beat should be sampled.
87    pub next_after: Duration,
88}
89
90/// A periodic liveness heartbeat. The default implementation emits a beat on a
91/// fixed cadence; a custom one can embed domain-specific state, rate-limit, or
92/// suppress beats conditionally.
93#[async_trait]
94pub trait Heartbeat: Send + Sync {
95    /// Produce the next heartbeat report (blocking until it is due).
96    async fn next(&self) -> HeartbeatReport;
97}
98
99/// A drain hook: run arbitrary graceful-shutdown work when drain begins.
100///
101/// Wired in by the supervisor between "stop accepting new work" and "exit".
102/// Examples: close WebSocket frames, flush buffers, disconnect upstream pools,
103/// release locks. This is the seam where each application injects its own
104/// drain closure (the design doc §5.3 step 3–6).
105#[async_trait]
106pub trait DrainHook: Send + Sync {
107    /// Perform drain work, completing within `budget` where possible.
108    async fn drain(&self, budget: Duration);
109}