Skip to main content

moirai_sync/sync/
spin_lock.rs

1use moirai_utils::CacheAligned;
2use std::cell::UnsafeCell;
3use std::fmt;
4use std::hint;
5use std::ops::{Deref, DerefMut};
6use std::sync::atomic::{AtomicBool, Ordering};
7
8/// Maximum backoff iterations for `SpinLock` (TBB-inspired)
9const SPINLOCK_MAX_BACKOFF: usize = 64;
10
11/// Maximum spin attempts before yielding to scheduler
12const SPINLOCK_MAX_SPINS_BEFORE_YIELD: usize = 1000;
13
14// SpinLock backoff constants (TBB-inspired)
15/// Initial backoff iterations for SpinLock
16const SPINLOCK_INITIAL_BACKOFF: usize = 1;
17
18/// A spin lock for very short critical sections with TBB-inspired exponential backoff.
19///
20/// This implementation uses exponential backoff and adaptive yielding for better
21/// performance under contention. The lock is cache-line aligned to prevent false sharing.
22///
23/// Use only when you know the critical section is extremely short (< 1μs).
24///
25/// `locked` is wrapped in [`CacheAligned`], which both aligns the lock to
26/// `moirai_utils::DESTRUCTIVE_INTERFERENCE_SIZE` (so a neighbouring object
27/// cannot share its sector) and separates the contended flag from `data` (so a
28/// spinning acquirer does not invalidate the payload). The separation is
29/// per-target and owned by `moirai-utils`, not a literal here.
30pub struct SpinLock<T> {
31    locked: CacheAligned<AtomicBool>,
32    data: UnsafeCell<T>,
33}
34
35impl<T> fmt::Debug for SpinLock<T> {
36    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37        let locked = self.locked.load(Ordering::Relaxed);
38        f.debug_struct("SpinLock")
39            .field("locked", &locked)
40            .finish_non_exhaustive()
41    }
42}
43
44// SAFETY: guarded values move with the lock; no address-sensitive state
45// beyond `T`.
46unsafe impl<T: Send> Send for SpinLock<T> {}
47// SAFETY: the atomic `locked` bit serializes access, so data references via
48// guards are exclusive while held; `T: Send` covers cross-thread transfer.
49unsafe impl<T: Send> Sync for SpinLock<T> {}
50
51impl<T> SpinLock<T> {
52    /// Create a new spin lock.
53    pub const fn new(data: T) -> Self {
54        Self {
55            locked: CacheAligned::new(AtomicBool::new(false)),
56            data: UnsafeCell::new(data),
57        }
58    }
59
60    /// Lock the spin lock with TBB-inspired exponential backoff.
61    ///
62    /// This implementation uses:
63    /// - Read-before-CAS to reduce memory contention
64    /// - Exponential backoff starting from 1 iteration up to 64
65    /// - Adaptive yielding after prolonged spinning
66    pub fn lock(&self) -> SpinLockGuard<'_, T> {
67        let mut backoff = SPINLOCK_INITIAL_BACKOFF;
68        let mut total_spins = 0;
69
70        loop {
71            // Read-before-CAS: only attempt atomic write if lock is observed unlocked
72            if !self.locked.load(Ordering::Relaxed)
73                && self
74                    .locked
75                    .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
76                    .is_ok()
77            {
78                return SpinLockGuard {
79                    lock: self,
80                    _phantom: std::marker::PhantomData,
81                };
82            }
83
84            // Exponential backoff with CPU pause instructions
85            for _ in 0..backoff {
86                hint::spin_loop();
87            }
88
89            // Double the backoff up to maximum
90            if backoff < SPINLOCK_MAX_BACKOFF {
91                backoff = backoff.saturating_mul(2);
92            }
93
94            total_spins += backoff;
95
96            // After many attempts, yield to scheduler to be cooperative
97            if total_spins >= SPINLOCK_MAX_SPINS_BEFORE_YIELD {
98                std::thread::yield_now();
99                total_spins = 0;
100                backoff = SPINLOCK_INITIAL_BACKOFF; // Reset backoff after yielding
101            }
102        }
103    }
104
105    /// Try to lock without spinning.
106    pub fn try_lock(&self) -> Option<SpinLockGuard<'_, T>> {
107        if !self.locked.load(Ordering::Relaxed)
108            && self
109                .locked
110                .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
111                .is_ok()
112        {
113            Some(SpinLockGuard {
114                lock: self,
115                _phantom: std::marker::PhantomData,
116            })
117        } else {
118            None
119        }
120    }
121}
122
123/// Guard for SpinLock that automatically unlocks on drop.
124pub struct SpinLockGuard<'a, T> {
125    lock: &'a SpinLock<T>,
126    _phantom: std::marker::PhantomData<T>,
127}
128
129impl<'a, T> Drop for SpinLockGuard<'a, T> {
130    fn drop(&mut self) {
131        self.lock.locked.store(false, Ordering::Release);
132    }
133}
134
135impl<'a, T> Deref for SpinLockGuard<'a, T> {
136    type Target = T;
137
138    fn deref(&self) -> &Self::Target {
139        // SAFETY: holding the guard proves the lock bit is ours, so no other
140        // reference to `data` exists; shared reborrow cannot race a writer.
141        unsafe { &*self.lock.data.get() }
142    }
143}
144
145impl<'a, T> DerefMut for SpinLockGuard<'a, T> {
146    fn deref_mut(&mut self) -> &mut Self::Target {
147        // SAFETY: unique guard plus the lock bit exclude all other access to
148        // `data` for the guard's lifetime.
149        unsafe { &mut *self.lock.data.get() }
150    }
151}