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