shared_state/
ready.rs

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

shared_state/
ready.rs
