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