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