Skip to main content

moirai_sync/sync/
spin_lock.rs

1use std::cell::UnsafeCell;
2use std::fmt;
3use std::hint;
4use std::ops::{Deref, DerefMut};
5use std::sync::atomic::{AtomicBool, Ordering};
6
7/// Maximum backoff iterations for `SpinLock` (TBB-inspired)
8const SPINLOCK_MAX_BACKOFF: usize = 64;
9
10/// Maximum spin attempts before yielding to scheduler
11const SPINLOCK_MAX_SPINS_BEFORE_YIELD: usize = 1000;
12
13// SpinLock backoff constants (TBB-inspired)
14/// Initial backoff iterations for SpinLock
15const SPINLOCK_INITIAL_BACKOFF: usize = 1;
16
17/// A spin lock for very short critical sections with TBB-inspired exponential backoff.
18///
19/// This implementation uses exponential backoff and adaptive yielding for better
20/// performance under contention. The lock is cache-line aligned to prevent false sharing.
21///
22/// Use only when you know the critical section is extremely short (< 1μs).
23#[repr(align(64))] // Cache line alignment to prevent false sharing
24pub struct SpinLock<T> {
25    locked: AtomicBool,
26    data: UnsafeCell<T>,
27}
28
29impl<T> fmt::Debug for SpinLock<T> {
30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
31        let locked = self.locked.load(Ordering::Relaxed);
32        f.debug_struct("SpinLock")
33            .field("locked", &locked)
34            .finish_non_exhaustive()
35    }
36}
37
38unsafe impl<T: Send> Send for SpinLock<T> {}
39unsafe impl<T: Send> Sync for SpinLock<T> {}
40
41impl<T> SpinLock<T> {
42    /// Create a new spin lock.
43    pub const fn new(data: T) -> Self {
44        Self {
45            locked: AtomicBool::new(false),
46            data: UnsafeCell::new(data),
47        }
48    }
49
50    /// Lock the spin lock with TBB-inspired exponential backoff.
51    ///
52    /// This implementation uses:
53    /// - Read-before-CAS to reduce memory contention
54    /// - Exponential backoff starting from 1 iteration up to 64
55    /// - Adaptive yielding after prolonged spinning
56    pub fn lock(&self) -> SpinLockGuard<'_, T> {
57        let mut backoff = SPINLOCK_INITIAL_BACKOFF;
58        let mut total_spins = 0;
59
60        loop {
61            // Read-before-CAS: only attempt atomic write if lock is observed unlocked
62            if !self.locked.load(Ordering::Relaxed)
63                && self
64                    .locked
65                    .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
66                    .is_ok()
67            {
68                return SpinLockGuard {
69                    lock: self,
70                    _phantom: std::marker::PhantomData,
71                };
72            }
73
74            // Exponential backoff with CPU pause instructions
75            for _ in 0..backoff {
76                hint::spin_loop();
77            }
78
79            // Double the backoff up to maximum
80            if backoff < SPINLOCK_MAX_BACKOFF {
81                backoff = backoff.saturating_mul(2);
82            }
83
84            total_spins += backoff;
85
86            // After many attempts, yield to scheduler to be cooperative
87            if total_spins >= SPINLOCK_MAX_SPINS_BEFORE_YIELD {
88                std::thread::yield_now();
89                total_spins = 0;
90                backoff = SPINLOCK_INITIAL_BACKOFF; // Reset backoff after yielding
91            }
92        }
93    }
94
95    /// Try to lock without spinning.
96    pub fn try_lock(&self) -> Option<SpinLockGuard<'_, T>> {
97        if !self.locked.load(Ordering::Relaxed)
98            && self
99                .locked
100                .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
101                .is_ok()
102        {
103            Some(SpinLockGuard {
104                lock: self,
105                _phantom: std::marker::PhantomData,
106            })
107        } else {
108            None
109        }
110    }
111}
112
113/// Guard for SpinLock that automatically unlocks on drop.
114pub struct SpinLockGuard<'a, T> {
115    lock: &'a SpinLock<T>,
116    _phantom: std::marker::PhantomData<T>,
117}
118
119impl<'a, T> Drop for SpinLockGuard<'a, T> {
120    fn drop(&mut self) {
121        self.lock.locked.store(false, Ordering::Release);
122    }
123}
124
125impl<'a, T> Deref for SpinLockGuard<'a, T> {
126    type Target = T;
127
128    fn deref(&self) -> &Self::Target {
129        unsafe { &*self.lock.data.get() }
130    }
131}
132
133impl<'a, T> DerefMut for SpinLockGuard<'a, T> {
134    fn deref_mut(&mut self) -> &mut Self::Target {
135        unsafe { &mut *self.lock.data.get() }
136    }
137}