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