state/
init.rs

1use crate::shim::sync::atomic::{AtomicBool, Ordering::{AcqRel, Acquire, Release, Relaxed}};
2use crate::shim::thread::yield_now;
3
4/// An atomic initializer: mutual exclusion during initialization.
5pub struct Init {
6    started: AtomicBool,
7    done: AtomicBool
8}
9
10impl Init {
11    /// A ready-to-init initializer.
12    #[cfg(not(loom))]
13    pub const fn new() -> Init {
14        Init {
15            started: AtomicBool::new(false),
16            done: AtomicBool::new(false)
17        }
18    }
19
20    /// A ready-to-init initializer.
21    #[cfg(loom)]
22    pub fn new() -> Init {
23        Init {
24            started: AtomicBool::new(false),
25            done: AtomicBool::new(false)
26        }
27    }
28
29    /// Returns true if initialization has completed without blocking. If this
30    /// function returns false, it may be the case that initialization is
31    /// currently in progress. If this function returns `true`, intialization is
32    /// guaranteed to be completed.
33    #[inline(always)]
34    pub fn has_completed(&self) -> bool {
35        self.done.load(Acquire)
36    }
37
38    /// Mark this initialization as complete, unblocking all threads that may be
39    /// waiting. Only the caller that received `true` from `needed()` is
40    /// expected to call this method.
41    #[inline(always)]
42    pub fn mark_complete(&self) {
43        // If this is being called from outside of a `needed` block, we need to
44        // ensure that `started` is `true` to avoid racing with (return `true`
45        // to) future `needed` calls.
46        self.started.store(true, Release);
47        self.done.store(true, Release);
48    }
49
50    /// Blocks until initialization is marked as completed.
51    ///
52    ///
53    // NOTE: Internally, this waits for the the done flag.
54    #[inline(always)]
55    pub fn wait_until_complete(&self) {
56        while !self.done.load(Acquire) { yield_now() }
57    }
58
59    #[cold]
60    #[inline(always)]
61    fn try_to_need_init(&self) -> bool {
62        // Quickly check if initialization has already started elsewhere.
63        if self.started.load(Relaxed) {
64            // If it has, wait until it's finished before returning. Finishing
65            // is marked by calling `mark_complete`.
66            self.wait_until_complete();
67            return false;
68        }
69
70        // Try to be the first. If we lose (init_started is true), we wait.
71        if self.started.compare_exchange(false, true, AcqRel, Relaxed).is_err() {
72            // Another compare_and_swap won. Wait until they're done.
73            self.wait_until_complete();
74            return false;
75        }
76
77        true
78    }
79
80    /// If initialization is complete, returns `false`. Otherwise, returns
81    /// `true` exactly once, to the caller that should perform initialization.
82    /// All other calls block until initialization is marked completed, at which
83    /// point `false` is returned.
84    ///
85    /// If this function is called from multiple threads simulatenously, exactly
86    /// one thread is guaranteed to receive `true`. All other threads are
87    /// blocked until initialization is marked completed.
88    #[inline(always)]
89    pub fn needed(&self) -> bool {
90        // Quickly check if initialization has finished, and return if so.
91        if self.has_completed() {
92            return false;
93        }
94
95        // We call a different function to attempt the intialiaztion to use
96        // Rust's `cold` attribute to try let LLVM know that this is unlikely.
97        self.try_to_need_init()
98    }
99}