veilid_tools/
startup_lock.rs

1use super::*;
2
3#[derive(ThisError, Debug, Copy, Clone, PartialEq, Eq)]
4#[error("Already started")]
5pub struct StartupLockAlreadyStartedError;
6
7#[derive(ThisError, Debug, Copy, Clone, PartialEq, Eq)]
8#[error("Already shut down")]
9pub struct StartupLockAlreadyShutDownError;
10
11#[derive(ThisError, Debug, Copy, Clone, PartialEq, Eq)]
12#[error("Not started")]
13pub struct StartupLockNotStartedError;
14
15/// RAII-style lock for startup and shutdown operations
16/// Must call 'success()' on this lock to report a successful startup or shutdown
17/// Dropping this lock without calling 'success()' first indicates a failed
18/// startup or shutdown operation
19#[derive(Debug)]
20pub struct StartupLockGuard<'a> {
21    guard: AsyncRwLockWriteGuard<'a, bool>,
22    success_value: bool,
23}
24
25impl<'a> StartupLockGuard<'a> {
26    /// Call this function at the end of a successful startup or shutdown
27    /// operation to switch the state of the StartupLock.
28    pub fn success(mut self) {
29        *self.guard = self.success_value;
30    }
31}
32
33/// RAII-style lock for entry operations on a started-up region of code.
34#[derive(Debug)]
35pub struct StartupLockEnterGuard<'a> {
36    _guard: AsyncRwLockReadGuard<'a, bool>,
37    #[cfg(feature = "debug-locks")]
38    id: usize,
39    #[cfg(feature = "debug-locks")]
40    active_guards: Arc<Mutex<HashMap<usize, backtrace::Backtrace>>>,
41}
42
43#[cfg(feature = "debug-locks")]
44impl<'a> Drop for StartupLockEnterGuard<'a> {
45    fn drop(&mut self) {
46        self.active_guards.lock().remove(&self.id);
47    }
48}
49
50/// RAII-style lock for entry operations on a started-up region of code.
51#[derive(Debug)]
52pub struct StartupLockEnterGuardArc {
53    _guard: AsyncRwLockReadGuardArc<bool>,
54    #[cfg(feature = "debug-locks")]
55    id: usize,
56    #[cfg(feature = "debug-locks")]
57    active_guards: Arc<Mutex<HashMap<usize, backtrace::Backtrace>>>,
58}
59
60#[cfg(feature = "debug-locks")]
61impl Drop for StartupLockEnterGuardArc {
62    fn drop(&mut self) {
63        self.active_guards.lock().remove(&self.id);
64    }
65}
66
67#[cfg(feature = "debug-locks")]
68static GUARD_ID: AtomicUsize = AtomicUsize::new(0);
69
70/// Synchronization mechanism that tracks the startup and shutdown of a region of code.
71/// Guarantees that some code can only be started up once and shut down only if it is
72/// already started.
73/// Also tracks if the code is in-use and will wait for all 'entered' code to finish
74/// before shutting down. Once a shutdown is requested, future calls to 'enter' will
75/// fail, ensuring that nothing is 'entered' at the time of shutdown. This allows an
76/// asynchronous shutdown to wait for operations to finish before proceeding.
77#[derive(Debug)]
78pub struct StartupLock {
79    startup_state: Arc<AsyncRwLock<bool>>,
80    stop_source: Mutex<Option<StopSource>>,
81    #[cfg(feature = "debug-locks")]
82    active_guards: Arc<Mutex<HashMap<usize, backtrace::Backtrace>>>,
83}
84
85impl StartupLock {
86    pub fn new() -> Self {
87        Self {
88            startup_state: Arc::new(AsyncRwLock::new(false)),
89            stop_source: Mutex::new(None),
90            #[cfg(feature = "debug-locks")]
91            active_guards: Arc::new(Mutex::new(HashMap::new())),
92        }
93    }
94
95    /// Start up if things are not already started up
96    /// One must call 'success()' on the returned startup lock guard if startup was successful
97    /// otherwise the startup lock will not shift to the 'started' state.
98    pub fn startup(&self) -> Result<StartupLockGuard, StartupLockAlreadyStartedError> {
99        let guard =
100            asyncrwlock_try_write!(self.startup_state).ok_or(StartupLockAlreadyStartedError)?;
101        if *guard {
102            return Err(StartupLockAlreadyStartedError);
103        }
104        *self.stop_source.lock() = Some(StopSource::new());
105
106        Ok(StartupLockGuard {
107            guard,
108            success_value: true,
109        })
110    }
111
112    /// Get a stop token for this lock
113    /// One can wait on this to timeout operations when a shutdown is requested
114    pub fn stop_token(&self) -> Option<StopToken> {
115        self.stop_source.lock().as_ref().map(|ss| ss.token())
116    }
117
118    /// Check if this StartupLock is currently in a started state
119    /// Returns false is the state is in transition
120    pub fn is_started(&self) -> bool {
121        let Some(guard) = asyncrwlock_try_read!(self.startup_state) else {
122            return false;
123        };
124        *guard
125    }
126
127    /// Check if this StartupLock is currently in a shut down state
128    /// Returns false is the state is in transition
129    pub fn is_shut_down(&self) -> bool {
130        let Some(guard) = asyncrwlock_try_read!(self.startup_state) else {
131            return false;
132        };
133        !*guard
134    }
135
136    /// Wait for all 'entered' operations to finish before shutting down
137    /// One must call 'success()' on the returned startup lock guard if shutdown was successful
138    /// otherwise the startup lock will not shift to the 'stopped' state.
139    pub async fn shutdown(&self) -> Result<StartupLockGuard, StartupLockAlreadyShutDownError> {
140        // Drop the stop source to ensure we can detect shutdown has been requested
141        *self.stop_source.lock() = None;
142
143        cfg_if! {
144            if #[cfg(feature = "debug-locks")] {
145                let guard = match timeout(30000, self.startup_state.write()).await {
146                    Ok(v) => v,
147                    Err(_) => {
148                        eprintln!("active guards: {:#?}", self.active_guards.lock().values().collect::<Vec<_>>());
149                        panic!("shutdown deadlock");
150                    }
151                };
152            } else {
153                let guard = self.startup_state.write().await;
154            }
155        }
156        if !*guard {
157            return Err(StartupLockAlreadyShutDownError);
158        }
159        Ok(StartupLockGuard {
160            guard,
161            success_value: false,
162        })
163    }
164
165    /// Enter an operation in a started-up module.
166    /// If this module has not yet started up or is in the process of startup or shutdown
167    /// this will fail.
168    pub fn enter(&self) -> Result<StartupLockEnterGuard, StartupLockNotStartedError> {
169        let guard = asyncrwlock_try_read!(self.startup_state).ok_or(StartupLockNotStartedError)?;
170        if !*guard {
171            return Err(StartupLockNotStartedError);
172        }
173        let out = StartupLockEnterGuard {
174            _guard: guard,
175            #[cfg(feature = "debug-locks")]
176            id: GUARD_ID.fetch_add(1, Ordering::AcqRel),
177            #[cfg(feature = "debug-locks")]
178            active_guards: self.active_guards.clone(),
179        };
180
181        #[cfg(feature = "debug-locks")]
182        self.active_guards
183            .lock()
184            .insert(out.id, backtrace::Backtrace::new());
185
186        Ok(out)
187    }
188
189    /// Enter an operation in a started-up module, using an owned lock.
190    /// If this module has not yet started up or is in the process of startup or shutdown
191    /// this will fail.
192    pub fn enter_arc(&self) -> Result<StartupLockEnterGuardArc, StartupLockNotStartedError> {
193        let guard =
194            asyncrwlock_try_read_arc!(self.startup_state).ok_or(StartupLockNotStartedError)?;
195        if !*guard {
196            return Err(StartupLockNotStartedError);
197        }
198        let out = StartupLockEnterGuardArc {
199            _guard: guard,
200            #[cfg(feature = "debug-locks")]
201            id: GUARD_ID.fetch_add(1, Ordering::AcqRel),
202            #[cfg(feature = "debug-locks")]
203            active_guards: self.active_guards.clone(),
204        };
205
206        #[cfg(feature = "debug-locks")]
207        self.active_guards
208            .lock()
209            .insert(out.id, backtrace::Backtrace::new());
210
211        Ok(out)
212    }
213}
214
215impl Default for StartupLock {
216    fn default() -> Self {
217        Self::new()
218    }
219}
veilid_tools/startup_lock.rs

veilid_tools/
startup_lock.rs
