[PATCH] Rust: Implement a unique reference type URef supplementing ARef

Oliver Mangold
2025-02-21T08:04:43Z

For usage with block-mq, we found that having a variant of ARef which is guaranteed to be unique is useful. As chances are it is useful in general, I implemented it as kernel::types::URef. The difference between ARef and URef is basically the same as between Arc and UniqueArc.

Signed-off-by: Oliver Mangold <oliver.mangold@pm.me>
---
rust/kernel/block/mq/request.rs | 97 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++------------------------------------
1 file changed, 61 insertions(+), 36 deletions(-)

diff --git a/rust/kernel/types.rs b/rust/kernel/types.rs
index 55ddd50e8aaa..80dc9edef1b9 100644
--- a/rust/kernel/types.rs
+++ b/rust/kernel/types.rs
@@ -543,6 +543,12 @@ fn from(b: &T) -> Self {
}
}

+impl<T: UniqueRefCounted> From<URef<T>> for ARef<T> {
+ fn from(b: URef<T>) -> Self {
+ UniqueRefCounted::unique_to_shared(b)
+ }
+}
+
impl<T: AlwaysRefCounted> Drop for ARef<T> {
fn drop(&mut self) {
// SAFETY: The type invariants guarantee that the `ARef` owns the reference we're about to
@@ -551,6 +557,119 @@ fn drop(&mut self) {
}
}

+/// Types that are `AlwaysRefCounted` and can be safely converted to an `URef`
+///
+pub unsafe trait UniqueRefCounted: AlwaysRefCounted + Sized {
+ /// Checks if the the [ARef] is unique and convert it
+ /// to an [URef] it that is that case.
+ /// Otherwise it returns again an [ARef] to the same
+ /// underlying object.
+ fn try_shared_to_unique(this: ARef<Self>) -> Result<URef<Self>,ARef<Self>>;
+ /// Converts the [URef] into an [ARef].
+ fn unique_to_shared(this: URef<Self>) -> ARef<Self>;
+}
+
+/// An unique owned reference to an always-reference-counted object.
+///
+/// It works the same ways as [`ARef`] but ensures that the reference is unique
+/// and thus can be dereferenced mutably.
+///
+/// # Invariants
+///
+/// The pointer stored in `ptr` is non-null and valid for the lifetime of the [`ARef`] instance. In
+/// particular, the [`ARef`] instance owns an increment on the underlying object's reference count.
+pub struct URef<T: UniqueRefCounted> {
+ ptr: NonNull<T>,
+ _p: PhantomData<T>,
+}
+
+// SAFETY: It is safe to send `URef<T>` to another thread when the underlying `T` is `Sync` because
+// it effectively means sharing `&T` (which is safe because `T` is `Sync`); additionally, it needs
+// `T` to be `Send` because any thread that has an `URef<T>` may ultimately access `T` using a
+// mutable reference, for example, when the reference count reaches zero and `T` is dropped.
+unsafe impl<T: UniqueRefCounted + Sync + Send> Send for URef<T> {}
+
+// SAFETY: It is safe to send `&URef<T>` to another thread when the underlying `T` is `Sync`
+// because it effectively means sharing `&T` (which is safe because `T` is `Sync`); additionally,
+// it needs `T` to be `Send` because any thread that has a `&URef<T>` may clone it and get an
+// `URef<T>` on that thread, so the thread may ultimately access `T` using a mutable reference, for
+// example, when the reference count reaches zero and `T` is dropped.
+unsafe impl<T: UniqueRefCounted + Sync + Send> Sync for URef<T> {}
+
+impl<T: UniqueRefCounted> URef<T> {
+ /// Creates a new instance of [`URef`].
+ ///
+ /// It takes over an increment of the reference count on the underlying object.
+ ///
+ /// # Safety
+ ///
+ /// Callers must ensure that the reference count is set to such a value
+ /// that a call to AlwaysRefCounted::dec_ref() will release the underlying object
+ /// in the way which is expected when the last reference is dropped.
+ /// Callers must not use the underlying object anymore --
+ /// it is only safe to do so via the newly created [`URef`].
+ pub unsafe fn from_raw(ptr: NonNull<T>) -> Self {
+ // INVARIANT: The safety requirements guarantee that the new instance now owns the
+ // increment on the refcount.
+ Self {
+ ptr,
+ _p: PhantomData,
+ }
+ }
+
+ /// Consumes the `URef`, returning a raw pointer.
+ ///
+ /// This function does not change the refcount. After calling this function, the caller is
+ /// responsible for the refcount previously managed by the `URef`.
+ pub fn into_raw(me: Self) -> NonNull<T> {
+ ManuallyDrop::new(me).ptr
+ }
+}
+
+impl<T: UniqueRefCounted> Deref for URef<T> {
+ type Target = T;
+
+ fn deref(&self) -> &Self::Target {
+ // SAFETY: The type invariants guarantee that the object is valid.
+ unsafe { self.ptr.as_ref() }
+ }
+}
+
+impl<T: UniqueRefCounted> DerefMut for URef<T> {
+
+ fn deref_mut(&mut self) -> &mut Self::Target {
+ // SAFETY: The type invariants guarantee that the object is valid.
+ unsafe { self.ptr.as_mut() }
+ }
+}
+
+impl<T: UniqueRefCounted> From<&T> for URef<T> {
+ /// Converts the [URef] into an [ARef] by calling [UniqueRefCounted::unique_to_shared] on it.
+ fn from(b: &T) -> Self {
+ b.inc_ref();
+ // SAFETY: We just incremented the refcount above.
+ unsafe { Self::from_raw(NonNull::from(b)) }
+ }
+}
+
+impl<T: UniqueRefCounted> TryFrom<ARef<T>> for URef<T> {
+ type Error = ARef<T>;
+ /// Tries to convert the [ARef] to an [URef] by calling [UniqueRefCounted::try_shared_to_unique].
+ /// In case the [ARef] is not unique it returns again an [ARef] to the same
+ /// underlying object.
+ fn try_from(b: ARef<T>) -> Result<URef<T>,Self::Error> {
+ UniqueRefCounted::try_shared_to_unique(b)
+ }
+}
+
+impl<T: UniqueRefCounted> Drop for URef<T> {
+ fn drop(&mut self) {
+ // SAFETY: The type invariants guarantee that the `URef` owns the reference we're about to
+ // decrement.
+ unsafe { T::dec_ref(self.ptr) };
+ }
+}
+
/// A sum type that always holds either a value of type `L` or `R`.
///
/// # Examples
---