Skip to main content

ostd/sync/
rwarc.rs

1// SPDX-License-Identifier: MPL-2.0
2
3use alloc::sync::Arc;
4use core::sync::atomic::{AtomicUsize, Ordering, fence};
5
6use super::{PreemptDisabled, RwLock, RwLockReadGuard, RwLockWriteGuard};
7
8/// A reference-counting pointer with read-write capabilities.
9///
10/// This is essentially `Arc<RwLock<T>>`, so it can provide read-write capabilities through
11/// [`RwArc::read`] and [`RwArc::write`].
12///
13/// In addition, this allows to derive another reference-counting pointer with read-only
14/// capabilities ([`RoArc`]) via [`RwArc::clone_ro`].
15///
16/// The purpose of having this type is to allow lockless (read) access to the underlying data when
17/// there is only one [`RwArc`] instance for the particular allocation (note that there can be any
18/// number of [`RoArc`] instances for that allocation). See the [`RwArc::get`] method for more
19/// details.
20pub struct RwArc<T>(Arc<Inner<T>>);
21
22/// A reference-counting pointer with read-only capabilities.
23///
24/// This type can be created from an existing [`RwArc`] using its [`RwArc::clone_ro`] method. See
25/// the type and method documentation for more details.
26pub struct RoArc<T>(Arc<Inner<T>>);
27
28struct Inner<T> {
29    data: RwLock<T>,
30    num_rw: AtomicUsize,
31}
32
33impl<T> RwArc<T> {
34    /// Creates a new `RwArc<T>`.
35    pub fn new(data: T) -> Self {
36        let inner = Inner {
37            data: RwLock::new(data),
38            num_rw: AtomicUsize::new(1),
39        };
40        Self(Arc::new(inner))
41    }
42
43    /// Acquires the read lock for immutable access.
44    pub fn read(&self) -> RwLockReadGuard<'_, T, PreemptDisabled> {
45        self.0.data.read()
46    }
47
48    /// Acquires the write lock for mutable access.
49    pub fn write(&self) -> RwLockWriteGuard<'_, T, PreemptDisabled> {
50        self.0.data.write()
51    }
52
53    /// Returns a raw pointer to the contained value.
54    pub fn as_ptr(&self) -> *const T {
55        self.0.data.as_ptr()
56    }
57
58    /// Returns an immutable reference if no other `RwArc` points to the same allocation.
59    ///
60    /// This method is cheap because it does not acquire a lock.
61    ///
62    /// It's still sound because:
63    /// - The mutable reference to `self` and the condition ensure that we are exclusively
64    ///   accessing the unique `RwArc` instance for the particular allocation.
65    /// - There may be any number of [`RoArc`]s pointing to the same allocation, but they may only
66    ///   produce immutable references to the underlying data.
67    pub fn get(&mut self) -> Option<&T> {
68        if self.0.num_rw.load(Ordering::Relaxed) > 1 {
69            return None;
70        }
71
72        // This will synchronize with `RwArc::drop` to make sure its changes are visible to us.
73        fence(Ordering::Acquire);
74
75        let data_ptr = self.0.data.as_ptr();
76
77        // SAFETY: The data is valid. During the lifetime, no one will be able to create a mutable
78        // reference to the data, so it's okay to create an immutable reference like the one below.
79        Some(unsafe { &*data_ptr })
80    }
81
82    /// Clones a [`RoArc`] that points to the same allocation.
83    pub fn clone_ro(&self) -> RoArc<T> {
84        RoArc(self.0.clone())
85    }
86}
87
88impl<T> Clone for RwArc<T> {
89    fn clone(&self) -> Self {
90        let inner = self.0.clone();
91
92        // Note that overflowing the counter will make it unsound. But not to worry: the above
93        // `Arc::clone` must have already aborted the kernel before this happens.
94        inner.num_rw.fetch_add(1, Ordering::Relaxed);
95
96        Self(inner)
97    }
98}
99
100impl<T> Drop for RwArc<T> {
101    fn drop(&mut self) {
102        self.0.num_rw.fetch_sub(1, Ordering::Release);
103    }
104}
105
106impl<T: Clone> RwArc<T> {
107    /// Returns the contained value by cloning it.
108    pub fn get_cloned(&self) -> T {
109        let guard = self.read();
110        guard.clone()
111    }
112}
113
114impl<T> RoArc<T> {
115    /// Acquires the read lock for immutable access.
116    pub fn read(&self) -> RwLockReadGuard<'_, T, PreemptDisabled> {
117        self.0.data.read()
118    }
119}
120
121#[cfg(ktest)]
122mod test {
123    use super::*;
124    use crate::prelude::*;
125
126    #[ktest]
127    fn lockless_get() {
128        let mut rw1 = RwArc::new(1u32);
129        assert_eq!(rw1.get(), Some(1).as_ref());
130
131        let _ro = rw1.clone_ro();
132        assert_eq!(rw1.get(), Some(1).as_ref());
133
134        let rw2 = rw1.clone();
135        assert_eq!(rw1.get(), None);
136
137        drop(rw2);
138        assert_eq!(rw1.get(), Some(1).as_ref());
139    }
140}