1 | use crate::sync::atomic::{ |
2 | self, |
3 | Ordering::{Acquire, Relaxed, Release}, |
4 | }; |
5 | use crate::sys::futex::{futex_wait, futex_wake}; |
6 | |
7 | cfg_if::cfg_if! { |
8 | if #[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 | |
18 | pub struct Mutex { |
19 | futex: Atomic, |
20 | } |
21 | |
22 | const UNLOCKED: State = 0; |
23 | const LOCKED: State = 1; // locked, no other threads waiting |
24 | const CONTENDED: State = 2; // locked, and other threads waiting (contended) |
25 | |
26 | impl 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 | |