[PATCH v2 1/3] rust: implement `kernel::sync::Refcount`

[PATCH v2 1/3] rust: implement `kernel::sync::Refcount`

[Gary Guo]

This is a wrapping layer of `include/linux/refcount.h`. Currently the
kernel refcount has already been used in `Arc`, however it calls into
FFI directly.

Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/helpers/refcount.c      | 10 ++++
 rust/kernel/sync.rs          |  2 +
 rust/kernel/sync/refcount.rs | 97 ++++++++++++++++++++++++++++++++++++
 3 files changed, 109 insertions(+)
 create mode 100644 rust/kernel/sync/refcount.rs

diff --git a/rust/helpers/refcount.c b/rust/helpers/refcount.c
index d6adbd2e45a1..d175898ad7b8 100644
--- a/rust/helpers/refcount.c
+++ b/rust/helpers/refcount.c
@@ -7,11 +7,21 @@ refcount_t rust_helper_REFCOUNT_INIT(int n)
 	return (refcount_t)REFCOUNT_INIT(n);
 }
 
+void rust_helper_refcount_set(refcount_t *r, int n)
+{
+	refcount_set(r, n);
+}
+
 void rust_helper_refcount_inc(refcount_t *r)
 {
 	refcount_inc(r);
 }
 
+void rust_helper_refcount_dec(refcount_t *r)
+{
+	refcount_dec(r);
+}
+
 bool rust_helper_refcount_dec_and_test(refcount_t *r)
 {
 	return refcount_dec_and_test(r);
diff --git a/rust/kernel/sync.rs b/rust/kernel/sync.rs
index 1eab7ebf25fd..b76b04e16eac 100644
--- a/rust/kernel/sync.rs
+++ b/rust/kernel/sync.rs
@@ -12,6 +12,7 @@
 pub mod lock;
 mod locked_by;
 pub mod poll;
+mod refcount;
 
 pub use arc::{Arc, ArcBorrow, UniqueArc};
 pub use condvar::{new_condvar, CondVar, CondVarTimeoutResult};
@@ -19,6 +20,7 @@
 pub use lock::mutex::{new_mutex, Mutex};
 pub use lock::spinlock::{new_spinlock, SpinLock};
 pub use locked_by::LockedBy;
+pub use refcount::Refcount;
 
 /// Represents a lockdep class. It's a wrapper around C's `lock_class_key`.
 #[repr(transparent)]
diff --git a/rust/kernel/sync/refcount.rs b/rust/kernel/sync/refcount.rs
new file mode 100644
index 000000000000..2198b1598b60
--- /dev/null
+++ b/rust/kernel/sync/refcount.rs
@@ -0,0 +1,97 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Atomic reference counting.
+//!
+//! C header: [`include/linux/refcount.h`](srctree/include/linux/refcount.h)
+
+use crate::types::Opaque;
+use core::sync::atomic::AtomicI32;
+
+/// Atomic reference counter.
+///
+/// This type is conceptually an atomic integer, but provides saturation semantics compared to
+/// normal atomic integers. Values in the negative range when viewed as a signed integer are
+/// saturation (bad) values. For details about the saturation semantics, please refer to top of
+/// [`include/linux/refcount.h`](srctree/include/refcount.h).
+///
+/// Wraps the kernel's C `refcount_t`.
+#[repr(transparent)]
+pub struct Refcount(Opaque<bindings::refcount_t>);
+
+impl Refcount {
+    /// Construct a new [`Refcount`] from an initial value.
+    #[inline]
+    pub fn new(value: i32) -> Self {
+        // SAFETY: There are no safety requirements for this FFI call.
+        Self(Opaque::new(unsafe { bindings::REFCOUNT_INIT(value) }))
+    }
+
+    #[inline]
+    fn as_ptr(&self) -> *mut bindings::refcount_t {
+        self.0.get()
+    }
+
+    /// Set a refcount's value.
+    #[inline]
+    pub fn set(&self, value: i32) {
+        // SAFETY: `self.as_ptr()` is valid.
+        unsafe { bindings::refcount_set(self.as_ptr(), value) }
+    }
+
+    /// Increment a refcount.
+    ///
+    /// It will saturate if overflows and `WARN`. It will also `WARN` if the refcount is 0, as this
+    /// represents a possible use-after-free condition.
+    ///
+    /// Provides no memory ordering, it is assumed that caller already has a reference on the
+    /// object.
+    #[inline]
+    pub fn inc(&self) {
+        // SAFETY: self is valid.
+        unsafe { bindings::refcount_inc(self.as_ptr()) }
+    }
+
+    /// Decrement a refcount.
+    ///
+    /// It will `WARN` on underflow and fail to decrement when saturated.
+    ///
+    /// Provides release memory ordering, such that prior loads and stores are done
+    /// before.
+    #[inline]
+    pub fn dec(&self) {
+        // SAFETY: `self.as_ptr()` is valid.
+        unsafe { bindings::refcount_dec(self.as_ptr()) }
+    }
+
+    /// Decrement a refcount and test if it is 0.
+    ///
+    /// It will `WARN` on underflow and fail to decrement when saturated.
+    ///
+    /// Provides release memory ordering, such that prior loads and stores are done
+    /// before, and provides an acquire ordering on success such that memory deallocation
+    /// must come after.
+    ///
+    /// Returns true if the resulting refcount is 0, false otherwise.
+    #[inline]
+    #[must_use = "use `dec` instead you do not need to test if it is 0"]
+    pub fn dec_and_test(&self) -> bool {
+        // SAFETY: `self.as_ptr()` is valid.
+        unsafe { bindings::refcount_dec_and_test(self.as_ptr()) }
+    }
+
+    /// Decrement a refcount if it is not 1.
+    ///
+    /// Returns true if the decrement operation was successful, false otherwise.
+    #[inline]
+    #[must_use]
+    pub fn dec_not_one(&self) -> bool {
+        // SAFETY: `self.as_ptr()` is valid.
+        unsafe { bindings::refcount_dec_not_one(self.as_ptr()) }
+    }
+}
+
+// SAFETY: `refcount_t` is thread-safe.
+unsafe impl Send for Refcount {}
+
+// SAFETY: `refcount_t` is thread-safe.
+unsafe impl Sync for Refcount {}
-- 
2.47.0

Re: [PATCH v2 1/3] rust: implement `kernel::sync::Refcount`

[Alice Ryhl]

On Sat, Dec 21, 2024 at 7:31 PM Gary Guo <gary@garyguo.net> wrote:
>
> This is a wrapping layer of `include/linux/refcount.h`. Currently the
> kernel refcount has already been used in `Arc`, however it calls into
> FFI directly.
>
> Signed-off-by: Gary Guo <gary@garyguo.net>

Reviewed-by: Alice Ryhl <aliceryhl@google.com>

Re: [PATCH v2 1/3] rust: implement `kernel::sync::Refcount`

[Boqun Feng]

On Sat, Dec 21, 2024 at 06:29:44PM +0000, Gary Guo wrote:
> This is a wrapping layer of `include/linux/refcount.h`. Currently the
> kernel refcount has already been used in `Arc`, however it calls into
> FFI directly.
> 
> Signed-off-by: Gary Guo <gary@garyguo.net>
> ---
>  rust/helpers/refcount.c      | 10 ++++
>  rust/kernel/sync.rs          |  2 +
>  rust/kernel/sync/refcount.rs | 97 ++++++++++++++++++++++++++++++++++++
>  3 files changed, 109 insertions(+)
>  create mode 100644 rust/kernel/sync/refcount.rs
> 
> diff --git a/rust/helpers/refcount.c b/rust/helpers/refcount.c
> index d6adbd2e45a1..d175898ad7b8 100644
> --- a/rust/helpers/refcount.c
> +++ b/rust/helpers/refcount.c
> @@ -7,11 +7,21 @@ refcount_t rust_helper_REFCOUNT_INIT(int n)
>  	return (refcount_t)REFCOUNT_INIT(n);
>  }
>  
> +void rust_helper_refcount_set(refcount_t *r, int n)
> +{
> +	refcount_set(r, n);
> +}
> +
>  void rust_helper_refcount_inc(refcount_t *r)
>  {
>  	refcount_inc(r);
>  }
>  
> +void rust_helper_refcount_dec(refcount_t *r)
> +{
> +	refcount_dec(r);
> +}
> +
>  bool rust_helper_refcount_dec_and_test(refcount_t *r)
>  {
>  	return refcount_dec_and_test(r);
> diff --git a/rust/kernel/sync.rs b/rust/kernel/sync.rs
> index 1eab7ebf25fd..b76b04e16eac 100644
> --- a/rust/kernel/sync.rs
> +++ b/rust/kernel/sync.rs
> @@ -12,6 +12,7 @@
>  pub mod lock;
>  mod locked_by;
>  pub mod poll;
> +mod refcount;
>  
>  pub use arc::{Arc, ArcBorrow, UniqueArc};
>  pub use condvar::{new_condvar, CondVar, CondVarTimeoutResult};
> @@ -19,6 +20,7 @@
>  pub use lock::mutex::{new_mutex, Mutex};
>  pub use lock::spinlock::{new_spinlock, SpinLock};
>  pub use locked_by::LockedBy;
> +pub use refcount::Refcount;
>  
>  /// Represents a lockdep class. It's a wrapper around C's `lock_class_key`.
>  #[repr(transparent)]
> diff --git a/rust/kernel/sync/refcount.rs b/rust/kernel/sync/refcount.rs
> new file mode 100644
> index 000000000000..2198b1598b60
> --- /dev/null
> +++ b/rust/kernel/sync/refcount.rs

Could you add this file into the entry of "ATOMIC INFRASTRUCTURE"? I
would also suggest you to add yourself as a reviewer or maintainer in
that entry (given your expertise on atomic and related compiler
behaviors), but that's totally up to you.

> @@ -0,0 +1,97 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +//! Atomic reference counting.
> +//!
> +//! C header: [`include/linux/refcount.h`](srctree/include/linux/refcount.h)
> +
> +use crate::types::Opaque;
> +use core::sync::atomic::AtomicI32;
> +

Could you move this "use" into patch #3, where it gets used?

Reviewed-by: Boqun Feng <boqun.feng@gmail.com>

Regards,
Boqun

[...]

[PATCH v2 2/3] rust: convert `Arc` to use `Refcount`

[Gary Guo]

With `Refcount` type created, `Arc` can use `Refcount` instead of
calling into FFI directly.

Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/kernel/sync/arc.rs | 64 ++++++++++++++++-------------------------
 1 file changed, 25 insertions(+), 39 deletions(-)

diff --git a/rust/kernel/sync/arc.rs b/rust/kernel/sync/arc.rs
index 9f0b04400e8e..34598ec1f977 100644
--- a/rust/kernel/sync/arc.rs
+++ b/rust/kernel/sync/arc.rs
@@ -8,7 +8,7 @@
 //! threads.
 //!
 //! It is different from the standard library's [`Arc`] in a few ways:
-//! 1. It is backed by the kernel's `refcount_t` type.
+//! 1. It is backed by the kernel's [`Refcount`] type.
 //! 2. It does not support weak references, which allows it to be half the size.
 //! 3. It saturates the reference count instead of aborting when it goes over a threshold.
 //! 4. It does not provide a `get_mut` method, so the ref counted object is pinned.
@@ -18,10 +18,10 @@
 
 use crate::{
     alloc::{AllocError, Flags, KBox},
-    bindings,
     init::{self, InPlaceInit, Init, PinInit},
+    sync::Refcount,
     try_init,
-    types::{ForeignOwnable, Opaque},
+    types::ForeignOwnable,
 };
 use core::{
     alloc::Layout,
@@ -141,7 +141,7 @@ pub struct Arc<T: ?Sized> {
 #[pin_data]
 #[repr(C)]
 struct ArcInner<T: ?Sized> {
-    refcount: Opaque<bindings::refcount_t>,
+    refcount: Refcount,
     data: T,
 }
 
@@ -153,7 +153,7 @@ impl<T: ?Sized> ArcInner<T> {
     /// `ptr` must have been returned by a previous call to [`Arc::into_raw`], and the `Arc` must
     /// not yet have been destroyed.
     unsafe fn container_of(ptr: *const T) -> NonNull<ArcInner<T>> {
-        let refcount_layout = Layout::new::<bindings::refcount_t>();
+        let refcount_layout = Layout::new::<Refcount>();
         // SAFETY: The caller guarantees that the pointer is valid.
         let val_layout = Layout::for_value(unsafe { &*ptr });
         // SAFETY: We're computing the layout of a real struct that existed when compiling this
@@ -203,8 +203,7 @@ impl<T> Arc<T> {
     pub fn new(contents: T, flags: Flags) -> Result<Self, AllocError> {
         // INVARIANT: The refcount is initialised to a non-zero value.
         let value = ArcInner {
-            // SAFETY: There are no safety requirements for this FFI call.
-            refcount: Opaque::new(unsafe { bindings::REFCOUNT_INIT(1) }),
+            refcount: Refcount::new(1),
             data: contents,
         };
 
@@ -285,7 +284,7 @@ pub fn ptr_eq(this: &Self, other: &Self) -> bool {
     /// use kernel::sync::{Arc, UniqueArc};
     ///
     /// let arc = Arc::new(42, GFP_KERNEL)?;
-    /// let unique_arc = arc.into_unique_or_drop();
+    /// let unique_arc = Arc::into_unique_or_drop(arc);
     ///
     /// // The above conversion should succeed since refcount of `arc` is 1.
     /// assert!(unique_arc.is_some());
@@ -301,35 +300,25 @@ pub fn ptr_eq(this: &Self, other: &Self) -> bool {
     /// let arc = Arc::new(42, GFP_KERNEL)?;
     /// let another = arc.clone();
     ///
-    /// let unique_arc = arc.into_unique_or_drop();
+    /// let unique_arc = Arc::into_unique_or_drop(arc);
     ///
     /// // The above conversion should fail since refcount of `arc` is >1.
     /// assert!(unique_arc.is_none());
     ///
     /// # Ok::<(), Error>(())
     /// ```
-    pub fn into_unique_or_drop(self) -> Option<Pin<UniqueArc<T>>> {
+    pub fn into_unique_or_drop(this: Self) -> Option<Pin<UniqueArc<T>>> {
         // We will manually manage the refcount in this method, so we disable the destructor.
-        let me = ManuallyDrop::new(self);
+        let this = ManuallyDrop::new(this);
         // SAFETY: We own a refcount, so the pointer is still valid.
-        let refcount = unsafe { me.ptr.as_ref() }.refcount.get();
+        let refcount = unsafe { &this.ptr.as_ref().refcount };
 
-        // If the refcount reaches a non-zero value, then we have destroyed this `Arc` and will
-        // return without further touching the `Arc`. If the refcount reaches zero, then there are
-        // no other arcs, and we can create a `UniqueArc`.
-        //
-        // SAFETY: We own a refcount, so the pointer is not dangling.
-        let is_zero = unsafe { bindings::refcount_dec_and_test(refcount) };
-        if is_zero {
-            // SAFETY: We have exclusive access to the arc, so we can perform unsynchronized
-            // accesses to the refcount.
-            unsafe { core::ptr::write(refcount, bindings::REFCOUNT_INIT(1)) };
-
-            // INVARIANT: We own the only refcount to this arc, so we may create a `UniqueArc`. We
-            // must pin the `UniqueArc` because the values was previously in an `Arc`, and they pin
-            // their values.
+        if !refcount.dec_not_one() {
+            // INVARIANT: If the refcount failed to decrement because it is 1, then we have the
+            // exclusive ownership, so we may create a `UniqueArc`. We must pin the `UniqueArc`
+            // because the values was previously in an `Arc`, and they pin their values.
             Some(Pin::from(UniqueArc {
-                inner: ManuallyDrop::into_inner(me),
+                inner: ManuallyDrop::into_inner(this),
             }))
         } else {
             None
@@ -380,10 +369,10 @@ fn as_ref(&self) -> &T {
 
 impl<T: ?Sized> Clone for Arc<T> {
     fn clone(&self) -> Self {
-        // INVARIANT: C `refcount_inc` saturates the refcount, so it cannot overflow to zero.
+        // INVARIANT: `Refcount` saturates the refcount, so it cannot overflow to zero.
         // SAFETY: By the type invariant, there is necessarily a reference to the object, so it is
         // safe to increment the refcount.
-        unsafe { bindings::refcount_inc(self.ptr.as_ref().refcount.get()) };
+        unsafe { self.ptr.as_ref().refcount.inc() };
 
         // SAFETY: We just incremented the refcount. This increment is now owned by the new `Arc`.
         unsafe { Self::from_inner(self.ptr) }
@@ -392,16 +381,14 @@ fn clone(&self) -> Self {
 
 impl<T: ?Sized> Drop for Arc<T> {
     fn drop(&mut self) {
-        // SAFETY: By the type invariant, there is necessarily a reference to the object. We cannot
-        // touch `refcount` after it's decremented to a non-zero value because another thread/CPU
-        // may concurrently decrement it to zero and free it. It is ok to have a raw pointer to
-        // freed/invalid memory as long as it is never dereferenced.
-        let refcount = unsafe { self.ptr.as_ref() }.refcount.get();
-
         // INVARIANT: If the refcount reaches zero, there are no other instances of `Arc`, and
         // this instance is being dropped, so the broken invariant is not observable.
-        // SAFETY: Also by the type invariant, we are allowed to decrement the refcount.
-        let is_zero = unsafe { bindings::refcount_dec_and_test(refcount) };
+        // SAFETY: By the type invariant, there is necessarily a reference to the object.
+        // NOTE: we cannot touch `refcount` after it's decremented to a non-zero value because
+        // another thread/CPU may concurrently decrement it to zero and free it. However it is okay
+        // to have a transient reference to decrement the refcount, see
+        // https://github.com/rust-lang/rust/issues/55005.
+        let is_zero = unsafe { self.ptr.as_ref().refcount.dec_and_test() };
         if is_zero {
             // The count reached zero, we must free the memory.
             //
@@ -650,8 +637,7 @@ pub fn new_uninit(flags: Flags) -> Result<UniqueArc<MaybeUninit<T>>, AllocError>
         // INVARIANT: The refcount is initialised to a non-zero value.
         let inner = KBox::try_init::<AllocError>(
             try_init!(ArcInner {
-                // SAFETY: There are no safety requirements for this FFI call.
-                refcount: Opaque::new(unsafe { bindings::REFCOUNT_INIT(1) }),
+                refcount: Refcount::new(1),
                 data <- init::uninit::<T, AllocError>(),
             }? AllocError),
             flags,
-- 
2.47.0

Re: [PATCH v2 2/3] rust: convert `Arc` to use `Refcount`

[Alice Ryhl]

On Sat, Dec 21, 2024 at 7:31 PM Gary Guo <gary@garyguo.net> wrote:
>
> With `Refcount` type created, `Arc` can use `Refcount` instead of
> calling into FFI directly.
>
> Signed-off-by: Gary Guo <gary@garyguo.net>

[...]

> -    pub fn into_unique_or_drop(self) -> Option<Pin<UniqueArc<T>>> {
> +    pub fn into_unique_or_drop(this: Self) -> Option<Pin<UniqueArc<T>>> {
>          // We will manually manage the refcount in this method, so we disable the destructor.
> -        let me = ManuallyDrop::new(self);
> +        let this = ManuallyDrop::new(this);
>          // SAFETY: We own a refcount, so the pointer is still valid.
> -        let refcount = unsafe { me.ptr.as_ref() }.refcount.get();
> +        let refcount = unsafe { &this.ptr.as_ref().refcount };
>
> -        // If the refcount reaches a non-zero value, then we have destroyed this `Arc` and will
> -        // return without further touching the `Arc`. If the refcount reaches zero, then there are
> -        // no other arcs, and we can create a `UniqueArc`.
> -        //
> -        // SAFETY: We own a refcount, so the pointer is not dangling.
> -        let is_zero = unsafe { bindings::refcount_dec_and_test(refcount) };
> +        if !refcount.dec_not_one() {

This is wrong. The into_unique_or_drop function must establish an
acqrel ordering when a UniqueArc is created, but dec_not_one() does
not do so. You need to use refcount_dec_and_test() instead.

Alice

Re: [PATCH v2 2/3] rust: convert `Arc` to use `Refcount`

[Gary Guo]

On Tue, 14 Jan 2025 11:02:25 +0100
Alice Ryhl <aliceryhl@google.com> wrote:

> On Sat, Dec 21, 2024 at 7:31 PM Gary Guo <gary@garyguo.net> wrote:
> >
> > With `Refcount` type created, `Arc` can use `Refcount` instead of
> > calling into FFI directly.
> >
> > Signed-off-by: Gary Guo <gary@garyguo.net>  
> 
> [...]
> 
> > -    pub fn into_unique_or_drop(self) -> Option<Pin<UniqueArc<T>>> {
> > +    pub fn into_unique_or_drop(this: Self) -> Option<Pin<UniqueArc<T>>> {
> >          // We will manually manage the refcount in this method, so we disable the destructor.
> > -        let me = ManuallyDrop::new(self);
> > +        let this = ManuallyDrop::new(this);
> >          // SAFETY: We own a refcount, so the pointer is still valid.
> > -        let refcount = unsafe { me.ptr.as_ref() }.refcount.get();
> > +        let refcount = unsafe { &this.ptr.as_ref().refcount };
> >
> > -        // If the refcount reaches a non-zero value, then we have destroyed this `Arc` and will
> > -        // return without further touching the `Arc`. If the refcount reaches zero, then there are
> > -        // no other arcs, and we can create a `UniqueArc`.
> > -        //
> > -        // SAFETY: We own a refcount, so the pointer is not dangling.
> > -        let is_zero = unsafe { bindings::refcount_dec_and_test(refcount) };
> > +        if !refcount.dec_not_one() {  
> 
> This is wrong. The into_unique_or_drop function must establish an
> acqrel ordering when a UniqueArc is created, but dec_not_one() does
> not do so. You need to use refcount_dec_and_test() instead.
> 
> Alice

Ah, good catch. In this case I think an acquire fence in the unique
path would be sufficient? Or would you prefer to use `dec_and_test` and
`set`?

Best,
Gary

Re: [PATCH v2 2/3] rust: convert `Arc` to use `Refcount`

[Boqun Feng]

On Wed, Jan 15, 2025 at 12:32:34PM +0000, Gary Guo wrote:
> On Tue, 14 Jan 2025 11:02:25 +0100
> Alice Ryhl <aliceryhl@google.com> wrote:
> 
> > On Sat, Dec 21, 2024 at 7:31 PM Gary Guo <gary@garyguo.net> wrote:
> > >
> > > With `Refcount` type created, `Arc` can use `Refcount` instead of
> > > calling into FFI directly.
> > >
> > > Signed-off-by: Gary Guo <gary@garyguo.net>  
> > 
> > [...]
> > 
> > > -    pub fn into_unique_or_drop(self) -> Option<Pin<UniqueArc<T>>> {
> > > +    pub fn into_unique_or_drop(this: Self) -> Option<Pin<UniqueArc<T>>> {
> > >          // We will manually manage the refcount in this method, so we disable the destructor.
> > > -        let me = ManuallyDrop::new(self);
> > > +        let this = ManuallyDrop::new(this);
> > >          // SAFETY: We own a refcount, so the pointer is still valid.
> > > -        let refcount = unsafe { me.ptr.as_ref() }.refcount.get();
> > > +        let refcount = unsafe { &this.ptr.as_ref().refcount };
> > >
> > > -        // If the refcount reaches a non-zero value, then we have destroyed this `Arc` and will
> > > -        // return without further touching the `Arc`. If the refcount reaches zero, then there are
> > > -        // no other arcs, and we can create a `UniqueArc`.
> > > -        //
> > > -        // SAFETY: We own a refcount, so the pointer is not dangling.
> > > -        let is_zero = unsafe { bindings::refcount_dec_and_test(refcount) };
> > > +        if !refcount.dec_not_one() {  
> > 
> > This is wrong. The into_unique_or_drop function must establish an
> > acqrel ordering when a UniqueArc is created, but dec_not_one() does
> > not do so. You need to use refcount_dec_and_test() instead.
> > 
> > Alice
> 
> Ah, good catch. In this case I think an acquire fence in the unique

Note that we don't have an acquire fence in LKMM, so you may need to use
a smp_mb() here.

> path would be sufficient? Or would you prefer to use `dec_and_test` and
> `set`?
> 

using `dec_and_test()` + `set` should be cheaper than dec_not_one() +
smp_mb().

Regards,
Boqun

> Best,
> Gary

Re: [PATCH v2 2/3] rust: convert `Arc` to use `Refcount`

[Alice Ryhl]

On Wed, Jan 15, 2025 at 1:32 PM Gary Guo <gary@garyguo.net> wrote:
>
> On Tue, 14 Jan 2025 11:02:25 +0100
> Alice Ryhl <aliceryhl@google.com> wrote:
>
> > On Sat, Dec 21, 2024 at 7:31 PM Gary Guo <gary@garyguo.net> wrote:
> > >
> > > With `Refcount` type created, `Arc` can use `Refcount` instead of
> > > calling into FFI directly.
> > >
> > > Signed-off-by: Gary Guo <gary@garyguo.net>
> >
> > [...]
> >
> > > -    pub fn into_unique_or_drop(self) -> Option<Pin<UniqueArc<T>>> {
> > > +    pub fn into_unique_or_drop(this: Self) -> Option<Pin<UniqueArc<T>>> {
> > >          // We will manually manage the refcount in this method, so we disable the destructor.
> > > -        let me = ManuallyDrop::new(self);
> > > +        let this = ManuallyDrop::new(this);
> > >          // SAFETY: We own a refcount, so the pointer is still valid.
> > > -        let refcount = unsafe { me.ptr.as_ref() }.refcount.get();
> > > +        let refcount = unsafe { &this.ptr.as_ref().refcount };
> > >
> > > -        // If the refcount reaches a non-zero value, then we have destroyed this `Arc` and will
> > > -        // return without further touching the `Arc`. If the refcount reaches zero, then there are
> > > -        // no other arcs, and we can create a `UniqueArc`.
> > > -        //
> > > -        // SAFETY: We own a refcount, so the pointer is not dangling.
> > > -        let is_zero = unsafe { bindings::refcount_dec_and_test(refcount) };
> > > +        if !refcount.dec_not_one() {
> >
> > This is wrong. The into_unique_or_drop function must establish an
> > acqrel ordering when a UniqueArc is created, but dec_not_one() does
> > not do so. You need to use refcount_dec_and_test() instead.
> >
> > Alice
>
> Ah, good catch. In this case I think an acquire fence in the unique
> path would be sufficient? Or would you prefer to use `dec_and_test` and
> `set`?

The existing approach is what I prefer.

Alice

[PATCH v2 3/3] rust: block: convert `block::mq` to use `Refcount`

[Gary Guo]

Currently there's a custom reference counting in `block::mq`, which uses
`AtomicU64` Rust atomics, and this type doesn't exist on some 32-bit
architectures. We cannot just change it to use 32-bit atomics, because
doing so will make it vulnerable to refcount overflow. So switch it to
use the kernel refcount `kernel::sync::Refcount` instead.

There is an operation needed by `block::mq`, atomically decreasing
refcount from 2 to 0, which is not available through refcount.h, so
I exposed `Refcount::as_atomic` which allows accessing the refcount
directly.

Acked-by: Andreas Hindborg <a.hindborg@kernel.org>
Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/kernel/block/mq/operations.rs |  7 +--
 rust/kernel/block/mq/request.rs    | 70 ++++++++++--------------------
 rust/kernel/sync/refcount.rs       | 12 +++++
 3 files changed, 38 insertions(+), 51 deletions(-)

diff --git a/rust/kernel/block/mq/operations.rs b/rust/kernel/block/mq/operations.rs
index c8646d0d9866..8314ce94910a 100644
--- a/rust/kernel/block/mq/operations.rs
+++ b/rust/kernel/block/mq/operations.rs
@@ -9,9 +9,10 @@
     block::mq::request::RequestDataWrapper,
     block::mq::Request,
     error::{from_result, Result},
+    sync::Refcount,
     types::ARef,
 };
-use core::{marker::PhantomData, sync::atomic::AtomicU64, sync::atomic::Ordering};
+use core::marker::PhantomData;
 
 /// Implement this trait to interface blk-mq as block devices.
 ///
@@ -77,7 +78,7 @@ impl<T: Operations> OperationsVTable<T> {
         let request = unsafe { &*(*bd).rq.cast::<Request<T>>() };
 
         // One refcount for the ARef, one for being in flight
-        request.wrapper_ref().refcount().store(2, Ordering::Relaxed);
+        request.wrapper_ref().refcount().set(2);
 
         // SAFETY:
         //  - We own a refcount that we took above. We pass that to `ARef`.
@@ -186,7 +187,7 @@ impl<T: Operations> OperationsVTable<T> {
 
             // SAFETY: The refcount field is allocated but not initialized, so
             // it is valid for writes.
-            unsafe { RequestDataWrapper::refcount_ptr(pdu.as_ptr()).write(AtomicU64::new(0)) };
+            unsafe { RequestDataWrapper::refcount_ptr(pdu.as_ptr()).write(Refcount::new(0)) };
 
             Ok(0)
         })
diff --git a/rust/kernel/block/mq/request.rs b/rust/kernel/block/mq/request.rs
index 7943f43b9575..7c782d70935e 100644
--- a/rust/kernel/block/mq/request.rs
+++ b/rust/kernel/block/mq/request.rs
@@ -8,12 +8,13 @@
     bindings,
     block::mq::Operations,
     error::Result,
+    sync::Refcount,
     types::{ARef, AlwaysRefCounted, Opaque},
 };
 use core::{
     marker::PhantomData,
     ptr::{addr_of_mut, NonNull},
-    sync::atomic::{AtomicU64, Ordering},
+    sync::atomic::Ordering,
 };
 
 /// A wrapper around a blk-mq [`struct request`]. This represents an IO request.
@@ -37,6 +38,9 @@
 /// We need to track 3 and 4 to ensure that it is safe to end the request and hand
 /// back ownership to the block layer.
 ///
+/// Note that driver can still obtain new `ARef` even if there is no `ARef`s in existence by using
+/// `tag_to_rq`, hence the need to distinct B and C.
+///
 /// The states are tracked through the private `refcount` field of
 /// `RequestDataWrapper`. This structure lives in the private data area of the C
 /// [`struct request`].
@@ -98,13 +102,17 @@ pub(crate) unsafe fn start_unchecked(this: &ARef<Self>) {
     ///
     /// [`struct request`]: srctree/include/linux/blk-mq.h
     fn try_set_end(this: ARef<Self>) -> Result<*mut bindings::request, ARef<Self>> {
-        // We can race with `TagSet::tag_to_rq`
-        if let Err(_old) = this.wrapper_ref().refcount().compare_exchange(
-            2,
-            0,
-            Ordering::Relaxed,
-            Ordering::Relaxed,
-        ) {
+        // To hand back the ownership, we need the current refcount to be 2.
+        // Since we can race with `TagSet::tag_to_rq`, this needs to atomically reduce
+        // refcount to 0. `Refcount` does not provide a way to do this, so use the underlying
+        // atomics directly.
+        if this
+            .wrapper_ref()
+            .refcount()
+            .as_atomic()
+            .compare_exchange(2, 0, Ordering::Relaxed, Ordering::Relaxed)
+            .is_err()
+        {
             return Err(this);
         }
 
@@ -168,13 +176,13 @@ pub(crate) struct RequestDataWrapper {
     /// - 0: The request is owned by C block layer.
     /// - 1: The request is owned by Rust abstractions but there are no [`ARef`] references to it.
     /// - 2+: There are [`ARef`] references to the request.
-    refcount: AtomicU64,
+    refcount: Refcount,
 }
 
 impl RequestDataWrapper {
     /// Return a reference to the refcount of the request that is embedding
     /// `self`.
-    pub(crate) fn refcount(&self) -> &AtomicU64 {
+    pub(crate) fn refcount(&self) -> &Refcount {
         &self.refcount
     }
 
@@ -184,7 +192,7 @@ pub(crate) fn refcount(&self) -> &AtomicU64 {
     /// # Safety
     ///
     /// - `this` must point to a live allocation of at least the size of `Self`.
-    pub(crate) unsafe fn refcount_ptr(this: *mut Self) -> *mut AtomicU64 {
+    pub(crate) unsafe fn refcount_ptr(this: *mut Self) -> *mut Refcount {
         // SAFETY: Because of the safety requirements of this function, the
         // field projection is safe.
         unsafe { addr_of_mut!((*this).refcount) }
@@ -200,47 +208,13 @@ unsafe impl<T: Operations> Send for Request<T> {}
 // mutate `self` are internally synchronized`
 unsafe impl<T: Operations> Sync for Request<T> {}
 
-/// Store the result of `op(target.load())` in target, returning new value of
-/// target.
-fn atomic_relaxed_op_return(target: &AtomicU64, op: impl Fn(u64) -> u64) -> u64 {
-    let old = target.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |x| Some(op(x)));
-
-    // SAFETY: Because the operation passed to `fetch_update` above always
-    // return `Some`, `old` will always be `Ok`.
-    let old = unsafe { old.unwrap_unchecked() };
-
-    op(old)
-}
-
-/// Store the result of `op(target.load)` in `target` if `target.load() !=
-/// pred`, returning [`true`] if the target was updated.
-fn atomic_relaxed_op_unless(target: &AtomicU64, op: impl Fn(u64) -> u64, pred: u64) -> bool {
-    target
-        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |x| {
-            if x == pred {
-                None
-            } else {
-                Some(op(x))
-            }
-        })
-        .is_ok()
-}
-
 // SAFETY: All instances of `Request<T>` are reference counted. This
 // implementation of `AlwaysRefCounted` ensure that increments to the ref count
 // keeps the object alive in memory at least until a matching reference count
 // decrement is executed.
 unsafe impl<T: Operations> AlwaysRefCounted for Request<T> {
     fn inc_ref(&self) {
-        let refcount = &self.wrapper_ref().refcount();
-
-        #[cfg_attr(not(CONFIG_DEBUG_MISC), allow(unused_variables))]
-        let updated = atomic_relaxed_op_unless(refcount, |x| x + 1, 0);
-
-        #[cfg(CONFIG_DEBUG_MISC)]
-        if !updated {
-            panic!("Request refcount zero on clone")
-        }
+        self.wrapper_ref().refcount().inc();
     }
 
     unsafe fn dec_ref(obj: core::ptr::NonNull<Self>) {
@@ -252,10 +226,10 @@ unsafe fn dec_ref(obj: core::ptr::NonNull<Self>) {
         let refcount = unsafe { &*RequestDataWrapper::refcount_ptr(wrapper_ptr) };
 
         #[cfg_attr(not(CONFIG_DEBUG_MISC), allow(unused_variables))]
-        let new_refcount = atomic_relaxed_op_return(refcount, |x| x - 1);
+        let is_zero = refcount.dec_and_test();
 
         #[cfg(CONFIG_DEBUG_MISC)]
-        if new_refcount == 0 {
+        if is_zero {
             panic!("Request reached refcount zero in Rust abstractions");
         }
     }
diff --git a/rust/kernel/sync/refcount.rs b/rust/kernel/sync/refcount.rs
index 2198b1598b60..ed525c719c12 100644
--- a/rust/kernel/sync/refcount.rs
+++ b/rust/kernel/sync/refcount.rs
@@ -31,6 +31,18 @@ fn as_ptr(&self) -> *mut bindings::refcount_t {
         self.0.get()
     }
 
+    /// Get the underlying atomic counter that backs the refcount.
+    ///
+    /// NOTE: This will be changed to LKMM atomic in the future.
+    #[inline]
+    pub fn as_atomic(&self) -> &AtomicI32 {
+        let ptr = self.0.get() as *const AtomicI32;
+        // SAFETY: `refcount_t` is a transparent wrapper of `atomic_t`, which is an atomic 32-bit
+        // integer that is layout-wise compatible with `AtomicI32`. All values are valid for
+        // `refcount_t`, despite some of the values are considered saturated and "bad".
+        unsafe { &*ptr }
+    }
+
     /// Set a refcount's value.
     #[inline]
     pub fn set(&self, value: i32) {
-- 
2.47.0