Skip to main content

nuclei/sync/local/
mutex.rs

1//! A first-come, first-served mutex for one shard: a one-permit
2//! [`Semaphore`](super::Semaphore) around a value.
3//!
4//! An unlock with waiters hands the lock to the front waiter at once, so the
5//! lock is never free while someone waits and nobody barges. It does not
6//! poison: a panic in a holder unlocks it.
7
8#![allow(unsafe_code)]
9
10use core::cell::UnsafeCell;
11use core::fmt;
12use core::future::Future;
13use core::ops::{Deref, DerefMut};
14use core::pin::Pin;
15use core::task::{Context, Poll};
16
17use super::semaphore::{Acquire, Semaphore};
18use crate::sync::TryLockError;
19
20/// A mutual-exclusion lock whose wait is a future, for the tasks of one shard.
21///
22/// Locking in arrival order: a [`lock`](Self::lock) that has to wait is
23/// granted the lock by the unlock that frees it, and [`try_lock`](Self::try_lock)
24/// fails while anyone waits. It is `!Send` and `!Sync`: share it between the
25/// tasks of a shard with an `Rc`.
26///
27/// ```
28/// use nuclei::sync::local::Mutex;
29///
30/// let counter = Mutex::new(0u32);
31/// *counter.try_lock().expect("uncontended") += 1;
32/// assert_eq!(counter.into_inner(), 1);
33/// ```
34pub struct Mutex<T: ?Sized> {
35    sem: Semaphore,
36    value: UnsafeCell<T>,
37}
38
39impl<T> Mutex<T> {
40    /// An unlocked mutex holding `value`.
41    pub const fn new(value: T) -> Self {
42        Self {
43            sem: Semaphore::new(1),
44            value: UnsafeCell::new(value),
45        }
46    }
47
48    /// Consumes the mutex and returns its value.
49    pub fn into_inner(self) -> T {
50        self.value.into_inner()
51    }
52}
53
54impl<T: ?Sized> Mutex<T> {
55    /// Waits for the lock, and answers a guard that unlocks when dropped.
56    pub fn lock(&self) -> Lock<'_, T> {
57        Lock {
58            mutex: self,
59            acquire: self.sem.acquire(),
60        }
61    }
62
63    /// Locks the mutex if it is free and nobody waits.
64    ///
65    /// # Errors
66    ///
67    /// [`TryLockError`] when the mutex is locked or owed to a waiter.
68    pub fn try_lock(&self) -> Result<MutexGuard<'_, T>, TryLockError> {
69        match self.sem.try_acquire() {
70            Ok(permit) => {
71                permit.forget();
72                Ok(MutexGuard { mutex: self })
73            }
74            Err(_) => Err(TryLockError(())),
75        }
76    }
77
78    /// The value, through a unique borrow: no lock is needed.
79    pub fn get_mut(&mut self) -> &mut T {
80        self.value.get_mut()
81    }
82}
83
84impl<T> From<T> for Mutex<T> {
85    fn from(value: T) -> Self {
86        Self::new(value)
87    }
88}
89
90impl<T: Default> Default for Mutex<T> {
91    fn default() -> Self {
92        Self::new(T::default())
93    }
94}
95
96impl<T: ?Sized + fmt::Debug> fmt::Debug for Mutex<T> {
97    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
98        let mut d = f.debug_struct("Mutex");
99        match self.try_lock() {
100            Ok(inner) => d.field("data", &&*inner),
101            Err(_) => d.field("data", &format_args!("<locked>")),
102        };
103        d.finish()
104    }
105}
106
107/// The wait of [`Mutex::lock`].
108///
109/// Cancel safe: dropping it gives up its place in line; if the lock was
110/// already handed to it, the lock passes to the next waiter.
111#[must_use = "futures do nothing unless polled"]
112pub struct Lock<'a, T: ?Sized> {
113    mutex: &'a Mutex<T>,
114    acquire: Acquire<'a>,
115}
116
117impl<'a, T: ?Sized> Future for Lock<'a, T> {
118    type Output = MutexGuard<'a, T>;
119
120    #[expect(
121        clippy::unreachable,
122        reason = "the mutex's semaphore is never closed, so its acquire never fails"
123    )]
124    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
125        let this = self.get_mut();
126        match Pin::new(&mut this.acquire).poll(cx) {
127            Poll::Ready(Ok(permit)) => {
128                permit.forget();
129                Poll::Ready(MutexGuard { mutex: this.mutex })
130            }
131            Poll::Ready(Err(_)) => unreachable!("the mutex's semaphore is never closed"),
132            Poll::Pending => Poll::Pending,
133        }
134    }
135}
136
137impl<T: ?Sized> fmt::Debug for Lock<'_, T> {
138    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
139        f.debug_struct("Lock").finish_non_exhaustive()
140    }
141}
142
143/// Holds a [`Mutex`] locked; unlocks it when dropped.
144#[must_use = "dropping the guard unlocks the mutex at once"]
145pub struct MutexGuard<'a, T: ?Sized> {
146    mutex: &'a Mutex<T>,
147}
148
149impl<T: ?Sized> Deref for MutexGuard<'_, T> {
150    type Target = T;
151
152    fn deref(&self) -> &T {
153        // SAFETY: this guard holds the mutex's only permit.
154        unsafe { &*self.mutex.value.get() }
155    }
156}
157
158impl<T: ?Sized> DerefMut for MutexGuard<'_, T> {
159    fn deref_mut(&mut self) -> &mut T {
160        // SAFETY: as in `deref`, and `&mut self` makes this borrow unique.
161        unsafe { &mut *self.mutex.value.get() }
162    }
163}
164
165impl<T: ?Sized> Drop for MutexGuard<'_, T> {
166    fn drop(&mut self) {
167        self.mutex.sem.release(1);
168    }
169}
170
171impl<T: ?Sized + fmt::Debug> fmt::Debug for MutexGuard<'_, T> {
172    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
173        fmt::Debug::fmt(&**self, f)
174    }
175}