state/
init.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
use crate::shim::sync::atomic::{AtomicBool, Ordering::{AcqRel, Acquire, Release, Relaxed}};
use crate::shim::thread::yield_now;

/// An atomic initializer: mutual exclusion during initialization.
pub struct Init {
    started: AtomicBool,
    done: AtomicBool
}

impl Init {
    /// A ready-to-init initializer.
    #[cfg(not(loom))]
    pub const fn new() -> Init {
        Init {
            started: AtomicBool::new(false),
            done: AtomicBool::new(false)
        }
    }

    /// A ready-to-init initializer.
    #[cfg(loom)]
    pub fn new() -> Init {
        Init {
            started: AtomicBool::new(false),
            done: AtomicBool::new(false)
        }
    }

    /// Returns true if initialization has completed without blocking. If this
    /// function returns false, it may be the case that initialization is
    /// currently in progress. If this function returns `true`, intialization is
    /// guaranteed to be completed.
    #[inline(always)]
    pub fn has_completed(&self) -> bool {
        self.done.load(Acquire)
    }

    /// Mark this initialization as complete, unblocking all threads that may be
    /// waiting. Only the caller that received `true` from `needed()` is
    /// expected to call this method.
    #[inline(always)]
    pub fn mark_complete(&self) {
        // If this is being called from outside of a `needed` block, we need to
        // ensure that `started` is `true` to avoid racing with (return `true`
        // to) future `needed` calls.
        self.started.store(true, Release);
        self.done.store(true, Release);
    }

    /// Blocks until initialization is marked as completed.
    ///
    ///
    // NOTE: Internally, this waits for the the done flag.
    #[inline(always)]
    pub fn wait_until_complete(&self) {
        while !self.done.load(Acquire) { yield_now() }
    }

    #[cold]
    #[inline(always)]
    fn try_to_need_init(&self) -> bool {
        // Quickly check if initialization has already started elsewhere.
        if self.started.load(Relaxed) {
            // If it has, wait until it's finished before returning. Finishing
            // is marked by calling `mark_complete`.
            self.wait_until_complete();
            return false;
        }

        // Try to be the first. If we lose (init_started is true), we wait.
        if self.started.compare_exchange(false, true, AcqRel, Relaxed).is_err() {
            // Another compare_and_swap won. Wait until they're done.
            self.wait_until_complete();
            return false;
        }

        true
    }

    /// If initialization is complete, returns `false`. Otherwise, returns
    /// `true` exactly once, to the caller that should perform initialization.
    /// All other calls block until initialization is marked completed, at which
    /// point `false` is returned.
    ///
    /// If this function is called from multiple threads simulatenously, exactly
    /// one thread is guaranteed to receive `true`. All other threads are
    /// blocked until initialization is marked completed.
    #[inline(always)]
    pub fn needed(&self) -> bool {
        // Quickly check if initialization has finished, and return if so.
        if self.has_completed() {
            return false;
        }

        // We call a different function to attempt the intialiaztion to use
        // Rust's `cold` attribute to try let LLVM know that this is unlikely.
        self.try_to_need_init()
    }
}