nuclei/sync/
reentrant_rwlock.rs

1//! An asynchronous reader-writer lock its holders may lock again, keyed
2//! by [`Owner`].
3//!
4//! The lock counts owners, not calls. An owner that already reads reads
5//! again at once, even with a writer owed a handoff, because that writer
6//! waits for this owner's read anyway; an owner that writes may also read.
7//! An owner that reads without writing and asks to write is refused with
8//! [`ReadHeldByOwner`]: the write would wait for its own read. So an owner
9//! never deadlocks on itself.
10//!
11//! Each reading owner has a record (owner id, read depth) in a lock-free
12//! list that only grows: a record is reused once its owner stops reading,
13//! so the list is as long as the most owners that ever read at once.
14
15#![allow(unsafe_code)]
16
17use core::fmt;
18use core::future::Future;
19use core::ops::Deref;
20use core::pin::Pin;
21use core::ptr;
22use core::task::{Context, Poll};
23
24use super::internal::{AtomicPtr, AtomicU64, AtomicUsize, Ordering, UnsafeCell};
25use super::owner::Owner;
26use super::raw_rwlock::{READ, RawRwLock, RwWait, WRITE};
27
28/// One reading owner: its id (zero when free) and how many reads it holds.
29struct Record {
30    owner: AtomicU64,
31    reads: AtomicUsize,
32    /// Set before the record is published and never again.
33    next: *mut Record,
34}
35
36struct Records {
37    head: AtomicPtr<Record>,
38}
39
40impl Records {
41    loom_const_fn! {
42        fn new() -> Self {
43            Self { head: AtomicPtr::new(ptr::null_mut()) }
44        }
45    }
46
47    fn iter(&self) -> impl Iterator<Item = &Record> {
48        let mut cursor = self.head.load(Ordering::Acquire);
49        core::iter::from_fn(move || {
50            // SAFETY: records are freed only when the list is dropped.
51            let record = unsafe { cursor.as_ref() }?;
52            cursor = record.next;
53            Some(record)
54        })
55    }
56
57    /// The record `id` holds. Only an owner stores its own id, so only
58    /// that owner can find it.
59    fn find(&self, id: u64) -> Option<&Record> {
60        self.iter()
61            .find(|record| record.owner.load(Ordering::Relaxed) == id)
62    }
63
64    /// A free record, now held by `id`.
65    fn claim(&self, id: u64) -> &Record {
66        if let Some(record) = self.iter().find(|record| {
67            record.owner.load(Ordering::Relaxed) == 0
68                && record
69                    .owner
70                    .compare_exchange(0, id, Ordering::Acquire, Ordering::Relaxed)
71                    .is_ok()
72        }) {
73            return record;
74        }
75        let fresh = Box::into_raw(Box::new(Record {
76            owner: AtomicU64::new(id),
77            reads: AtomicUsize::new(0),
78            next: ptr::null_mut(),
79        }));
80        let mut head = self.head.load(Ordering::Relaxed);
81        loop {
82            // SAFETY: `fresh` is unpublished, so this thread owns it.
83            unsafe { (*fresh).next = head };
84            match self
85                .head
86                .compare_exchange_weak(head, fresh, Ordering::Release, Ordering::Relaxed)
87            {
88                // SAFETY: published; it lives until the list is dropped.
89                Ok(_) => return unsafe { &*fresh },
90                Err(actual) => head = actual,
91            }
92        }
93    }
94}
95
96impl Drop for Records {
97    fn drop(&mut self) {
98        let mut cursor = self.head.load(Ordering::Acquire);
99        while !cursor.is_null() {
100            // SAFETY: `&mut self` leaves the list unreachable to others.
101            let record = unsafe { Box::from_raw(cursor) };
102            cursor = record.next;
103        }
104    }
105}
106
107/// A reader-writer lock its holders may lock again, with the bounded
108/// bypass of [`RwLock`](super::RwLock).
109///
110/// Both guards give `&T`, since one owner may hold several at once; the
111/// write side excludes every other owner, so mutate through interior
112/// mutability that is `Sync` (an atomic, a [`Mutex`](super::Mutex)), as
113/// readers on other threads share the value. Guards and lock futures
114/// borrow their [`Owner`] and stay on its thread.
115///
116/// ```
117/// use nuclei::sync::{Owner, ReentrantRwLock};
118/// use std::sync::atomic::{AtomicU32, Ordering};
119///
120/// let lock = ReentrantRwLock::new(AtomicU32::new(0));
121/// let me = Owner::new();
122/// let write = lock.try_write(&me).expect("free");
123/// let read = lock.try_read(&me).expect("a writer may also read");
124/// write.store(7, Ordering::Relaxed);
125/// drop(write);
126/// assert!(lock.try_write(&Owner::new()).is_none(), "the read is still held");
127/// assert_eq!(read.load(Ordering::Relaxed), 7);
128/// ```
129pub struct ReentrantRwLock<T: ?Sized> {
130    raw: RawRwLock,
131    /// The writing owner's id, or zero; only the writer stores its own id.
132    writer: AtomicU64,
133    /// How many write guards the writer holds; touched only by the writer.
134    write_depth: UnsafeCell<usize>,
135    readers: Records,
136    value: core::cell::UnsafeCell<T>,
137}
138
139// SAFETY: the value moves with the lock.
140unsafe impl<T: ?Sized + Send> Send for ReentrantRwLock<T> {}
141// SAFETY: readers on several threads share `&T`, so `T: Sync`; the writer
142// gets nothing more than `&T`.
143unsafe impl<T: ?Sized + Send + Sync> Sync for ReentrantRwLock<T> {}
144
145impl<T> ReentrantRwLock<T> {
146    loom_const_fn! {
147        /// An unlocked lock holding `value`.
148        pub fn new(value: T) -> Self {
149            Self {
150                raw: RawRwLock::new(),
151                writer: AtomicU64::new(0),
152                write_depth: UnsafeCell::new(0),
153                readers: Records::new(),
154                value: core::cell::UnsafeCell::new(value),
155            }
156        }
157    }
158
159    /// Consumes the lock and returns its value.
160    pub fn into_inner(self) -> T {
161        self.value.into_inner()
162    }
163}
164
165impl<T: ?Sized> ReentrantRwLock<T> {
166    /// Reads if `owner` already reads or writes, or if no writer holds the
167    /// lock and no handoff is owed.
168    pub fn try_read<'a>(&'a self, owner: &'a Owner) -> Option<ReentrantReadGuard<'a, T>> {
169        let record = match self.reenter_read(owner) {
170            Some(record) => record,
171            None if self.raw.try_kind(READ) => self.enter_read(owner),
172            None => return None,
173        };
174        Some(ReentrantReadGuard::new(self, owner, record))
175    }
176
177    /// Writes if `owner` already writes, or if nobody holds the lock and
178    /// no handoff is owed. `None` too while `owner` reads without writing.
179    pub fn try_write<'a>(&'a self, owner: &'a Owner) -> Option<ReentrantWriteGuard<'a, T>> {
180        if self.reenter_write(owner) {
181            return Some(ReentrantWriteGuard::new(self, owner));
182        }
183        (self.readers.find(owner.id()).is_none() && self.raw.try_kind(WRITE)).then(|| {
184            self.enter_write(owner);
185            ReentrantWriteGuard::new(self, owner)
186        })
187    }
188
189    /// Waits to read, at once if `owner` already reads or writes. A
190    /// contended wait is lock-free, not wait-free: see
191    /// [Progress](super#progress).
192    pub fn read<'a>(&'a self, owner: &'a Owner) -> ReentrantRead<'a, T> {
193        ReentrantRead {
194            lock: self,
195            owner,
196            wait: RwWait::new(),
197        }
198    }
199
200    /// Waits to write, at once if `owner` already writes. Fails with
201    /// [`ReadHeldByOwner`] when `owner` reads without writing: the write
202    /// would wait for that read, which only `owner` can end. A contended
203    /// wait is lock-free, not wait-free: see [Progress](super#progress).
204    pub fn write<'a>(&'a self, owner: &'a Owner) -> ReentrantWrite<'a, T> {
205        ReentrantWrite {
206            lock: self,
207            owner,
208            wait: RwWait::new(),
209        }
210    }
211
212    /// The value, borrowed mutably; `&mut self` proves no guard exists.
213    pub fn get_mut(&mut self) -> &mut T {
214        self.value.get_mut()
215    }
216
217    fn reenter_read(&self, owner: &Owner) -> Option<&Record> {
218        let id = owner.id();
219        if let Some(record) = self.readers.find(id) {
220            record.reads.fetch_add(1, Ordering::Relaxed);
221            return Some(record);
222        }
223        (self.writer.load(Ordering::Relaxed) == id).then(|| self.enter_read(owner))
224    }
225
226    /// Records a first read by `owner`, which holds a share or the write.
227    fn enter_read(&self, owner: &Owner) -> &Record {
228        let record = self.readers.claim(owner.id());
229        record.reads.store(1, Ordering::Relaxed);
230        record
231    }
232
233    fn read_unlock(&self, id: u64, record: &Record) {
234        if record.reads.fetch_sub(1, Ordering::Relaxed) == 1 {
235            record.owner.store(0, Ordering::Release);
236            // While the owner writes, its reads hold no share of their own.
237            if self.writer.load(Ordering::Relaxed) != id {
238                self.raw.read_unlock();
239            }
240        }
241    }
242
243    #[expect(
244        clippy::panic,
245        reason = "a count that cannot overflow in practice, checked rather than wrapped"
246    )]
247    fn reenter_write(&self, owner: &Owner) -> bool {
248        if self.writer.load(Ordering::Relaxed) != owner.id() {
249            return false;
250        }
251        // SAFETY: only the writer touches the depth, and this is it.
252        self.write_depth.with_mut(|depth| unsafe {
253            *depth = (*depth)
254                .checked_add(1)
255                .unwrap_or_else(|| panic!("a ReentrantRwLock was write-locked usize::MAX times"));
256        });
257        true
258    }
259
260    fn enter_write(&self, owner: &Owner) {
261        // SAFETY: the write side makes this caller the only writer.
262        self.write_depth.with_mut(|depth| unsafe { *depth = 1 });
263        self.writer.store(owner.id(), Ordering::Relaxed);
264    }
265
266    fn write_unlock(&self, id: u64) {
267        // SAFETY: only the writer drops a write guard.
268        let left = self.write_depth.with_mut(|depth| unsafe {
269            *depth -= 1;
270            *depth
271        });
272        if left == 0 {
273            self.writer.store(0, Ordering::Relaxed);
274            if self.readers.find(id).is_some() {
275                self.raw.downgrade();
276            } else {
277                self.raw.write_unlock();
278            }
279        }
280    }
281}
282
283impl<T: Default> Default for ReentrantRwLock<T> {
284    fn default() -> Self {
285        Self::new(T::default())
286    }
287}
288
289impl<T: ?Sized> fmt::Debug for ReentrantRwLock<T> {
290    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
291        f.debug_struct("ReentrantRwLock").finish_non_exhaustive()
292    }
293}
294
295/// One read of a [`ReentrantRwLock`]; drops one level of the owner's
296/// reads. The drop of the last level releases the read, which is
297/// lock-free, not wait-free: with waiters it returns after one pass over
298/// the queue, plus at most one for each operation of another thread that
299/// lands on the queue meanwhile (see [Progress](super#progress)).
300#[must_use = "dropping the guard releases one read at once"]
301pub struct ReentrantReadGuard<'a, T: ?Sized> {
302    lock: &'a ReentrantRwLock<T>,
303    owner: &'a Owner,
304    record: &'a Record,
305}
306
307impl<'a, T: ?Sized> ReentrantReadGuard<'a, T> {
308    fn new(lock: &'a ReentrantRwLock<T>, owner: &'a Owner, record: &'a Record) -> Self {
309        Self {
310            lock,
311            owner,
312            record,
313        }
314    }
315}
316
317impl<T: ?Sized> Deref for ReentrantReadGuard<'_, T> {
318    type Target = T;
319
320    fn deref(&self) -> &T {
321        // SAFETY: a read excludes every other owner's write, and the
322        // writer, if it is this owner, only ever has `&T` too.
323        unsafe { &*self.lock.value.get() }
324    }
325}
326
327impl<T: ?Sized> Drop for ReentrantReadGuard<'_, T> {
328    fn drop(&mut self) {
329        self.lock.read_unlock(self.owner.id(), self.record);
330    }
331}
332
333impl<T: ?Sized + fmt::Debug> fmt::Debug for ReentrantReadGuard<'_, T> {
334    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
335        fmt::Debug::fmt(&**self, f)
336    }
337}
338
339/// One write of a [`ReentrantRwLock`]; drops one level of the owner's
340/// writes. The drop of the last level releases the write, with the
341/// progress [`ReentrantReadGuard`]'s release makes.
342#[must_use = "dropping the guard releases one write at once"]
343pub struct ReentrantWriteGuard<'a, T: ?Sized> {
344    lock: &'a ReentrantRwLock<T>,
345    owner: &'a Owner,
346}
347
348impl<'a, T: ?Sized> ReentrantWriteGuard<'a, T> {
349    fn new(lock: &'a ReentrantRwLock<T>, owner: &'a Owner) -> Self {
350        Self { lock, owner }
351    }
352}
353
354impl<T: ?Sized> Deref for ReentrantWriteGuard<'_, T> {
355    type Target = T;
356
357    fn deref(&self) -> &T {
358        // SAFETY: the write excludes every other owner, and this owner's
359        // guards only ever hand out `&T`.
360        unsafe { &*self.lock.value.get() }
361    }
362}
363
364impl<T: ?Sized> Drop for ReentrantWriteGuard<'_, T> {
365    fn drop(&mut self) {
366        self.lock.write_unlock(self.owner.id());
367    }
368}
369
370impl<T: ?Sized + fmt::Debug> fmt::Debug for ReentrantWriteGuard<'_, T> {
371    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
372        fmt::Debug::fmt(&**self, f)
373    }
374}
375
376/// The future [`ReentrantRwLock::read`] returns.
377#[must_use = "futures do nothing unless polled"]
378pub struct ReentrantRead<'a, T: ?Sized> {
379    lock: &'a ReentrantRwLock<T>,
380    owner: &'a Owner,
381    wait: RwWait,
382}
383
384impl<'a, T: ?Sized> Future for ReentrantRead<'a, T> {
385    type Output = ReentrantReadGuard<'a, T>;
386
387    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<ReentrantReadGuard<'a, T>> {
388        let this = self.get_mut();
389        if !this.wait.is_queued()
390            && let Some(record) = this.lock.reenter_read(this.owner)
391        {
392            // Completed at once: it spends one operation of the poll's
393            // budget, or, the budget spent, the read goes and it yields.
394            let guard = ReentrantReadGuard::new(this.lock, this.owner, record);
395            if crate::coop::spend(cx.waker()) {
396                return Poll::Ready(guard);
397            }
398            drop(guard);
399            return Poll::Pending;
400        }
401        this.wait.poll(&this.lock.raw, READ, cx.waker()).map(|()| {
402            let record = this.lock.enter_read(this.owner);
403            ReentrantReadGuard::new(this.lock, this.owner, record)
404        })
405    }
406}
407
408impl<T: ?Sized> Drop for ReentrantRead<'_, T> {
409    fn drop(&mut self) {
410        self.wait.cancel(&self.lock.raw, READ);
411    }
412}
413
414impl<T: ?Sized> fmt::Debug for ReentrantRead<'_, T> {
415    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
416        f.debug_struct("ReentrantRead")
417            .field("owner", self.owner)
418            .finish_non_exhaustive()
419    }
420}
421
422/// Why a [`ReentrantRwLock::write`] failed: its owner reads without
423/// writing, so the write would wait for a read only that owner can end.
424/// Release the reads first, or take the write before reading.
425#[derive(Clone, Copy, Debug, PartialEq, Eq)]
426#[non_exhaustive]
427pub struct ReadHeldByOwner;
428
429impl fmt::Display for ReadHeldByOwner {
430    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
431        f.write_str("this owner holds a read of the lock; release it before writing")
432    }
433}
434
435impl std::error::Error for ReadHeldByOwner {}
436
437/// The future [`ReentrantRwLock::write`] returns.
438#[must_use = "futures do nothing unless polled"]
439pub struct ReentrantWrite<'a, T: ?Sized> {
440    lock: &'a ReentrantRwLock<T>,
441    owner: &'a Owner,
442    wait: RwWait,
443}
444
445impl<'a, T: ?Sized> Future for ReentrantWrite<'a, T> {
446    type Output = Result<ReentrantWriteGuard<'a, T>, ReadHeldByOwner>;
447
448    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
449        let this = self.get_mut();
450        if !this.wait.is_queued() {
451            // Either completes at once: it spends one operation of the poll's
452            // budget, or, the budget spent, the level goes and it yields.
453            if this.lock.reenter_write(this.owner) {
454                let guard = ReentrantWriteGuard::new(this.lock, this.owner);
455                if crate::coop::spend(cx.waker()) {
456                    return Poll::Ready(Ok(guard));
457                }
458                drop(guard);
459                return Poll::Pending;
460            }
461            if this.lock.readers.find(this.owner.id()).is_some() {
462                return if crate::coop::spend(cx.waker()) {
463                    Poll::Ready(Err(ReadHeldByOwner))
464                } else {
465                    Poll::Pending
466                };
467            }
468        }
469        this.wait.poll(&this.lock.raw, WRITE, cx.waker()).map(|()| {
470            this.lock.enter_write(this.owner);
471            Ok(ReentrantWriteGuard::new(this.lock, this.owner))
472        })
473    }
474}
475
476impl<T: ?Sized> Drop for ReentrantWrite<'_, T> {
477    fn drop(&mut self) {
478        self.wait.cancel(&self.lock.raw, WRITE);
479    }
480}
481
482impl<T: ?Sized> fmt::Debug for ReentrantWrite<'_, T> {
483    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
484        f.debug_struct("ReentrantWrite")
485            .field("owner", self.owner)
486            .finish_non_exhaustive()
487    }
488}
489
490#[cfg(all(test, not(loom)))]
491mod tests {
492    use super::*;
493    use crate::sync::MAX_BYPASS;
494    use crate::sync::test_support::Polled;
495
496    #[test]
497    fn a_reader_reads_again_while_a_writer_waits() {
498        let lock = ReentrantRwLock::new(());
499        let (me, writer) = (Owner::new(), Owner::new());
500        let first = lock.try_read(&me).expect("free");
501        let mut write = Polled::new(lock.write(&writer));
502        write.pending();
503        let second = Polled::new(lock.read(&me)).ready();
504        drop(first);
505        drop(second);
506        assert!(write.wakes() >= 1, "the last read out nudges the writer");
507        let held = write.ready().expect("the writer holds no read");
508        assert!(lock.try_read(&Owner::new()).is_none());
509        drop(held);
510    }
511
512    #[test]
513    fn an_owned_read_never_waits_behind_an_owed_writer() {
514        let lock = ReentrantRwLock::new(());
515        let (me, writer) = (Owner::new(), Owner::new());
516        let others: Vec<Owner> = (0..=MAX_BYPASS).map(|_| Owner::new()).collect();
517        let mine = lock.try_read(&me).expect("free");
518        let mut write = Polled::new(lock.write(&writer));
519        write.pending();
520        let mut other = lock.try_read(&others[0]).expect("barge");
521        for owner in &others[1..] {
522            let next = lock.try_read(owner).expect("a new owner barges");
523            drop(other);
524            other = next;
525            write.pending();
526        }
527        assert!(lock.try_read(&Owner::new()).is_none(), "the writer is owed");
528        let again = Polled::new(lock.read(&me)).ready();
529        drop((other, again, mine));
530        assert!(write.ready().is_ok());
531    }
532
533    #[test]
534    fn a_writer_reads_and_keeps_its_read_after_the_write() {
535        let lock = ReentrantRwLock::new(());
536        let (me, other) = (Owner::new(), Owner::new());
537        let write = lock.try_write(&me).expect("free");
538        let nested = lock.try_write(&me).expect("the writer re-enters");
539        let read = lock.try_read(&me).expect("a writer may also read");
540        drop((nested, write));
541        assert!(
542            lock.try_write(&other).is_none(),
543            "the read survives the write"
544        );
545        let shared = lock.try_read(&other).expect("the write became a read");
546        drop((read, shared));
547        assert!(lock.try_write(&other).is_some());
548    }
549
550    #[test]
551    fn a_reader_asking_to_write_is_refused_rather_than_deadlocked() {
552        let lock = ReentrantRwLock::new(());
553        let me = Owner::new();
554        let read = lock.try_read(&me).expect("free");
555        assert!(lock.try_write(&me).is_none());
556        assert_eq!(
557            Polled::new(lock.write(&me)).ready().err(),
558            Some(ReadHeldByOwner)
559        );
560        drop(read);
561        assert!(Polled::new(lock.write(&me)).ready().is_ok());
562    }
563}
nuclei/sync/reentrant_rwlock.rs

nuclei/sync/
reentrant_rwlock.rs
