xref: /DragonOS/kernel/src/libs/once.rs (revision 7a29d4fcbcd89a226289c7bf541c2c78623de3ad)
1 use core::{
2     fmt::{self, Debug, Formatter},
3     sync::atomic::Ordering,
4 };
5 
6 use atomic_enum::atomic_enum;
7 
8 pub struct Once {
9     inner: AtomicOnceState,
10 }
11 
12 #[atomic_enum]
13 #[derive(PartialEq, Eq)]
14 pub enum OnceState {
15     Incomplete,
16     Posioned,
17     Complete,
18 }
19 
20 #[allow(dead_code)]
21 impl Once {
22     pub const fn new() -> Self {
23         Self {
24             inner: AtomicOnceState::new(OnceState::Incomplete),
25         }
26     }
27 
28     #[track_caller]
29     pub fn call_once<F: FnOnce()>(&self, f: F) {
30         if self.is_completed() {
31             return;
32         }
33 
34         // set initialized
35         let r = self.inner.compare_exchange(
36             OnceState::Incomplete,
37             OnceState::Posioned,
38             Ordering::SeqCst,
39             Ordering::SeqCst,
40         );
41         if r.is_err() {
42             return;
43         }
44         // call function
45         f();
46         // set completed
47         self.inner.store(OnceState::Complete, Ordering::SeqCst);
48     }
49 
50     /// Performs the same function as [`call_once()`] except ignores poisoning.
51     ///
52     /// Unlike [`call_once()`], if this [`Once`] has been poisoned (i.e., a previous
53     /// call to [`call_once()`] or [`call_once_force()`] caused a panic), calling
54     /// [`call_once_force()`] will still invoke the closure `f` and will _not_
55     /// result in an immediate panic. If `f` panics, the [`Once`] will remain
56     /// in a poison state. If `f` does _not_ panic, the [`Once`] will no
57     /// longer be in a poison state and all future calls to [`call_once()`] or
58     /// [`call_once_force()`] will be no-ops.
59     ///
60     /// The closure `f` is yielded a [`OnceState`] structure which can be used
61     /// to query the poison status of the [`Once`].
62     ///
63     /// [`call_once()`]: Once::call_once
64     /// [`call_once_force()`]: Once::call_once_force
65     ///
66     /// # Examples
67     ///
68     /// ```
69     /// use std::sync::Once;
70     /// use std::thread;
71     ///
72     /// static INIT: Once = Once::new();
73     ///
74     /// // poison the once
75     /// let handle = thread::spawn(|| {
76     ///     INIT.call_once(|| panic!());
77     /// });
78     /// assert!(handle.join().is_err());
79     ///
80     /// // poisoning propagates
81     /// let handle = thread::spawn(|| {
82     ///     INIT.call_once(|| {});
83     /// });
84     /// assert!(handle.join().is_err());
85     ///
86     /// // call_once_force will still run and reset the poisoned state
87     /// INIT.call_once_force(|state| {
88     ///     assert!(state.is_poisoned());
89     /// });
90     ///
91     /// // once any success happens, we stop propagating the poison
92     /// INIT.call_once(|| {});
93     /// ```
94     pub fn call_once_force<F>(&self, f: F)
95     where
96         F: FnOnce(&OnceState),
97     {
98         // fast path check
99         if self.is_completed() {
100             return;
101         }
102 
103         // set poisoned
104         self.inner
105             .compare_exchange(
106                 OnceState::Incomplete,
107                 OnceState::Posioned,
108                 Ordering::SeqCst,
109                 Ordering::SeqCst,
110             )
111             .ok();
112 
113         // call function
114         f(&self.inner.load(Ordering::SeqCst));
115 
116         // set initialized
117         self.inner.store(OnceState::Complete, Ordering::SeqCst);
118     }
119 
120     /// Fast path check
121     #[inline]
122     pub fn is_completed(&self) -> bool {
123         self.inner.load(Ordering::SeqCst) == OnceState::Complete
124     }
125 
126     /// Returns the current state of the `Once` instance.
127     #[inline]
128     pub fn state(&self) -> OnceState {
129         self.inner.load(Ordering::SeqCst)
130     }
131 }
132 
133 impl Debug for Once {
134     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
135         f.debug_struct("Once").finish_non_exhaustive()
136     }
137 }
138 
139 #[allow(dead_code)]
140 impl OnceState {
141     /// Returns `true` if the associated [`Once`] was poisoned prior to the
142     /// invocation of the closure passed to [`Once::call_once_force()`].
143     ///
144     /// # Examples
145     ///
146     /// A poisoned [`Once`]:
147     ///
148     /// ```
149     /// use std::sync::Once;
150     /// use std::thread;
151     ///
152     /// static INIT: Once = Once::new();
153     ///
154     /// // poison the once
155     /// let handle = thread::spawn(|| {
156     ///     INIT.call_once(|| panic!());
157     /// });
158     /// assert!(handle.join().is_err());
159     ///
160     /// INIT.call_once_force(|state| {
161     ///     assert!(state.is_poisoned());
162     /// });
163     /// ```
164     ///
165     /// An unpoisoned [`Once`]:
166     ///
167     /// ```
168     /// use std::sync::Once;
169     ///
170     /// static INIT: Once = Once::new();
171     ///
172     /// INIT.call_once_force(|state| {
173     ///     assert!(!state.is_poisoned());
174     /// });
175     #[inline]
176     pub fn is_poisoned(&self) -> bool {
177         *self == OnceState::Posioned
178     }
179 }
180