Skip to main content

autumn_plugin_pingora/
lifecycle.rs

1//! The proxy lifecycle as a finite state machine.
2//!
3//! ```mermaid
4//! stateDiagram-v2
5//!     [*] --> Idle
6//!     Idle --> Serving: Bound
7//!     Idle --> Failed: StartFailed
8//!     Idle --> Stopped: ShutdownRequested
9//!     Serving --> Draining: ShutdownRequested
10//!     Serving --> Failed: ServerExited
11//!     Draining --> Stopped: Drained / ServerExited
12//! ```
13//!
14//! [`Lifecycle::next`] is the only transition function. It is pure and
15//! total. An illegal event returns `None` and does not change the state.
16
17use std::fmt;
18use std::sync::atomic::{AtomicU8, Ordering};
19
20/// The state of the proxy.
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
22#[non_exhaustive]
23pub enum Lifecycle {
24    /// Not started yet.
25    Idle,
26    /// The listener is bound and accepts connections.
27    Serving,
28    /// Shutdown started. No new connections. Open requests finish.
29    Draining,
30    /// The proxy stopped. Terminal.
31    Stopped,
32    /// The proxy could not start, or stopped unexpectedly. Terminal.
33    Failed,
34}
35
36/// An input to the state machine.
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
38#[non_exhaustive]
39pub enum LifecycleEvent {
40    /// The listener is bound.
41    Bound,
42    /// The proxy could not start.
43    StartFailed,
44    /// The app asked the proxy to stop.
45    ShutdownRequested,
46    /// All connections ended, or the grace period expired.
47    Drained,
48    /// The accept loop ended.
49    ServerExited,
50}
51
52impl Lifecycle {
53    /// All states.
54    pub const ALL: [Self; 5] = [
55        Self::Idle,
56        Self::Serving,
57        Self::Draining,
58        Self::Stopped,
59        Self::Failed,
60    ];
61
62    /// The next state after `event`, or `None` if `event` is illegal here.
63    #[must_use]
64    pub const fn next(self, event: LifecycleEvent) -> Option<Self> {
65        use LifecycleEvent::{Bound, Drained, ServerExited, ShutdownRequested, StartFailed};
66        match (self, event) {
67            (Self::Idle, Bound) => Some(Self::Serving),
68            (Self::Serving, ShutdownRequested) => Some(Self::Draining),
69            (Self::Idle, StartFailed) | (Self::Serving, ServerExited) => Some(Self::Failed),
70            (Self::Idle, ShutdownRequested) | (Self::Draining, Drained | ServerExited) => {
71                Some(Self::Stopped)
72            }
73            _ => None,
74        }
75    }
76
77    /// `true` only in [`Serving`](Self::Serving).
78    #[must_use]
79    pub const fn is_ready(self) -> bool {
80        matches!(self, Self::Serving)
81    }
82
83    /// `true` in [`Stopped`](Self::Stopped) and [`Failed`](Self::Failed).
84    #[must_use]
85    pub const fn is_terminal(self) -> bool {
86        matches!(self, Self::Stopped | Self::Failed)
87    }
88
89    /// Lower-case name.
90    #[must_use]
91    pub const fn as_str(self) -> &'static str {
92        match self {
93            Self::Idle => "idle",
94            Self::Serving => "serving",
95            Self::Draining => "draining",
96            Self::Stopped => "stopped",
97            Self::Failed => "failed",
98        }
99    }
100
101    const fn to_u8(self) -> u8 {
102        match self {
103            Self::Idle => 0,
104            Self::Serving => 1,
105            Self::Draining => 2,
106            Self::Stopped => 3,
107            Self::Failed => 4,
108        }
109    }
110
111    const fn from_u8(value: u8) -> Self {
112        match value {
113            0 => Self::Idle,
114            1 => Self::Serving,
115            2 => Self::Draining,
116            3 => Self::Stopped,
117            _ => Self::Failed,
118        }
119    }
120}
121
122impl LifecycleEvent {
123    /// All events.
124    pub const ALL: [Self; 5] = [
125        Self::Bound,
126        Self::StartFailed,
127        Self::ShutdownRequested,
128        Self::Drained,
129        Self::ServerExited,
130    ];
131}
132
133impl fmt::Display for Lifecycle {
134    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
135        f.write_str(self.as_str())
136    }
137}
138
139/// A shared, atomic [`Lifecycle`]. The only way to change the state.
140#[derive(Debug)]
141pub struct LifecycleCell(AtomicU8);
142
143impl Default for LifecycleCell {
144    fn default() -> Self {
145        Self::new()
146    }
147}
148
149impl LifecycleCell {
150    /// A cell in [`Lifecycle::Idle`].
151    #[must_use]
152    pub const fn new() -> Self {
153        Self(AtomicU8::new(Lifecycle::Idle.to_u8()))
154    }
155
156    /// The current state.
157    #[must_use]
158    pub fn get(&self) -> Lifecycle {
159        Lifecycle::from_u8(self.0.load(Ordering::Acquire))
160    }
161
162    /// Apply `event` atomically.
163    ///
164    /// # Errors
165    ///
166    /// Returns the unchanged state when `event` is illegal.
167    pub fn apply(&self, event: LifecycleEvent) -> Result<Lifecycle, Lifecycle> {
168        let mut current = self.get();
169        loop {
170            let Some(next) = current.next(event) else {
171                return Err(current);
172            };
173            match self.0.compare_exchange(
174                current.to_u8(),
175                next.to_u8(),
176                Ordering::AcqRel,
177                Ordering::Acquire,
178            ) {
179                Ok(_) => return Ok(next),
180                Err(actual) => current = Lifecycle::from_u8(actual),
181            }
182        }
183    }
184}