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