1 use crate::std::sync::atomic::{ 2 AtomicU32, 3 Ordering::{Acquire, Relaxed, Release}, 4 }; 5 use crate::std::sys::futex::{futex_wait, futex_wake}; 6 7 pub struct Mutex { 8 /// 0: unlocked 9 /// 1: locked, no other threads waiting 10 /// 2: locked, and other threads waiting (contended) 11 futex: AtomicU32, 12 } 13 14 impl Mutex { 15 #[inline] 16 pub const fn new() -> Self { 17 Self { 18 futex: AtomicU32::new(0), 19 } 20 } 21 22 #[inline] 23 pub fn try_lock(&self) -> bool { 24 self.futex.compare_exchange(0, 1, Acquire, Relaxed).is_ok() 25 } 26 27 #[inline] 28 pub fn lock(&self) { 29 if self.futex.compare_exchange(0, 1, Acquire, Relaxed).is_err() { 30 self.lock_contended(); 31 } 32 } 33 34 #[cold] 35 fn lock_contended(&self) { 36 // Spin first to speed things up if the lock is released quickly. 37 let mut state = self.spin(); 38 39 // If it's unlocked now, attempt to take the lock 40 // without marking it as contended. 41 if state == 0 { 42 match self.futex.compare_exchange(0, 1, Acquire, Relaxed) { 43 Ok(_) => return, // Locked! 44 Err(s) => state = s, 45 } 46 } 47 48 loop { 49 // Put the lock in contended state. 50 // We avoid an unnecessary write if it as already set to 2, 51 // to be friendlier for the caches. 52 if state != 2 && self.futex.swap(2, Acquire) == 0 { 53 // We changed it from 0 to 2, so we just successfully locked it. 54 return; 55 } 56 57 // Wait for the futex to change state, assuming it is still 2. 58 futex_wait(&self.futex, 2, None); 59 60 // Spin again after waking up. 61 state = self.spin(); 62 } 63 } 64 65 fn spin(&self) -> u32 { 66 let mut spin = 100; 67 loop { 68 // We only use `load` (and not `swap` or `compare_exchange`) 69 // while spinning, to be easier on the caches. 70 let state = self.futex.load(Relaxed); 71 72 // We stop spinning when the mutex is unlocked (0), 73 // but also when it's contended (2). 74 if state != 1 || spin == 0 { 75 return state; 76 } 77 78 crate::std::hint::spin_loop(); 79 spin -= 1; 80 } 81 } 82 83 #[inline] 84 pub unsafe fn unlock(&self) { 85 if self.futex.swap(0, Release) == 2 { 86 // We only wake up one thread. When that thread locks the mutex, it 87 // will mark the mutex as contended (2) (see lock_contended above), 88 // which makes sure that any other waiting threads will also be 89 // woken up eventually. 90 self.wake(); 91 } 92 } 93 94 #[cold] 95 fn wake(&self) { 96 futex_wake(&self.futex); 97 } 98 } 99