rustfs_common/
readiness.rs

1// Copyright 2024 RustFS Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::atomic::{AtomicU8, Ordering};
16
17/// Represents the various stages of system startup
18#[repr(u8)]
19pub enum SystemStage {
20    Booting = 0,
21    StorageReady = 1, // Disks online, Quorum met
22    IamReady = 2,     // Users and Policies loaded into cache
23    FullReady = 3,    // System ready to serve all traffic
24}
25
26/// Global readiness tracker for the service
27/// This struct uses atomic operations to track the readiness status of various components
28/// of the service in a thread-safe manner.
29pub struct GlobalReadiness {
30    status: AtomicU8,
31}
32
33impl Default for GlobalReadiness {
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39impl GlobalReadiness {
40    /// Create a new GlobalReadiness instance with initial status as Starting
41    /// # Returns
42    /// A new instance of GlobalReadiness
43    pub fn new() -> Self {
44        Self {
45            status: AtomicU8::new(SystemStage::Booting as u8),
46        }
47    }
48
49    /// Update the system to a new stage
50    ///
51    /// # Arguments
52    /// * `step` - The SystemStage step to mark as ready
53    pub fn mark_stage(&self, step: SystemStage) {
54        self.status.fetch_max(step as u8, Ordering::SeqCst);
55    }
56
57    /// Check if the service is fully ready
58    /// # Returns
59    /// `true` if the service is fully ready, `false` otherwise
60    pub fn is_ready(&self) -> bool {
61        self.status.load(Ordering::SeqCst) == SystemStage::FullReady as u8
62    }
63
64    /// Get the current system stage
65    /// # Returns
66    /// The current SystemStage of the service
67    pub fn current_stage(&self) -> SystemStage {
68        match self.status.load(Ordering::SeqCst) {
69            0 => SystemStage::Booting,
70            1 => SystemStage::StorageReady,
71            2 => SystemStage::IamReady,
72            3 => SystemStage::FullReady,
73            invalid => {
74                debug_assert!(false, "GlobalReadiness::current_stage: invalid status value {}", invalid);
75                // Fallback to the most conservative stage on invalid values
76                SystemStage::Booting
77            }
78        }
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use std::sync::Arc;
86    use std::thread;
87
88    #[test]
89    fn test_initial_state() {
90        let readiness = GlobalReadiness::new();
91        assert!(!readiness.is_ready());
92        assert_eq!(readiness.status.load(Ordering::SeqCst), SystemStage::Booting as u8);
93    }
94
95    #[test]
96    fn test_mark_stage_progression() {
97        let readiness = GlobalReadiness::new();
98        readiness.mark_stage(SystemStage::StorageReady);
99        assert!(!readiness.is_ready());
100        assert_eq!(readiness.status.load(Ordering::SeqCst), SystemStage::StorageReady as u8);
101
102        readiness.mark_stage(SystemStage::IamReady);
103        assert!(!readiness.is_ready());
104        assert_eq!(readiness.status.load(Ordering::SeqCst), SystemStage::IamReady as u8);
105
106        readiness.mark_stage(SystemStage::FullReady);
107        assert!(readiness.is_ready());
108    }
109
110    #[test]
111    fn test_no_regression() {
112        let readiness = GlobalReadiness::new();
113        readiness.mark_stage(SystemStage::FullReady);
114        readiness.mark_stage(SystemStage::IamReady); // Should not regress
115        assert!(readiness.is_ready());
116    }
117
118    #[test]
119    fn test_concurrent_marking() {
120        let readiness = Arc::new(GlobalReadiness::new());
121        let mut handles = vec![];
122
123        for _ in 0..10 {
124            let r = Arc::clone(&readiness);
125            handles.push(thread::spawn(move || {
126                r.mark_stage(SystemStage::StorageReady);
127                r.mark_stage(SystemStage::IamReady);
128                r.mark_stage(SystemStage::FullReady);
129            }));
130        }
131
132        for h in handles {
133            h.join().unwrap();
134        }
135
136        assert!(readiness.is_ready());
137    }
138
139    #[test]
140    fn test_is_ready_only_at_full_ready() {
141        let readiness = GlobalReadiness::new();
142        assert!(!readiness.is_ready());
143
144        readiness.mark_stage(SystemStage::StorageReady);
145        assert!(!readiness.is_ready());
146
147        readiness.mark_stage(SystemStage::IamReady);
148        assert!(!readiness.is_ready());
149
150        readiness.mark_stage(SystemStage::FullReady);
151        assert!(readiness.is_ready());
152    }
153}
rustfs_common/readiness.rs

rustfs_common/
readiness.rs
