shared_state/
ready.rs

1//! Readiness gate for service health probing.
2//!
3//! The atomic type is swapped for loom's equivalent under `cfg(loom)`
4//! (see `loom_tests`) so the Release/Acquire discipline is model-checked.
5
6#[cfg(loom)]
7use loom::sync::atomic::{AtomicBool, Ordering};
8#[cfg(not(loom))]
9use std::sync::atomic::{AtomicBool, Ordering};
10use std::time::Duration;
11
12/// A thread-safe readiness gate.
13///
14/// Services use this to signal when they are ready to accept traffic.
15/// Kubernetes liveness/readiness probes poll `is_ready()`.
16pub struct ReadyGate {
17    ready: AtomicBool,
18}
19
20impl Default for ReadyGate {
21    fn default() -> Self {
22        Self::new()
23    }
24}
25
26impl ReadyGate {
27    /// Create a new gate in the NOT-ready state.
28    pub fn new() -> Self {
29        Self {
30            ready: AtomicBool::new(false),
31        }
32    }
33
34    /// Mark the service as ready.
35    pub fn set_ready(&self) {
36        self.ready.store(true, Ordering::Release);
37    }
38
39    /// Mark the service as NOT ready.
40    pub fn set_not_ready(&self) {
41        self.ready.store(false, Ordering::Release);
42    }
43
44    /// Check if the service is ready (non-blocking).
45    pub fn is_ready(&self) -> bool {
46        self.ready.load(Ordering::Acquire)
47    }
48
49    /// Wait until the service is ready, polling every `interval`.
50    pub async fn wait_until_ready(&self) {
51        loop {
52            if self.is_ready() {
53                return;
54            }
55            tokio::time::sleep(Duration::from_millis(100)).await;
56        }
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63
64    #[tokio::test]
65    async fn starts_not_ready() {
66        let gate = ReadyGate::new();
67        assert!(!gate.is_ready());
68    }
69
70    #[tokio::test]
71    async fn becomes_ready() {
72        let gate = ReadyGate::new();
73        gate.set_ready();
74        assert!(gate.is_ready());
75    }
76
77    #[tokio::test]
78    async fn can_become_not_ready() {
79        let gate = ReadyGate::new();
80        gate.set_ready();
81        gate.set_not_ready();
82        assert!(!gate.is_ready());
83    }
84
85    #[tokio::test]
86    async fn wait_until_ready() {
87        let gate = ReadyGate::new();
88        let gate_clone = std::sync::Arc::new(gate);
89        let g = gate_clone.clone();
90        tokio::spawn(async move {
91            tokio::time::sleep(Duration::from_millis(50)).await;
92            g.set_ready();
93        });
94        gate_clone.wait_until_ready().await;
95        assert!(gate_clone.is_ready());
96    }
97}
shared_state/ready.rs

shared_state/
ready.rs
