Skip to main content

net_io/
control.rs

1//! The control line: a process-wide signal every worker reads with one atomic load.
2
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::{Mutex, PoisonError};
5use std::time::Instant;
6
7use crate::info::Lane;
8
9/// A process-wide signal (shed, lockdown, drain, stop) that net-io's workers read
10/// at the top of each loop, before each accept batch, and before each admission.
11///
12/// The epoch changes on every state change; the state is read only when it moved.
13pub trait ControlLine: Send + Sync + 'static {
14    /// The current epoch (an `Acquire` load).
15    fn epoch(&self) -> u64;
16    /// The current state; read only when the epoch moved.
17    fn state(&self) -> ControlState;
18    /// The deadline for `Draining`, if the line carries one.
19    fn deadline(&self) -> Option<Instant>;
20}
21
22/// What the control line asks for.
23///
24/// ```
25/// use net_io::{ControlState, Lane};
26/// let shed = ControlState::Shed { from_lane: Lane(2) };
27/// assert_ne!(shed, ControlState::Normal);
28/// ```
29#[non_exhaustive]
30#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
31pub enum ControlState {
32    /// No reaction.
33    #[default]
34    Normal,
35    /// Refuse new connections in `from_lane` and every less important lane.
36    Shed {
37        /// The most important lane that is refused.
38        from_lane: Lane,
39    },
40    /// Refuse new connections on listeners not marked `lockdown_allowed`.
41    Lockdown,
42    /// Close the door and drain.
43    Draining,
44    /// Kill every remaining connection.
45    Stopping,
46}
47
48/// A control line that never changes: `epoch()` is always 0.
49///
50/// ```
51/// use net_io::{ControlLine, NoControl};
52/// assert_eq!(NoControl.epoch(), 0);
53/// ```
54#[derive(Clone, Copy, Debug, Default)]
55pub struct NoControl;
56
57impl ControlLine for NoControl {
58    #[inline]
59    fn epoch(&self) -> u64 {
60        0
61    }
62    fn state(&self) -> ControlState {
63        ControlState::Normal
64    }
65    fn deadline(&self) -> Option<Instant> {
66        None
67    }
68}
69
70/// A control line for programs without their own.
71///
72/// ```
73/// use net_io::{ControlLine, ControlState, SimpleControl};
74/// let line = SimpleControl::new();
75/// let before = line.epoch();
76/// line.set(ControlState::Lockdown, None);
77/// assert!(line.epoch() > before);
78/// assert_eq!(line.state(), ControlState::Lockdown);
79/// ```
80#[derive(Debug, Default)]
81pub struct SimpleControl {
82    epoch: AtomicU64,
83    inner: Mutex<(ControlState, Option<Instant>)>,
84}
85
86impl SimpleControl {
87    /// A line in `Normal`.
88    #[must_use]
89    pub fn new() -> SimpleControl {
90        SimpleControl::default()
91    }
92
93    /// Changes the state and bumps the epoch.
94    pub fn set(&self, state: ControlState, deadline: Option<Instant>) {
95        *self.inner.lock().unwrap_or_else(PoisonError::into_inner) = (state, deadline);
96        self.epoch.fetch_add(1, Ordering::Release);
97    }
98}
99
100impl ControlLine for SimpleControl {
101    #[inline]
102    fn epoch(&self) -> u64 {
103        self.epoch.load(Ordering::Acquire)
104    }
105    fn state(&self) -> ControlState {
106        self.inner.lock().unwrap_or_else(PoisonError::into_inner).0
107    }
108    fn deadline(&self) -> Option<Instant> {
109        self.inner.lock().unwrap_or_else(PoisonError::into_inner).1
110    }
111}
112
113/// A worker's cached view of the line: one load per check, a state read only on change.
114#[derive(Debug)]
115pub(crate) struct ControlView {
116    epoch: u64,
117    state: ControlState,
118    deadline: Option<Instant>,
119}
120
121impl ControlView {
122    pub(crate) fn new() -> ControlView {
123        ControlView {
124            epoch: 0,
125            state: ControlState::Normal,
126            deadline: None,
127        }
128    }
129
130    /// Refreshes from `line`; returns `true` when the state changed.
131    #[inline]
132    pub(crate) fn refresh(&mut self, line: &dyn ControlLine) -> bool {
133        let epoch = line.epoch();
134        if epoch == self.epoch {
135            return false;
136        }
137        self.epoch = epoch;
138        let state = line.state();
139        self.deadline = line.deadline();
140        let changed = state != self.state;
141        self.state = state;
142        changed
143    }
144
145    pub(crate) fn state(&self) -> ControlState {
146        self.state
147    }
148
149    pub(crate) fn deadline(&self) -> Option<Instant> {
150        self.deadline
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn test_view_sees_only_final_state() {
160        let line = SimpleControl::new();
161        let mut view = ControlView::new();
162        assert!(!view.refresh(&line));
163        line.set(ControlState::Shed { from_lane: Lane(1) }, None);
164        line.set(ControlState::Normal, None);
165        // The epoch moved but the final state equals the cached one.
166        assert!(!view.refresh(&line));
167        assert_eq!(view.state(), ControlState::Normal);
168        line.set(ControlState::Draining, None);
169        assert!(view.refresh(&line));
170        assert_eq!(view.state(), ControlState::Draining);
171    }
172
173    #[test]
174    fn test_no_control_never_changes() {
175        let mut view = ControlView::new();
176        assert!(!view.refresh(&NoControl));
177    }
178}