1use core::sync::atomic::{AtomicU64, Ordering};
4
5use arceos_api::task::{self as api, AxWaitQueueHandle};
6
7pub struct RawMutex {
13 wq: AxWaitQueueHandle,
14 owner_id: AtomicU64,
15}
16
17impl RawMutex {
18 #[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 #[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 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 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 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 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
91pub type Mutex<T> = lock_api::Mutex<RawMutex, T>;
93pub type MutexGuard<'a, T> = lock_api::MutexGuard<'a, RawMutex, T>;