1 // Copyright 2016 Amanieu d'Antras 2 // 3 // Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or 4 // http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or 5 // http://opensource.org/licenses/MIT>, at your option. This file may not be 6 // copied, modified, or distributed except according to those terms. 7 8 use crate::util::UncheckedOptionExt; 9 use core::{ 10 fmt, mem, 11 sync::atomic::{fence, AtomicU8, Ordering}, 12 }; 13 use parking_lot_core::{self, SpinWait, DEFAULT_PARK_TOKEN, DEFAULT_UNPARK_TOKEN}; 14 15 const DONE_BIT: u8 = 1; 16 const POISON_BIT: u8 = 2; 17 const LOCKED_BIT: u8 = 4; 18 const PARKED_BIT: u8 = 8; 19 20 /// Current state of a `Once`. 21 #[derive(Copy, Clone, Eq, PartialEq, Debug)] 22 pub enum OnceState { 23 /// A closure has not been executed yet 24 New, 25 26 /// A closure was executed but panicked. 27 Poisoned, 28 29 /// A thread is currently executing a closure. 30 InProgress, 31 32 /// A closure has completed successfully. 33 Done, 34 } 35 36 impl OnceState { 37 /// Returns whether the associated `Once` has been poisoned. 38 /// 39 /// Once an initialization routine for a `Once` has panicked it will forever 40 /// indicate to future forced initialization routines that it is poisoned. 41 #[inline] poisoned(self) -> bool42 pub fn poisoned(self) -> bool { 43 match self { 44 OnceState::Poisoned => true, 45 _ => false, 46 } 47 } 48 49 /// Returns whether the associated `Once` has successfully executed a 50 /// closure. 51 #[inline] done(self) -> bool52 pub fn done(self) -> bool { 53 match self { 54 OnceState::Done => true, 55 _ => false, 56 } 57 } 58 } 59 60 /// A synchronization primitive which can be used to run a one-time 61 /// initialization. Useful for one-time initialization for globals, FFI or 62 /// related functionality. 63 /// 64 /// # Differences from the standard library `Once` 65 /// 66 /// - Only requires 1 byte of space, instead of 1 word. 67 /// - Not required to be `'static`. 68 /// - Relaxed memory barriers in the fast path, which can significantly improve 69 /// performance on some architectures. 70 /// - Efficient handling of micro-contention using adaptive spinning. 71 /// 72 /// # Examples 73 /// 74 /// ``` 75 /// use parking_lot::Once; 76 /// 77 /// static START: Once = Once::new(); 78 /// 79 /// START.call_once(|| { 80 /// // run initialization here 81 /// }); 82 /// ``` 83 pub struct Once(AtomicU8); 84 85 impl Once { 86 /// Creates a new `Once` value. 87 #[inline] new() -> Once88 pub const fn new() -> Once { 89 Once(AtomicU8::new(0)) 90 } 91 92 /// Returns the current state of this `Once`. 93 #[inline] state(&self) -> OnceState94 pub fn state(&self) -> OnceState { 95 let state = self.0.load(Ordering::Acquire); 96 if state & DONE_BIT != 0 { 97 OnceState::Done 98 } else if state & LOCKED_BIT != 0 { 99 OnceState::InProgress 100 } else if state & POISON_BIT != 0 { 101 OnceState::Poisoned 102 } else { 103 OnceState::New 104 } 105 } 106 107 /// Performs an initialization routine once and only once. The given closure 108 /// will be executed if this is the first time `call_once` has been called, 109 /// and otherwise the routine will *not* be invoked. 110 /// 111 /// This method will block the calling thread if another initialization 112 /// routine is currently running. 113 /// 114 /// When this function returns, it is guaranteed that some initialization 115 /// has run and completed (it may not be the closure specified). It is also 116 /// guaranteed that any memory writes performed by the executed closure can 117 /// be reliably observed by other threads at this point (there is a 118 /// happens-before relation between the closure and code executing after the 119 /// return). 120 /// 121 /// # Examples 122 /// 123 /// ``` 124 /// use parking_lot::Once; 125 /// 126 /// static mut VAL: usize = 0; 127 /// static INIT: Once = Once::new(); 128 /// 129 /// // Accessing a `static mut` is unsafe much of the time, but if we do so 130 /// // in a synchronized fashion (e.g. write once or read all) then we're 131 /// // good to go! 132 /// // 133 /// // This function will only call `expensive_computation` once, and will 134 /// // otherwise always return the value returned from the first invocation. 135 /// fn get_cached_val() -> usize { 136 /// unsafe { 137 /// INIT.call_once(|| { 138 /// VAL = expensive_computation(); 139 /// }); 140 /// VAL 141 /// } 142 /// } 143 /// 144 /// fn expensive_computation() -> usize { 145 /// // ... 146 /// # 2 147 /// } 148 /// ``` 149 /// 150 /// # Panics 151 /// 152 /// The closure `f` will only be executed once if this is called 153 /// concurrently amongst many threads. If that closure panics, however, then 154 /// it will *poison* this `Once` instance, causing all future invocations of 155 /// `call_once` to also panic. 156 #[inline] call_once<F>(&self, f: F) where F: FnOnce(),157 pub fn call_once<F>(&self, f: F) 158 where 159 F: FnOnce(), 160 { 161 if self.0.load(Ordering::Acquire) == DONE_BIT { 162 return; 163 } 164 165 let mut f = Some(f); 166 self.call_once_slow(false, &mut |_| unsafe { f.take().unchecked_unwrap()() }); 167 } 168 169 /// Performs the same function as `call_once` except ignores poisoning. 170 /// 171 /// If this `Once` has been poisoned (some initialization panicked) then 172 /// this function will continue to attempt to call initialization functions 173 /// until one of them doesn't panic. 174 /// 175 /// The closure `f` is yielded a structure which can be used to query the 176 /// state of this `Once` (whether initialization has previously panicked or 177 /// not). 178 #[inline] call_once_force<F>(&self, f: F) where F: FnOnce(OnceState),179 pub fn call_once_force<F>(&self, f: F) 180 where 181 F: FnOnce(OnceState), 182 { 183 if self.0.load(Ordering::Acquire) == DONE_BIT { 184 return; 185 } 186 187 let mut f = Some(f); 188 self.call_once_slow(true, &mut |state| unsafe { 189 f.take().unchecked_unwrap()(state) 190 }); 191 } 192 193 // This is a non-generic function to reduce the monomorphization cost of 194 // using `call_once` (this isn't exactly a trivial or small implementation). 195 // 196 // Additionally, this is tagged with `#[cold]` as it should indeed be cold 197 // and it helps let LLVM know that calls to this function should be off the 198 // fast path. Essentially, this should help generate more straight line code 199 // in LLVM. 200 // 201 // Finally, this takes an `FnMut` instead of a `FnOnce` because there's 202 // currently no way to take an `FnOnce` and call it via virtual dispatch 203 // without some allocation overhead. 204 #[cold] call_once_slow(&self, ignore_poison: bool, f: &mut dyn FnMut(OnceState))205 fn call_once_slow(&self, ignore_poison: bool, f: &mut dyn FnMut(OnceState)) { 206 let mut spinwait = SpinWait::new(); 207 let mut state = self.0.load(Ordering::Relaxed); 208 loop { 209 // If another thread called the closure, we're done 210 if state & DONE_BIT != 0 { 211 // An acquire fence is needed here since we didn't load the 212 // state with Ordering::Acquire. 213 fence(Ordering::Acquire); 214 return; 215 } 216 217 // If the state has been poisoned and we aren't forcing, then panic 218 if state & POISON_BIT != 0 && !ignore_poison { 219 // Need the fence here as well for the same reason 220 fence(Ordering::Acquire); 221 panic!("Once instance has previously been poisoned"); 222 } 223 224 // Grab the lock if it isn't locked, even if there is a queue on it. 225 // We also clear the poison bit since we are going to try running 226 // the closure again. 227 if state & LOCKED_BIT == 0 { 228 match self.0.compare_exchange_weak( 229 state, 230 (state | LOCKED_BIT) & !POISON_BIT, 231 Ordering::Acquire, 232 Ordering::Relaxed, 233 ) { 234 Ok(_) => break, 235 Err(x) => state = x, 236 } 237 continue; 238 } 239 240 // If there is no queue, try spinning a few times 241 if state & PARKED_BIT == 0 && spinwait.spin() { 242 state = self.0.load(Ordering::Relaxed); 243 continue; 244 } 245 246 // Set the parked bit 247 if state & PARKED_BIT == 0 { 248 if let Err(x) = self.0.compare_exchange_weak( 249 state, 250 state | PARKED_BIT, 251 Ordering::Relaxed, 252 Ordering::Relaxed, 253 ) { 254 state = x; 255 continue; 256 } 257 } 258 259 // Park our thread until we are woken up by the thread that owns the 260 // lock. 261 unsafe { 262 let addr = self as *const _ as usize; 263 let validate = || self.0.load(Ordering::Relaxed) == LOCKED_BIT | PARKED_BIT; 264 let before_sleep = || {}; 265 let timed_out = |_, _| unreachable!(); 266 parking_lot_core::park( 267 addr, 268 validate, 269 before_sleep, 270 timed_out, 271 DEFAULT_PARK_TOKEN, 272 None, 273 ); 274 } 275 276 // Loop back and check if the done bit was set 277 spinwait.reset(); 278 state = self.0.load(Ordering::Relaxed); 279 } 280 281 struct PanicGuard<'a>(&'a Once); 282 impl<'a> Drop for PanicGuard<'a> { 283 fn drop(&mut self) { 284 // Mark the state as poisoned, unlock it and unpark all threads. 285 let once = self.0; 286 let state = once.0.swap(POISON_BIT, Ordering::Release); 287 if state & PARKED_BIT != 0 { 288 unsafe { 289 let addr = once as *const _ as usize; 290 parking_lot_core::unpark_all(addr, DEFAULT_UNPARK_TOKEN); 291 } 292 } 293 } 294 } 295 296 // At this point we have the lock, so run the closure. Make sure we 297 // properly clean up if the closure panicks. 298 let guard = PanicGuard(self); 299 let once_state = if state & POISON_BIT != 0 { 300 OnceState::Poisoned 301 } else { 302 OnceState::New 303 }; 304 f(once_state); 305 mem::forget(guard); 306 307 // Now unlock the state, set the done bit and unpark all threads 308 let state = self.0.swap(DONE_BIT, Ordering::Release); 309 if state & PARKED_BIT != 0 { 310 unsafe { 311 let addr = self as *const _ as usize; 312 parking_lot_core::unpark_all(addr, DEFAULT_UNPARK_TOKEN); 313 } 314 } 315 } 316 } 317 318 impl Default for Once { 319 #[inline] default() -> Once320 fn default() -> Once { 321 Once::new() 322 } 323 } 324 325 impl fmt::Debug for Once { fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result326 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 327 f.debug_struct("Once") 328 .field("state", &self.state()) 329 .finish() 330 } 331 } 332 333 #[cfg(test)] 334 mod tests { 335 use crate::Once; 336 use std::panic; 337 use std::sync::mpsc::channel; 338 use std::thread; 339 340 #[test] smoke_once()341 fn smoke_once() { 342 static O: Once = Once::new(); 343 let mut a = 0; 344 O.call_once(|| a += 1); 345 assert_eq!(a, 1); 346 O.call_once(|| a += 1); 347 assert_eq!(a, 1); 348 } 349 350 #[test] stampede_once()351 fn stampede_once() { 352 static O: Once = Once::new(); 353 static mut RUN: bool = false; 354 355 let (tx, rx) = channel(); 356 for _ in 0..10 { 357 let tx = tx.clone(); 358 thread::spawn(move || { 359 for _ in 0..4 { 360 thread::yield_now() 361 } 362 unsafe { 363 O.call_once(|| { 364 assert!(!RUN); 365 RUN = true; 366 }); 367 assert!(RUN); 368 } 369 tx.send(()).unwrap(); 370 }); 371 } 372 373 unsafe { 374 O.call_once(|| { 375 assert!(!RUN); 376 RUN = true; 377 }); 378 assert!(RUN); 379 } 380 381 for _ in 0..10 { 382 rx.recv().unwrap(); 383 } 384 } 385 386 #[test] poison_bad()387 fn poison_bad() { 388 static O: Once = Once::new(); 389 390 // poison the once 391 let t = panic::catch_unwind(|| { 392 O.call_once(|| panic!()); 393 }); 394 assert!(t.is_err()); 395 396 // poisoning propagates 397 let t = panic::catch_unwind(|| { 398 O.call_once(|| {}); 399 }); 400 assert!(t.is_err()); 401 402 // we can subvert poisoning, however 403 let mut called = false; 404 O.call_once_force(|p| { 405 called = true; 406 assert!(p.poisoned()) 407 }); 408 assert!(called); 409 410 // once any success happens, we stop propagating the poison 411 O.call_once(|| {}); 412 } 413 414 #[test] wait_for_force_to_finish()415 fn wait_for_force_to_finish() { 416 static O: Once = Once::new(); 417 418 // poison the once 419 let t = panic::catch_unwind(|| { 420 O.call_once(|| panic!()); 421 }); 422 assert!(t.is_err()); 423 424 // make sure someone's waiting inside the once via a force 425 let (tx1, rx1) = channel(); 426 let (tx2, rx2) = channel(); 427 let t1 = thread::spawn(move || { 428 O.call_once_force(|p| { 429 assert!(p.poisoned()); 430 tx1.send(()).unwrap(); 431 rx2.recv().unwrap(); 432 }); 433 }); 434 435 rx1.recv().unwrap(); 436 437 // put another waiter on the once 438 let t2 = thread::spawn(|| { 439 let mut called = false; 440 O.call_once(|| { 441 called = true; 442 }); 443 assert!(!called); 444 }); 445 446 tx2.send(()).unwrap(); 447 448 assert!(t1.join().is_ok()); 449 assert!(t2.join().is_ok()); 450 } 451 452 #[test] test_once_debug()453 fn test_once_debug() { 454 static O: Once = Once::new(); 455 456 assert_eq!(format!("{:?}", O), "Once { state: New }"); 457 } 458 } 459