nuclei/sync/
reentrant_mutex.rs

1//! A fair asynchronous mutex its holder may lock again, keyed by
2//! [`Owner`].
3
4#![allow(unsafe_code)]
5
6use core::fmt;
7use core::future::Future;
8use core::marker::PhantomData;
9use core::ops::Deref;
10use core::pin::Pin;
11use core::task::{Context, Poll};
12
13use super::internal::{AtomicU64, Ordering, UnsafeCell};
14use super::owner::Owner;
15use super::raw_semaphore::{PermitWait, RawSemaphore};
16
17/// A mutex the owner holding it can lock again without waiting.
18///
19/// Each lock by the holding owner deepens a count; the mutex unlocks when
20/// the last of its guards drops, and then wakes the oldest waiting
21/// owner, with the bounded bypass of [`Mutex`](super::Mutex).
22///
23/// A guard gives `&T`, since two guards of one owner are alive at once;
24/// mutate through a `Cell` or `RefCell` inside. Guards and lock futures
25/// borrow their [`Owner`], so they stay on its thread (see the `Owner`
26/// docs): a task that holds or awaits one is not `Send` meanwhile.
27///
28/// ```
29/// use nuclei::sync::{Owner, ReentrantMutex};
30/// use std::cell::Cell;
31///
32/// let lock = ReentrantMutex::new(Cell::new(0));
33/// let me = Owner::new();
34/// let a = lock.try_lock(&me).expect("free");
35/// let b = lock.try_lock(&me).expect("re-entered");
36/// b.set(a.get() + 1);
37/// drop((a, b));
38/// assert_eq!(lock.into_inner().get(), 1);
39/// ```
40pub struct ReentrantMutex<T: ?Sized> {
41    sem: RawSemaphore,
42    /// The holding owner's id, or zero. Only the holder stores its own id,
43    /// so a holder reading its own id back is the only way to match.
44    holder: AtomicU64,
45    /// How many guards the holder has; touched only by the holder.
46    depth: UnsafeCell<usize>,
47    value: core::cell::UnsafeCell<T>,
48}
49
50// SAFETY: the value moves with the lock.
51unsafe impl<T: ?Sized + Send> Send for ReentrantMutex<T> {}
52// SAFETY: the semaphore admits one owner at a time and every guard of that
53// owner is on its thread (see `Owner`), so `&T` never reaches two threads
54// at once and `T: Sync` is not needed.
55unsafe impl<T: ?Sized + Send> Sync for ReentrantMutex<T> {}
56
57impl<T> ReentrantMutex<T> {
58    loom_const_fn! {
59        /// An unlocked mutex holding `value`.
60        pub fn new(value: T) -> Self {
61            Self {
62                sem: RawSemaphore::new(1),
63                holder: AtomicU64::new(0),
64                depth: UnsafeCell::new(0),
65                value: core::cell::UnsafeCell::new(value),
66            }
67        }
68    }
69
70    /// Consumes the mutex and returns its value.
71    pub fn into_inner(self) -> T {
72        self.value.into_inner()
73    }
74}
75
76impl<T: ?Sized> ReentrantMutex<T> {
77    /// Locks the mutex if `owner` holds it, or if it is free and no
78    /// handoff is owed.
79    pub fn try_lock<'a>(&'a self, owner: &'a Owner) -> Option<ReentrantMutexGuard<'a, T>> {
80        if self.reenter(owner) || self.sem.try_acquire(1) && self.enter(owner) {
81            return Some(ReentrantMutexGuard::new(self));
82        }
83        None
84    }
85
86    /// Waits for the lock, at once if `owner` already holds it. A
87    /// contended wait is lock-free, not wait-free: see
88    /// [Progress](super#progress).
89    pub fn lock<'a>(&'a self, owner: &'a Owner) -> ReentrantMutexLock<'a, T> {
90        ReentrantMutexLock {
91            mutex: self,
92            owner,
93            wait: PermitWait::new(),
94        }
95    }
96
97    /// The value, borrowed mutably; `&mut self` proves no guard exists.
98    pub fn get_mut(&mut self) -> &mut T {
99        self.value.get_mut()
100    }
101
102    /// Deepens the hold if `owner` has it.
103    #[expect(
104        clippy::panic,
105        reason = "a count that cannot overflow in practice, checked rather than wrapped"
106    )]
107    fn reenter(&self, owner: &Owner) -> bool {
108        if self.holder.load(Ordering::Relaxed) != owner.id() {
109            return false;
110        }
111        // SAFETY: only the holder touches the depth, and this is it.
112        self.depth.with_mut(|depth| unsafe {
113            *depth = (*depth)
114                .checked_add(1)
115                .unwrap_or_else(|| panic!("a ReentrantMutex was locked usize::MAX times"));
116        });
117        true
118    }
119
120    /// Records `owner` as the holder after the permit was taken.
121    fn enter(&self, owner: &Owner) -> bool {
122        // SAFETY: the permit makes this caller the only holder.
123        self.depth.with_mut(|depth| unsafe { *depth = 1 });
124        self.holder.store(owner.id(), Ordering::Relaxed);
125        true
126    }
127
128    fn unlock(&self) {
129        // SAFETY: only the holder drops a guard, and only it touches the
130        // depth.
131        let left = self.depth.with_mut(|depth| unsafe {
132            *depth -= 1;
133            *depth
134        });
135        if left == 0 {
136            self.holder.store(0, Ordering::Relaxed);
137            self.sem.release(1);
138        }
139    }
140}
141
142impl<T: Default> Default for ReentrantMutex<T> {
143    fn default() -> Self {
144        Self::new(T::default())
145    }
146}
147
148impl<T: ?Sized> fmt::Debug for ReentrantMutex<T> {
149    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
150        f.debug_struct("ReentrantMutex").finish_non_exhaustive()
151    }
152}
153
154/// Shared access to a [`ReentrantMutex`]'s value; drops one level of the
155/// hold. The drop of the last level releases the lock, which is lock-free,
156/// not wait-free: with waiters it returns after one pass over the queue,
157/// plus at most one for each operation of another thread that lands on the
158/// queue meanwhile (see [Progress](super#progress)).
159#[must_use = "dropping the guard releases one level of the lock at once"]
160pub struct ReentrantMutexGuard<'a, T: ?Sized> {
161    mutex: &'a ReentrantMutex<T>,
162    // Borrowing an `Owner`, which is not `Sync`, keeps the guard on the
163    // owner's thread.
164    _owner: PhantomData<&'a Owner>,
165}
166
167impl<'a, T: ?Sized> ReentrantMutexGuard<'a, T> {
168    fn new(mutex: &'a ReentrantMutex<T>) -> Self {
169        Self {
170            mutex,
171            _owner: PhantomData,
172        }
173    }
174}
175
176impl<T: ?Sized> Deref for ReentrantMutexGuard<'_, T> {
177    type Target = T;
178
179    fn deref(&self) -> &T {
180        // SAFETY: the holder's guards are the only ones and they are all
181        // on its thread, so shared access cannot race.
182        unsafe { &*self.mutex.value.get() }
183    }
184}
185
186impl<T: ?Sized> Drop for ReentrantMutexGuard<'_, T> {
187    fn drop(&mut self) {
188        self.mutex.unlock();
189    }
190}
191
192impl<T: ?Sized + fmt::Debug> fmt::Debug for ReentrantMutexGuard<'_, T> {
193    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
194        fmt::Debug::fmt(&**self, f)
195    }
196}
197
198/// The future [`ReentrantMutex::lock`] returns.
199#[must_use = "futures do nothing unless polled"]
200pub struct ReentrantMutexLock<'a, T: ?Sized> {
201    mutex: &'a ReentrantMutex<T>,
202    owner: &'a Owner,
203    wait: PermitWait,
204}
205
206impl<'a, T: ?Sized> Future for ReentrantMutexLock<'a, T> {
207    type Output = ReentrantMutexGuard<'a, T>;
208
209    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<ReentrantMutexGuard<'a, T>> {
210        let this = self.get_mut();
211        if !this.wait.is_queued() && this.mutex.reenter(this.owner) {
212            // Completed at once: it spends one operation of the poll's
213            // budget, or, the budget spent, the level goes and it yields.
214            let guard = ReentrantMutexGuard::new(this.mutex);
215            if crate::coop::spend(cx.waker()) {
216                return Poll::Ready(guard);
217            }
218            drop(guard);
219            return Poll::Pending;
220        }
221        this.wait.poll(&this.mutex.sem, 1, cx.waker()).map(|()| {
222            this.mutex.enter(this.owner);
223            ReentrantMutexGuard::new(this.mutex)
224        })
225    }
226}
227
228impl<T: ?Sized> Drop for ReentrantMutexLock<'_, T> {
229    fn drop(&mut self) {
230        self.wait.cancel(&self.mutex.sem, 1);
231    }
232}
233
234impl<T: ?Sized> fmt::Debug for ReentrantMutexLock<'_, T> {
235    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
236        f.debug_struct("ReentrantMutexLock")
237            .field("owner", self.owner)
238            .finish_non_exhaustive()
239    }
240}
241
242#[cfg(all(test, not(loom)))]
243mod tests {
244    use super::*;
245    use crate::sync::test_support::Polled;
246    use core::cell::RefCell;
247
248    #[test]
249    fn the_holder_reenters_and_the_lock_opens_after_its_last_guard() {
250        let lock = ReentrantMutex::new(RefCell::new(Vec::new()));
251        let (me, other) = (Owner::new(), Owner::new());
252        let outer = lock.try_lock(&me).expect("free");
253        let inner = Polled::new(lock.lock(&me)).ready();
254        inner.borrow_mut().push(1);
255        let mut waiting = Polled::new(lock.lock(&other));
256        waiting.pending();
257        assert!(lock.try_lock(&other).is_none());
258        drop(inner);
259        assert_eq!(waiting.wakes(), 0, "one level is still held");
260        let again = lock.try_lock(&me).expect("still the holder");
261        drop((outer, again));
262        assert_eq!(waiting.wakes(), 1);
263        let theirs = waiting.ready();
264        theirs.borrow_mut().push(2);
265        assert!(lock.try_lock(&me).is_none(), "the lock changed hands");
266        drop(theirs);
267        drop(waiting);
268        assert_eq!(lock.into_inner().into_inner(), vec![1, 2]);
269    }
270
271    #[test]
272    fn a_thread_owner_is_the_same_on_one_thread_and_different_on_another() {
273        let lock = ReentrantMutex::new(());
274        let here = Owner::thread();
275        let guard = lock.try_lock(&here).expect("free");
276        let again = Owner::thread();
277        assert!(lock.try_lock(&again).is_some(), "the same thread re-enters");
278        std::thread::scope(|scope| {
279            scope.spawn(|| {
280                let there = Owner::thread();
281                assert!(
282                    lock.try_lock(&there).is_none(),
283                    "another thread is another owner"
284                );
285            });
286        });
287        drop(guard);
288    }
289
290    #[test]
291    fn a_withdrawn_lock_future_leaves_no_trace() {
292        let lock = ReentrantMutex::new(());
293        let (me, other) = (Owner::new(), Owner::new());
294        let held = lock.try_lock(&me).expect("free");
295        Polled::new(lock.lock(&other)).pending();
296        drop(held);
297        assert!(lock.try_lock(&other).is_some());
298    }
299}
nuclei/sync/reentrant_mutex.rs

nuclei/sync/
reentrant_mutex.rs
