Coverage Report

Created: 2025-08-26 07:09

/rust/registry/src/index.crates.io-6f17d22bba15001f/parking_lot-0.12.4/src/once.rs
Line
Count
Source (jump to first uncovered line)
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]
42
0
    pub fn poisoned(self) -> bool {
43
0
        matches!(self, OnceState::Poisoned)
44
0
    }
45
46
    /// Returns whether the associated `Once` has successfully executed a
47
    /// closure.
48
    #[inline]
49
0
    pub fn done(self) -> bool {
50
0
        matches!(self, OnceState::Done)
51
0
    }
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
/// ```
77
pub struct Once(AtomicU8);
78
79
impl Once {
80
    /// Creates a new `Once` value.
81
    #[inline]
82
0
    pub const fn new() -> Once {
83
0
        Once(AtomicU8::new(0))
84
0
    }
85
86
    /// Returns the current state of this `Once`.
87
    #[inline]
88
0
    pub fn state(&self) -> OnceState {
89
0
        let state = self.0.load(Ordering::Acquire);
90
0
        if state & DONE_BIT != 0 {
91
0
            OnceState::Done
92
0
        } else if state & LOCKED_BIT != 0 {
93
0
            OnceState::InProgress
94
0
        } else if state & POISON_BIT != 0 {
95
0
            OnceState::Poisoned
96
        } else {
97
0
            OnceState::New
98
        }
99
0
    }
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
0
    pub fn call_once<F>(&self, f: F)
152
0
    where
153
0
        F: FnOnce(),
154
0
    {
155
0
        if self.0.load(Ordering::Acquire) == DONE_BIT {
156
0
            return;
157
0
        }
158
0
159
0
        let mut f = Some(f);
160
0
        self.call_once_slow(false, &mut |_| unsafe { f.take().unchecked_unwrap()() });
161
0
    }
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
0
    pub fn call_once_force<F>(&self, f: F)
174
0
    where
175
0
        F: FnOnce(OnceState),
176
0
    {
177
0
        if self.0.load(Ordering::Acquire) == DONE_BIT {
178
0
            return;
179
0
        }
180
0
181
0
        let mut f = Some(f);
182
0
        self.call_once_slow(true, &mut |state| unsafe {
183
0
            f.take().unchecked_unwrap()(state)
184
0
        });
185
0
    }
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
0
    fn call_once_slow(&self, ignore_poison: bool, f: &mut dyn FnMut(OnceState)) {
200
0
        let mut spinwait = SpinWait::new();
201
0
        let mut state = self.0.load(Ordering::Relaxed);
202
        loop {
203
            // If another thread called the closure, we're done
204
0
            if state & DONE_BIT != 0 {
205
                // An acquire fence is needed here since we didn't load the
206
                // state with Ordering::Acquire.
207
0
                fence(Ordering::Acquire);
208
0
                return;
209
0
            }
210
0
211
0
            // If the state has been poisoned and we aren't forcing, then panic
212
0
            if state & POISON_BIT != 0 && !ignore_poison {
213
                // Need the fence here as well for the same reason
214
0
                fence(Ordering::Acquire);
215
0
                panic!("Once instance has previously been poisoned");
216
0
            }
217
0
218
0
            // Grab the lock if it isn't locked, even if there is a queue on it.
219
0
            // We also clear the poison bit since we are going to try running
220
0
            // the closure again.
221
0
            if state & LOCKED_BIT == 0 {
222
0
                match self.0.compare_exchange_weak(
223
0
                    state,
224
0
                    (state | LOCKED_BIT) & !POISON_BIT,
225
0
                    Ordering::Acquire,
226
0
                    Ordering::Relaxed,
227
0
                ) {
228
0
                    Ok(_) => break,
229
0
                    Err(x) => state = x,
230
0
                }
231
0
                continue;
232
0
            }
233
0
234
0
            // If there is no queue, try spinning a few times
235
0
            if state & PARKED_BIT == 0 && spinwait.spin() {
236
0
                state = self.0.load(Ordering::Relaxed);
237
0
                continue;
238
0
            }
239
0
240
0
            // Set the parked bit
241
0
            if state & PARKED_BIT == 0 {
242
0
                if let Err(x) = self.0.compare_exchange_weak(
243
0
                    state,
244
0
                    state | PARKED_BIT,
245
0
                    Ordering::Relaxed,
246
0
                    Ordering::Relaxed,
247
0
                ) {
248
0
                    state = x;
249
0
                    continue;
250
0
                }
251
0
            }
252
253
            // Park our thread until we are woken up by the thread that owns the
254
            // lock.
255
0
            let addr = self as *const _ as usize;
256
0
            let validate = || self.0.load(Ordering::Relaxed) == LOCKED_BIT | PARKED_BIT;
257
0
            let before_sleep = || {};
258
0
            let timed_out = |_, _| unreachable!();
259
0
            unsafe {
260
0
                parking_lot_core::park(
261
0
                    addr,
262
0
                    validate,
263
0
                    before_sleep,
264
0
                    timed_out,
265
0
                    DEFAULT_PARK_TOKEN,
266
0
                    None,
267
0
                );
268
0
            }
269
0
270
0
            // Loop back and check if the done bit was set
271
0
            spinwait.reset();
272
0
            state = self.0.load(Ordering::Relaxed);
273
        }
274
275
        struct PanicGuard<'a>(&'a Once);
276
        impl<'a> Drop for PanicGuard<'a> {
277
0
            fn drop(&mut self) {
278
0
                // Mark the state as poisoned, unlock it and unpark all threads.
279
0
                let once = self.0;
280
0
                let state = once.0.swap(POISON_BIT, Ordering::Release);
281
0
                if state & PARKED_BIT != 0 {
282
0
                    let addr = once as *const _ as usize;
283
0
                    unsafe {
284
0
                        parking_lot_core::unpark_all(addr, DEFAULT_UNPARK_TOKEN);
285
0
                    }
286
0
                }
287
0
            }
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
0
        let guard = PanicGuard(self);
293
0
        let once_state = if state & POISON_BIT != 0 {
294
0
            OnceState::Poisoned
295
        } else {
296
0
            OnceState::New
297
        };
298
0
        f(once_state);
299
0
        mem::forget(guard);
300
0
301
0
        // Now unlock the state, set the done bit and unpark all threads
302
0
        let state = self.0.swap(DONE_BIT, Ordering::Release);
303
0
        if state & PARKED_BIT != 0 {
304
0
            let addr = self as *const _ as usize;
305
0
            unsafe {
306
0
                parking_lot_core::unpark_all(addr, DEFAULT_UNPARK_TOKEN);
307
0
            }
308
0
        }
309
0
    }
310
}
311
312
impl Default for Once {
313
    #[inline]
314
0
    fn default() -> Once {
315
0
        Once::new()
316
0
    }
317
}
318
319
impl fmt::Debug for Once {
320
0
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
321
0
        f.debug_struct("Once")
322
0
            .field("state", &self.state())
323
0
            .finish()
324
0
    }
325
}
326
327
#[cfg(test)]
328
mod 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
}