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