axstd/sync/
mutex.rs

1//! A naïve sleeping mutex.
2
3use core::sync::atomic::{AtomicU64, Ordering};
4
5use arceos_api::task::{self as api, AxWaitQueueHandle};
6
7/// A [`lock_api::RawMutex`] implementation.
8///
9/// When the mutex is locked, the current task will block and be put into the
10/// wait queue. When the mutex is unlocked, all tasks waiting on the queue
11/// will be woken up.
12pub struct RawMutex {
13    wq: AxWaitQueueHandle,
14    owner_id: AtomicU64,
15}
16
17impl RawMutex {
18    /// Creates a [`RawMutex`].
19    #[inline(always)]
20    pub const fn new() -> Self {
21        Self {
22            wq: AxWaitQueueHandle::new(),
23            owner_id: AtomicU64::new(0),
24        }
25    }
26}
27
28unsafe impl lock_api::RawMutex for RawMutex {
29    /// Initial value for an unlocked mutex.
30    ///
31    /// A “non-constant” const item is a legacy way to supply an initialized value to downstream
32    /// static items. Can hopefully be replaced with `const fn new() -> Self` at some point.
33    #[allow(clippy::declare_interior_mutable_const)]
34    const INIT: Self = RawMutex::new();
35
36    type GuardMarker = lock_api::GuardSend;
37
38    #[inline(always)]
39    fn lock(&self) {
40        let current_id = api::ax_current_task_id();
41        loop {
42            // Can fail to lock even if the spinlock is not locked. May be more efficient than `try_lock`
43            // when called in a loop.
44            match self.owner_id.compare_exchange_weak(
45                0,
46                current_id,
47                Ordering::Acquire,
48                Ordering::Relaxed,
49            ) {
50                Ok(_) => break,
51                Err(owner_id) => {
52                    assert_ne!(
53                        owner_id, current_id,
54                        "Thread({current_id}) tried to acquire mutex it already owns.",
55                    );
56                    // Wait until the lock looks unlocked before retrying
57                    api::ax_wait_queue_wait_until(&self.wq, || !self.is_locked(), None);
58                }
59            }
60        }
61    }
62
63    #[inline(always)]
64    fn try_lock(&self) -> bool {
65        let current_id = api::ax_current_task_id();
66        // The reason for using a strong compare_exchange is explained here:
67        // https://github.com/Amanieu/parking_lot/pull/207#issuecomment-575869107
68        self.owner_id
69            .compare_exchange(0, current_id, Ordering::Acquire, Ordering::Relaxed)
70            .is_ok()
71    }
72
73    #[inline(always)]
74    unsafe fn unlock(&self) {
75        let owner_id = self.owner_id.swap(0, Ordering::Release);
76        let current_id = api::ax_current_task_id();
77        assert_eq!(
78            owner_id, current_id,
79            "Thread({current_id}) tried to release mutex it doesn't own",
80        );
81        // wake up one waiting thread.
82        api::ax_wait_queue_wake(&self.wq, 1);
83    }
84
85    #[inline(always)]
86    fn is_locked(&self) -> bool {
87        self.owner_id.load(Ordering::Relaxed) != 0
88    }
89}
90
91/// An alias of [`lock_api::Mutex`].
92pub type Mutex<T> = lock_api::Mutex<RawMutex, T>;
93/// An alias of [`lock_api::MutexGuard`].
94pub type MutexGuard<'a, T> = lock_api::MutexGuard<'a, RawMutex, T>;