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    const INIT: Self = RawMutex::new();
30
31    type GuardMarker = lock_api::GuardSend;
32
33    #[inline(always)]
34    fn lock(&self) {
35        let current_id = api::ax_current_task_id();
36        loop {
37            // Can fail to lock even if the spinlock is not locked. May be more efficient than `try_lock`
38            // when called in a loop.
39            match self.owner_id.compare_exchange_weak(
40                0,
41                current_id,
42                Ordering::Acquire,
43                Ordering::Relaxed,
44            ) {
45                Ok(_) => break,
46                Err(owner_id) => {
47                    assert_ne!(
48                        owner_id, current_id,
49                        "Thread({current_id}) tried to acquire mutex it already owns.",
50                    );
51                    // Wait until the lock looks unlocked before retrying
52                    api::ax_wait_queue_wait_until(&self.wq, || !self.is_locked(), None);
53                }
54            }
55        }
56    }
57
58    #[inline(always)]
59    fn try_lock(&self) -> bool {
60        let current_id = api::ax_current_task_id();
61        // The reason for using a strong compare_exchange is explained here:
62        // https://github.com/Amanieu/parking_lot/pull/207#issuecomment-575869107
63        self.owner_id
64            .compare_exchange(0, current_id, Ordering::Acquire, Ordering::Relaxed)
65            .is_ok()
66    }
67
68    #[inline(always)]
69    unsafe fn unlock(&self) {
70        let owner_id = self.owner_id.swap(0, Ordering::Release);
71        let current_id = api::ax_current_task_id();
72        assert_eq!(
73            owner_id, current_id,
74            "Thread({current_id}) tried to release mutex it doesn't own",
75        );
76        // wake up one waiting thread.
77        api::ax_wait_queue_wake(&self.wq, 1);
78    }
79
80    #[inline(always)]
81    fn is_locked(&self) -> bool {
82        self.owner_id.load(Ordering::Relaxed) != 0
83    }
84}
85
86/// An alias of [`lock_api::Mutex`].
87pub type Mutex<T> = lock_api::Mutex<RawMutex, T>;
88/// An alias of [`lock_api::MutexGuard`].
89pub type MutexGuard<'a, T> = lock_api::MutexGuard<'a, RawMutex, T>;
axstd/sync/mutex.rs

axstd/sync/
mutex.rs
