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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
//! Synchronous spinning-based synchronization primitives.
//!
//! The synchronization primitives in `maitake-sync` are _asynchronous_. They
//! are designed to be used with [`core::task`] and [`core::future`], and when
//! it is necessary to wait for another task to complete some work for the
//! current task to proceed, `maitake`'s synchronization primitives wait by
//! *yielding* to the asynchronous task scheduler to allow another task to
//! proceed.
//!
//! This module, on the other hand, provides _synchronous_ (or _blocking_)
//! synchronization primitives. Rather than yielding to the runtime, these
//! synchronization primitives will block the current CPU core (or thread, if
//! running in an environment with threads) until they are woken by other cores.
//! This is performed by *spinning*: issuing yield or pause instructions in a
//! loop until some value changes. These synchronization primitives are, in some
//! cases, necessary to implement the async synchronization primitives that form
//! `maitake-sync`'s core APIs. They are also exposed publicly so they can be
//! used in other projects, when a spinlock-based synchronization primitive is
//! needed.
//!
//! This module provides the following APIs:
//!
//! - [`Spinlock`]: a synchronous [mutual exclusion] spinlock, which implements
//!       the [`blocking::RawMutex`] and [`blocking::ScopedRawMutex`] traits.
//! - [`RwSpinlock`]: a synchronous [reader-writer] spinlock, which implements
//!       the [`blocking::RawRwLock`] trait.
//! - [`InitOnce`]: a cell storing a [`MaybeUninit`](core::mem::MaybeUninit)
//!       value which must be manually initialized prior to use.
//! - [`Lazy`]: an [`InitOnce`] cell coupled with an initializer function. The
//!       [`Lazy`] cell ensures the initializer is called to initialize the
//!       value the first time it is accessed.
//!
//! [mutual exclusion lock]: https://en.wikipedia.org/wiki/Mutual_exclusion
//! [reader-writer lock]: https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock
pub mod once;

pub use self::once::{InitOnce, Lazy};

use crate::{
    blocking,
    loom::sync::atomic::{AtomicBool, AtomicUsize, Ordering::*},
    util::{fmt, Backoff},
};
// This import is pulled out because we want to reference it in docs, and
// importing `use crate::blocking; use blocking::{RawMutex, RawRwLock};` makes
// `blocking` appear used even if it's not referenced directly, while
// `use crate::blocking::{self, RawMutex, RawRwLock};` causes the `self` to
// appear "unused" to rustc. Yes, this is stupid, but this workaround felt
// better than `allow(unused_imports)`.
use blocking::{RawMutex, RawRwLock};

/// A spinlock-based [`RawMutex`] implementation.
///
/// This mutex will spin with an exponential backoff while waiting for the lock
/// to become available.
///
/// This type implements the [`RawMutex`] and
/// [`ScopedRawMutex`](mutex_traits::ScopedRawMutex) traits from the
/// [`mutex-traits`] crate. This allows it to be used with the
/// [`blocking::Mutex`] type when a spinlock-based mutex is needed.
#[derive(Debug)]
pub struct Spinlock {
    locked: AtomicBool,
}

/// A spinlock-based [`RawRwLock`] implementation.
///
/// This type implements the [`blocking::RawRwLock`] trait. This allows it to be
/// used with the [`blocking::RwLock`] type when a spinlock-based reader-writer
/// lock is needed.
pub struct RwSpinlock {
    state: AtomicUsize,
}

// === impl RawSpinlock ===

impl Spinlock {
    loom_const_fn! {
        /// Returns a new `Spinlock`, in the unlocked state.
        pub fn new() -> Self {
            Self { locked: AtomicBool::new(false) }
        }
    }

    #[inline]
    #[must_use]
    fn is_locked(&self) -> bool {
        self.locked.load(Relaxed)
    }
}

impl Default for Spinlock {
    fn default() -> Self {
        Self::new()
    }
}

unsafe impl RawMutex for Spinlock {
    type GuardMarker = ();

    #[cfg_attr(test, track_caller)]
    fn lock(&self) {
        let mut boff = Backoff::default();
        while test_dbg!(self
            .locked
            .compare_exchange(false, true, Acquire, Acquire)
            .is_err())
        {
            while test_dbg!(self.is_locked()) {
                boff.spin();
            }
        }
    }

    #[cfg_attr(test, track_caller)]
    #[inline]
    fn try_lock(&self) -> bool {
        test_dbg!(self
            .locked
            .compare_exchange(false, true, Acquire, Acquire)
            .is_ok())
    }

    #[cfg_attr(test, track_caller)]
    #[inline]
    unsafe fn unlock(&self) {
        test_dbg!(self.locked.store(false, Release));
    }

    #[inline]
    fn is_locked(&self) -> bool {
        Spinlock::is_locked(self)
    }
}

#[cfg(not(loom))]
impl blocking::ConstInit for Spinlock {
    // As usual, clippy is totally wrong about this --- the whole point of this
    // constant is to create a *new* spinlock every time.
    #[allow(clippy::declare_interior_mutable_const)]
    const INIT: Self = Spinlock::new();
}

const UNLOCKED: usize = 0;
const WRITER: usize = 1 << 0;
const READER: usize = 1 << 1;

impl RwSpinlock {
    loom_const_fn! {
        pub(crate) fn new() -> Self {
            Self {
                state: AtomicUsize::new(UNLOCKED),
            }
        }
    }

    pub(crate) fn reader_count(&self) -> usize {
        self.state.load(Relaxed) >> 1
    }
}

unsafe impl RawRwLock for RwSpinlock {
    type GuardMarker = ();

    #[cfg_attr(test, track_caller)]
    fn lock_shared(&self) {
        let mut boff = Backoff::new();
        while !self.try_lock_shared() {
            boff.spin();
        }
    }

    #[cfg_attr(test, track_caller)]
    fn try_lock_shared(&self) -> bool {
        // Add a reader.
        let state = test_dbg!(self.state.fetch_add(READER, Acquire));

        // Ensure we don't overflow the reader count and clobber the lock's
        // state.
        assert!(
            state < usize::MAX - (READER * 2),
            "read lock counter overflow! this is very bad"
        );

        // Is the write lock held? If so, undo the increment and bail.
        if state & WRITER == 1 {
            test_dbg!(self.state.fetch_sub(READER, Release));
            false
        } else {
            true
        }
    }

    #[cfg_attr(test, track_caller)]
    #[inline]
    unsafe fn unlock_shared(&self) {
        let _val = test_dbg!(self.state.fetch_sub(READER, Release));
        debug_assert_eq!(
            _val & WRITER,
            0,
            "tried to drop a read guard while write locked, something is Very Wrong!"
        )
    }

    #[cfg_attr(test, track_caller)]
    fn lock_exclusive(&self) {
        let mut backoff = Backoff::new();

        // Wait for the lock to become available and set the `WRITER` bit.
        //
        // Note that, unlike the `lock_shared` method, we don't use the
        // `try_lock_exclusive` method here, as we would like to use
        // `compare_exchange_weak` to allow spurious failures for improved
        // performance. The `try_lock_exclusive` method  cannot use
        // `compare_exchange_weak`, because it will never retry, and
        // a spurious failure means we would incorrectly fail to lock the RwLock
        // when we should have successfully locked it.
        while test_dbg!(self
            .state
            .compare_exchange_weak(UNLOCKED, WRITER, Acquire, Relaxed))
        .is_err()
        {
            test_dbg!(backoff.spin());
        }
    }

    #[cfg_attr(test, track_caller)]
    #[inline]
    fn try_lock_exclusive(&self) -> bool {
        test_dbg!(self
            .state
            .compare_exchange(UNLOCKED, WRITER, Acquire, Relaxed))
        .is_ok()
    }

    #[cfg_attr(test, track_caller)]
    #[inline]
    unsafe fn unlock_exclusive(&self) {
        let _val = test_dbg!(self.state.swap(UNLOCKED, Release));
    }

    #[inline]
    fn is_locked(&self) -> bool {
        self.state.load(Relaxed) & (WRITER | READER) != 0
    }

    #[inline]
    fn is_locked_exclusive(&self) -> bool {
        self.state.load(Relaxed) & WRITER == 1
    }
}

#[cfg(not(loom))]
impl blocking::ConstInit for RwSpinlock {
    // As usual, clippy is totally wrong about this --- the whole point of this
    // constant is to create a *new* spinlock every time.
    #[allow(clippy::declare_interior_mutable_const)]
    const INIT: Self = RwSpinlock::new();
}

impl fmt::Debug for RwSpinlock {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let state = &self.state.load(Relaxed);
        f.debug_struct("RwSpinlock")
            // N.B.: this does *not* use the `reader_count` and `has_writer`
            // methods *intentionally*, because those two methods perform
            // independent reads of the lock's state, and may race with other
            // lock operations that occur concurrently. If, for example, we read
            // a non-zero reader count, and then read the state again to check
            // for a writer, the reader may have been released and a write lock
            // acquired between the two reads, resulting in the `Debug` impl
            // displaying an invalid state when the lock was not actually *in*
            // such a state!
            //
            // Therefore, we instead perform a single load to snapshot the state
            // and unpack both the reader count and the writer count from the
            // lock.
            .field("readers", &(state >> 1))
            .field("writer", &(state & WRITER))
            .finish()
    }
}