Implement RwArc

ostd/src/sync/mod.rs

@@ -7,6 +7,7 @@ mod mutex;
 // TODO: refactor this rcu implementation
 // Comment out this module since it raises lint error
 // mod rcu;
+mod rwarc;
 mod rwlock;
 mod rwmutex;
 mod spin;
@@ -17,6 +18,7 @@ pub(crate) use self::guard::GuardTransfer;
 pub use self::{
     guard::{LocalIrqDisabled, PreemptDisabled, WriteIrqDisabled},
     mutex::{ArcMutexGuard, Mutex, MutexGuard},
+    rwarc::{RoArc, RwArc},
     rwlock::{
         ArcRwLockReadGuard, ArcRwLockUpgradeableGuard, ArcRwLockWriteGuard, RwLock,
         RwLockReadGuard, RwLockUpgradeableGuard, RwLockWriteGuard,

ostd/src/sync/rwarc.rs

@@ -0,0 +1,127 @@
+// SPDX-License-Identifier: MPL-2.0
+
+use alloc::sync::Arc;
+use core::sync::atomic::{fence, AtomicUsize, Ordering};
+
+use super::{PreemptDisabled, RwLock, RwLockReadGuard, RwLockWriteGuard};
+
+/// A reference-counting pointer with read-write capabilities.
+///
+/// This is essentially `Arc<RwLock<T>>`, so it can provide read-write capabilities through
+/// [`RwArc::read`] and [`RwArc::write`].
+///
+/// In addition, this allows to derive another reference-counting pointer with read-only
+/// capabilities ([`RoArc`]) via [`RwArc::clone_ro`].
+///
+/// The purpose of having this type is to allow lockless (read) access to the underlying data when
+/// there is only one [`RwArc`] instance for the particular allocation (note that there can be any
+/// number of [`RoArc`] instances for that allocation). See the [`RwArc::get`] method for more
+/// details.
+pub struct RwArc<T>(Arc<Inner<T>>);
+
+/// A reference-counting pointer with read-only capabilities.
+///
+/// This type can be created from an existing [`RwArc`] using its [`RwArc::clone_ro`] method. See
+/// the type and method documentation for more details.
+pub struct RoArc<T>(Arc<Inner<T>>);
+
+struct Inner<T> {
+    data: RwLock<T>,
+    num_rw: AtomicUsize,
+}
+
+impl<T> RwArc<T> {
+    /// Creates a new `RwArc<T>`.
+    pub fn new(data: T) -> Self {
+        let inner = Inner {
+            data: RwLock::new(data),
+            num_rw: AtomicUsize::new(1),
+        };
+        Self(Arc::new(inner))
+    }
+
+    /// Acquires the read lock for immutable access.
+    pub fn read(&self) -> RwLockReadGuard<T, PreemptDisabled> {
+        self.0.data.read()
+    }
+
+    /// Acquires the write lock for mutable access.
+    pub fn write(&self) -> RwLockWriteGuard<T, PreemptDisabled> {
+        self.0.data.write()
+    }
+
+    /// Returns an immutable reference if no other `RwArc` points to the same allocation.
+    ///
+    /// This method is cheap because it does not acquire a lock.
+    ///
+    /// It's still sound because:
+    /// - The mutable reference to `self` and the condition ensure that we are exclusively
+    ///   accessing the unique `RwArc` instance for the particular allocation.
+    /// - There may be any number of [`RoArc`]s pointing to the same allocation, but they may only
+    ///   produce immutable references to the underlying data.
+    pub fn get(&mut self) -> Option<&T> {
+        if self.0.num_rw.load(Ordering::Relaxed) > 1 {
+            return None;
+        }
+
+        // This will synchronize with `RwArc::drop` to make sure its changes are visible to us.
+        fence(Ordering::Acquire);
+
+        let data_ptr = self.0.data.as_ptr();
+
+        // SAFETY: The data is valid. During the lifetime, no one will be able to create a mutable
+        // reference to the data, so it's okay to create an immutable reference like the one below.
+        Some(unsafe { &*data_ptr })
+    }
+
+    /// Clones a [`RoArc`] that points to the same allocation.
+    pub fn clone_ro(&self) -> RoArc<T> {
+        RoArc(self.0.clone())
+    }
+}
+
+impl<T> Clone for RwArc<T> {
+    fn clone(&self) -> Self {
+        let inner = self.0.clone();
+
+        // Note that overflowing the counter will make it unsound. But not to worry: the above
+        // `Arc::clone` must have already aborted the kernel before this happens.
+        inner.num_rw.fetch_add(1, Ordering::Relaxed);
+
+        Self(inner)
+    }
+}
+
+impl<T> Drop for RwArc<T> {
+    fn drop(&mut self) {
+        self.0.num_rw.fetch_sub(1, Ordering::Release);
+    }
+}
+
+impl<T> RoArc<T> {
+    /// Acquires the read lock for immutable access.
+    pub fn read(&self) -> RwLockReadGuard<T, PreemptDisabled> {
+        self.0.data.read()
+    }
+}
+
+#[cfg(ktest)]
+mod test {
+    use super::*;
+    use crate::prelude::*;
+
+    #[ktest]
+    fn lockless_get() {
+        let mut rw1 = RwArc::new(1u32);
+        assert_eq!(rw1.get(), Some(1).as_ref());
+
+        let _ro = rw1.clone_ro();
+        assert_eq!(rw1.get(), Some(1).as_ref());
+
+        let rw2 = rw1.clone();
+        assert_eq!(rw1.get(), None);
+
+        drop(rw2);
+        assert_eq!(rw1.get(), Some(1).as_ref());
+    }
+}

ostd/src/sync/rwlock.rs

@@ -339,6 +339,14 @@ impl<T: ?Sized, G: Guardian> RwLock<T, G> {
     pub fn get_mut(&mut self) -> &mut T {
         self.val.get_mut()
     }
+
+    /// Returns a raw pointer to the underlying data.
+    ///
+    /// This method is safe, but it's up to the caller to ensure that access to the data behind it
+    /// is still safe.
+    pub(super) fn as_ptr(&self) -> *mut T {
+        self.val.get()
+    }
 }
 
 impl<T: ?Sized + fmt::Debug, G> fmt::Debug for RwLock<T, G> {