Skip to main content

ostd/sync/
mutex.rs

1// SPDX-License-Identifier: MPL-2.0
2
3use core::{
4    cell::UnsafeCell,
5    fmt,
6    ops::{Deref, DerefMut},
7    sync::atomic::{AtomicBool, Ordering},
8};
9
10use super::WaitQueue;
11
12/// A mutex with waitqueue.
13pub struct Mutex<T: ?Sized> {
14    lock: AtomicBool,
15    queue: WaitQueue,
16    val: UnsafeCell<T>,
17}
18
19impl<T> Mutex<T> {
20    /// Creates a new mutex.
21    pub const fn new(val: T) -> Self {
22        Self {
23            lock: AtomicBool::new(false),
24            queue: WaitQueue::new(),
25            val: UnsafeCell::new(val),
26        }
27    }
28}
29
30impl<T: ?Sized> Mutex<T> {
31    /// Acquires the mutex.
32    ///
33    /// This method runs in a block way until the mutex can be acquired.
34    #[track_caller]
35    pub fn lock(&self) -> MutexGuard<'_, T> {
36        self.queue.wait_until(|| self.try_lock())
37    }
38
39    /// Tries to acquire the mutex immediately.
40    pub fn try_lock(&self) -> Option<MutexGuard<'_, T>> {
41        self.acquire_lock()
42            // We cannot use `then_some`. Otherwise, the possible dropping of the temporary
43            // guard will cause an unexpected unlock.
44            // SAFETY: The lock is successfully acquired when creating the guard.
45            .then(|| unsafe { MutexGuard::new(self) })
46    }
47
48    /// Returns a mutable reference to the underlying data.
49    ///
50    /// This method is zero-cost: By holding a mutable reference to the lock, the compiler has
51    /// already statically guaranteed that access to the data is exclusive.
52    pub fn get_mut(&mut self) -> &mut T {
53        self.val.get_mut()
54    }
55
56    /// Releases the mutex and wake up one thread which is blocked on this mutex.
57    ///
58    /// # Safety
59    ///
60    /// The caller must ensure that it has acquired the lock.
61    unsafe fn unlock(&self) {
62        // SAFETY: The safety is upheld by the caller.
63        unsafe { self.release_lock() };
64        self.queue.wake_one();
65    }
66
67    fn acquire_lock(&self) -> bool {
68        self.lock
69            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
70            .is_ok()
71    }
72
73    /// # Safety
74    ///
75    /// The caller must ensure that it has acquired the lock.
76    unsafe fn release_lock(&self) {
77        self.lock.store(false, Ordering::Release);
78    }
79}
80
81impl<T: ?Sized + fmt::Debug> fmt::Debug for Mutex<T> {
82    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
83        fmt::Debug::fmt(&self.val, f)
84    }
85}
86
87// SAFETY: Only a single lock holder is permitted to access the inner data of `Mutex`.
88unsafe impl<T: ?Sized + Send> Send for Mutex<T> {}
89unsafe impl<T: ?Sized + Send> Sync for Mutex<T> {}
90
91/// A guard that provides exclusive access to the data protected by a [`Mutex`].
92#[clippy::has_significant_drop]
93#[must_use]
94pub struct MutexGuard<'a, T: ?Sized> {
95    mutex: &'a Mutex<T>,
96}
97
98impl<'a, T: ?Sized> MutexGuard<'a, T> {
99    /// # Safety
100    ///
101    /// The caller must ensure that the given reference of [`Mutex`] lock has been successfully
102    /// acquired in the current context. When the created [`MutexGuard`] is dropped, it will unlock
103    /// the [`Mutex`].
104    unsafe fn new(mutex: &'a Mutex<T>) -> Self {
105        MutexGuard { mutex }
106    }
107
108    /// Returns the [`Mutex`] associated with this guard.
109    pub fn get_lock(guard: &Self) -> &'a Mutex<T> {
110        guard.mutex
111    }
112}
113
114impl<T: ?Sized> Deref for MutexGuard<'_, T> {
115    type Target = T;
116
117    fn deref(&self) -> &Self::Target {
118        // SAFETY: The mutex is locked, which provides exclusive access to the data.
119        unsafe { &*self.mutex.val.get() }
120    }
121}
122
123impl<T: ?Sized> DerefMut for MutexGuard<'_, T> {
124    fn deref_mut(&mut self) -> &mut Self::Target {
125        // SAFETY: The mutex is locked, which provides exclusive access to the data.
126        unsafe { &mut *self.mutex.val.get() }
127    }
128}
129
130impl<T: ?Sized> Drop for MutexGuard<'_, T> {
131    fn drop(&mut self) {
132        // SAFETY: The mutex is locked.
133        unsafe { self.mutex.unlock() };
134    }
135}
136
137impl<T: ?Sized + fmt::Debug> fmt::Debug for MutexGuard<'_, T> {
138    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
139        fmt::Debug::fmt(&**self, f)
140    }
141}
142
143impl<T: ?Sized> !Send for MutexGuard<'_, T> {}
144unsafe impl<T: ?Sized + Sync> Sync for MutexGuard<'_, T> {}
145
146#[cfg(ktest)]
147mod test {
148    use super::*;
149    use crate::prelude::*;
150
151    // A regression test for a bug fixed in [#1279](https://github.com/asterinas/asterinas/pull/1279).
152    #[ktest]
153    fn try_lock_does_not_unlock() {
154        let lock = Mutex::new(0);
155        assert!(!lock.lock.load(Ordering::Relaxed));
156
157        // A successful lock
158        let guard1 = lock.lock();
159        assert!(lock.lock.load(Ordering::Relaxed));
160
161        // A failed `try_lock` won't drop the lock
162        assert!(lock.try_lock().is_none());
163        assert!(lock.lock.load(Ordering::Relaxed));
164
165        // Ensure the lock is held until here
166        drop(guard1);
167    }
168}