rust: update Ref to use the kernel's refcount_t.

Saturate the refcount instead of panicking. It also leverages existing C
code instead of reimplementing functionality.

Signed-off-by: Wedson Almeida Filho <[email protected]>

Author: wedsonaf

drivers/android/process.rs

@@ -293,22 +293,24 @@ unsafe impl Sync for Process {}
 
 impl Process {
     fn new(ctx: Arc<Context>) -> Result<Ref<Self>> {
-        let mut proc_ref = Ref::try_new(Self {
-            ref_count: RefCount::new(),
-            ctx,
-            // SAFETY: `inner` is initialised in the call to `mutex_init` below.
-            inner: unsafe { Mutex::new(ProcessInner::new()) },
-            // SAFETY: `node_refs` is initialised in the call to `mutex_init` below.
-            node_refs: unsafe { Mutex::new(ProcessNodeRefs::new()) },
-        })?;
-        let process = Ref::get_mut(&mut proc_ref).ok_or(Error::EINVAL)?;
-        // SAFETY: `inner` is pinned behind the `Arc` reference.
-        let pinned = unsafe { Pin::new_unchecked(&process.inner) };
-        kernel::mutex_init!(pinned, "Process::inner");
-        // SAFETY: `node_refs` is pinned behind the `Arc` reference.
-        let pinned = unsafe { Pin::new_unchecked(&process.node_refs) };
-        kernel::mutex_init!(pinned, "Process::node_refs");
-        Ok(proc_ref)
+        Ref::try_new_and_init(
+            Self {
+                ref_count: RefCount::default(),
+                ctx,
+                // SAFETY: `inner` is initialised in the call to `mutex_init` below.
+                inner: unsafe { Mutex::new(ProcessInner::new()) },
+                // SAFETY: `node_refs` is initialised in the call to `mutex_init` below.
+                node_refs: unsafe { Mutex::new(ProcessNodeRefs::new()) },
+            },
+            |process| {
+                // SAFETY: `inner` is pinned behind the `Ref` reference.
+                let pinned = unsafe { Pin::new_unchecked(&process.inner) };
+                kernel::mutex_init!(pinned, "Process::inner");
+                // SAFETY: `node_refs` is pinned behind the `Ref` reference.
+                let pinned = unsafe { Pin::new_unchecked(&process.node_refs) };
+                kernel::mutex_init!(pinned, "Process::node_refs");
+            },
+        )
     }
 
     /// Attemps to fetch a work item from the process queue.

rust/helpers.c

@@ -130,6 +130,24 @@ void rust_helper_mutex_lock(struct mutex *lock)
 }
 EXPORT_SYMBOL_GPL(rust_helper_mutex_lock);
 
+void rust_helper_refcount_set(refcount_t *r, unsigned int n)
+{
+	refcount_set(r, n);
+}
+EXPORT_SYMBOL_GPL(rust_helper_refcount_set);
+
+void rust_helper_refcount_inc(refcount_t *r)
+{
+	refcount_inc(r);
+}
+EXPORT_SYMBOL_GPL(rust_helper_refcount_inc);
+
+bool rust_helper_refcount_dec_and_test(refcount_t *r)
+{
+	return refcount_dec_and_test(r);
+}
+EXPORT_SYMBOL_GPL(rust_helper_refcount_dec_and_test);
+
 /* We use bindgen's --size_t-is-usize option to bind the C size_t type
  * as the Rust usize type, so we can use it in contexts where Rust
  * expects a usize like slice (array) indices. usize is defined to be

rust/kernel/sync/arc.rs

@@ -7,20 +7,24 @@
 //! underlying object when it reaches zero. It is also safe to use concurrently from multiple
 //! threads.
 //!
-//! It is different from the standard library's [`Arc`] in two ways: it does not support weak
-//! references, which allows it to be smaller -- a single pointer-sized integer; it allows users to
-//! safely increment the reference count from a single reference to the underlying object.
+//! It is different from the standard library's [`Arc`] in at least three ways:
+//! 1. It does not support weak references, which allows it to be smaller: a single pointer-sized
+//!    integer;
+//! 2. It allows users to safely increment the reference count from a single reference to the
+//!    underlying object. That is, go from `&T` to [`Ref<T>`];
+//! 3. It saturates the reference count instead of aborting when it goes over a threshold.
 //!
 //! [`Arc`]: https://doc.rust-lang.org/std/sync/struct.Arc.html
 
-use crate::Result;
+use crate::{bindings, Result};
 use alloc::boxed::Box;
-use core::{
-    mem::ManuallyDrop,
-    ops::Deref,
-    ptr::NonNull,
-    sync::atomic::{fence, AtomicUsize, Ordering},
-};
+use core::{cell::UnsafeCell, convert::AsRef, mem::ManuallyDrop, ops::Deref, ptr::NonNull};
+
+extern "C" {
+    fn rust_helper_refcount_set(r: *mut bindings::refcount_t, n: u32);
+    fn rust_helper_refcount_inc(r: *mut bindings::refcount_t);
+    fn rust_helper_refcount_dec_and_test(r: *mut bindings::refcount_t) -> bool;
+}
 
 /// A reference-counted pointer to an instance of `T`.
 ///
@@ -30,7 +34,7 @@ use core::{
 /// # Invariants
 ///
 /// The value stored in [`RefCounted::get_count`] corresponds to the number of instances of [`Ref`]
-/// that point to that instance of `T`.
+/// that point to that instance of `T`. So it is always non-zero when a [`Ref`] exists.
 pub struct Ref<T: RefCounted + ?Sized> {
     ptr: NonNull<T>,
 }
@@ -49,8 +53,18 @@ unsafe impl<T: RefCounted + ?Sized + Sync + Send> Sync for Ref<T> {}
 impl<T: RefCounted> Ref<T> {
     /// Constructs a new reference counted instance of `T`.
     pub fn try_new(contents: T) -> Result<Self> {
-        let boxed = Box::try_new(contents)?;
-        boxed.get_count().count.store(1, Ordering::Relaxed);
+        Self::try_new_and_init(contents, |_| {})
+    }
+
+    /// Constructs a new reference counted instance of `T` and calls the initialisation function.
+    ///
+    /// This is useful because it provides a mutable reference to `T` at its final location.
+    pub fn try_new_and_init<U: FnOnce(&mut T)>(contents: T, init: U) -> Result<Self> {
+        let mut boxed = Box::try_new(contents)?;
+        // SAFETY: We just allocated the object, no other references exist to it, so it is safe to
+        // initialise the refcount to 1.
+        unsafe { rust_helper_refcount_set(boxed.get_count().count.get(), 1) };
+        init(&mut boxed);
         let ptr = NonNull::from(Box::leak(boxed));
         Ok(Ref { ptr })
     }
@@ -62,28 +76,14 @@ impl<T: RefCounted + ?Sized> Ref<T> {
     /// It works by incrementing the current reference count as part of constructing the new
     /// pointer.
     pub fn new_from(obj: &T) -> Self {
-        let ref_count = obj.get_count();
-        let cur = ref_count.count.fetch_add(1, Ordering::Relaxed);
-        if cur == usize::MAX {
-            panic!("Reference count overflowed");
-        }
+        // SAFETY: By the type invariant, there is necessarily a reference to the object, so it is
+        // safe to increment the refcount.
+        unsafe { rust_helper_refcount_inc(obj.get_count().count.get()) };
         Self {
             ptr: NonNull::from(obj),
         }
     }
 
-    /// Returns a mutable reference to `T` iff the reference count is one. Otherwise returns
-    /// [`None`].
-    pub fn get_mut(&mut self) -> Option<&mut T> {
-        // Synchronises with the decrement in `drop`.
-        if self.get_count().count.load(Ordering::Acquire) != 1 {
-            return None;
-        }
-        // SAFETY: Since there is only one reference, we know it isn't possible for another thread
-        // to concurrently call this.
-        Some(unsafe { self.ptr.as_mut() })
-    }
-
     /// Determines if two reference-counted pointers point to the same underlying instance of `T`.
     pub fn ptr_eq(a: &Self, b: &Self) -> bool {
         core::ptr::eq(a.ptr.as_ptr(), b.ptr.as_ptr())
@@ -126,22 +126,29 @@ impl<T: RefCounted + ?Sized> Clone for Ref<T> {
     }
 }
 
+impl<T: RefCounted + ?Sized> AsRef<T> for Ref<T> {
+    fn as_ref(&self) -> &T {
+        // SAFETY: By the type invariant, there is necessarily a reference to the object, so it is
+        // safe to dereference it.
+        unsafe { self.ptr.as_ref() }
+    }
+}
+
 impl<T: RefCounted + ?Sized> Drop for Ref<T> {
     fn drop(&mut self) {
         {
             // SAFETY: By the type invariant, there is necessarily a reference to the object.
             let obj = unsafe { self.ptr.as_ref() };
 
-            // Synchronises with the acquire below or with the acquire in `get_mut`.
-            if obj.get_count().count.fetch_sub(1, Ordering::Release) != 1 {
+            // SAFETY: Also by the type invariant, we are allowed to decrement the refcount. We
+            // also need to limit the scope of `obj` as it cannot be accessed after the refcount is
+            // decremented.
+            let is_zero = unsafe { rust_helper_refcount_dec_and_test(obj.get_count().count.get()) };
+            if !is_zero {
                 return;
             }
         }
 
-        // Synchronises with the release when decrementing above. This ensures that modifications
-        // from all previous threads/CPUs are visible to the underlying object's `drop`.
-        fence(Ordering::Acquire);
-
         // The count reached zero, we must free the memory.
         //
         // SAFETY: The pointer was initialised from the result of `Box::into_raw`.
@@ -154,7 +161,9 @@ impl<T: RefCounted + ?Sized> Drop for Ref<T> {
 /// # Safety
 ///
 /// Implementers of [`RefCounted`] must ensure that all of their constructors call
-/// [`Ref::try_new`].
+/// [`Ref::try_new`] or [`Ref::try_new_and_init`]. Constructing instances that are not wrapped in
+/// [`Ref`] results in unsafety because one could call [`Ref::new_from`], and the destructor would
+/// drop the underlying object as if it had been allocated through one of [`Ref`] constructors.
 pub unsafe trait RefCounted {
     /// Returns a pointer to the object field holds the reference count.
     fn get_count(&self) -> &RefCount;
@@ -164,21 +173,7 @@ pub unsafe trait RefCounted {
 ///
 /// It is meant to be embedded in objects to be reference-counted, with [`RefCounted::get_count`]
 /// returning a reference to it.
+#[derive(Default)]
 pub struct RefCount {
-    count: AtomicUsize,
-}
-
-impl RefCount {
-    /// Constructs a new instance of [`RefCount`].
-    pub fn new() -> Self {
-        Self {
-            count: AtomicUsize::new(1),
-        }
-    }
-}
-
-impl Default for RefCount {
-    fn default() -> Self {
-        Self::new()
-    }
+    count: UnsafeCell<bindings::refcount_t>,
 }