library/std/src/sys/unix/locks/futex_mutex.rs

   1 use crate::sync::atomic::{
   2     AtomicU32,
   3     Ordering::{Acquire, Relaxed, Release},
   4 };
   5 use crate::sys::futex::{futex_wait, futex_wake};
   6
   7 pub 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
  14 impl 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 }