parking_lot/
once.rs

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(#[automatically_derived]
impl ::core::marker::Copy for OnceState { }Copy, #[automatically_derived]
impl ::core::clone::Clone for OnceState {
    #[inline]
    fn clone(&self) -> OnceState { *self }
}Clone, #[automatically_derived]
impl ::core::cmp::Eq for OnceState {
    #[inline]
    #[doc(hidden)]
    #[coverage(off)]
    fn assert_receiver_is_total_eq(&self) -> () {}
}Eq, #[automatically_derived]
impl ::core::cmp::PartialEq for OnceState {
    #[inline]
    fn eq(&self, other: &OnceState) -> bool {
        let __self_discr = ::core::intrinsics::discriminant_value(self);
        let __arg1_discr = ::core::intrinsics::discriminant_value(other);
        __self_discr == __arg1_discr
    }
}PartialEq, #[automatically_derived]
impl ::core::fmt::Debug for OnceState {
    #[inline]
    fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
        ::core::fmt::Formatter::write_str(f,
            match self {
                OnceState::New => "New",
                OnceState::Poisoned => "Poisoned",
                OnceState::InProgress => "InProgress",
                OnceState::Done => "Done",
            })
    }
}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        #[allow(non_exhaustive_omitted_patterns)] match self {
    OnceState::Poisoned => true,
    _ => false,
}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        #[allow(non_exhaustive_omitted_patterns)] match self {
    OnceState::Done => true,
    _ => false,
}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                {
    ::core::panicking::panic_fmt(format_args!("Once instance has previously been poisoned"));
};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 = |_, _| ::core::panicking::panic("internal error: entered unreachable code")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.debug_struct("Once")
322            .field("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}
parking_lot/once.rs

parking_lot/
once.rs
