std/sys/sync/mutex/
futex.rs

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}