[PATCH v4 0/5] implement `kernel::sync::Refcount` and convert users

[PATCH v4 0/5] implement `kernel::sync::Refcount` and convert users

[Gary Guo]

From: Gary Guo <gary@garyguo.net>

Currently there're two refcount usage in rust/, `Arc` and `block::mq`.
`Arc` uses `refcount_t` with FFI calls directly, and `block::mq` use
Rust atomics and custom refcounting.

This series consolidate them to have a single `Refcount` which wraps
`refcount_t` and have it used by both.

With the removal of Rust `AtomicU64` this would also make the code work
for 32-bit systems. See [1] as a previous attempt to fix this issue.

Cc: Will Deacon <will@kernel.org>
Cc: Peter Zijlstra <peterz@infradead.org>
Cc: Boqun Feng <boqun.feng@gmail.com>
Cc: Mark Rutland <mark.rutland@arm.com>
Link: https://lore.kernel.org/rust-for-linux/20240905061214.3954271-1-davidgow@google.com/ [1]

Changes in v4:
- Move justification on why `&Refcount` is okay when used for reference
  coutning into the `Refcount::dec_and_test` documentation comment.
- Split `Arc::into_unique_or_drop` signature change into another patch.
- Wording/comment changes
- Link to v3: https://lore.kernel.org/rust-for-linux/20250219201602.1898383-1-gary@garyguo.net/

Changes in v3:
- Drop `dec_not_one` and revert to `dec_and_test`/`set` as the former
  lacks acquire semantics.
- Move a `use` statement from patch 1 to correct place (patch 3).
- Update maintainer entry to include refcount.rs and add myself as a
  reviewer, as suggested by Boqun.
- Link to v2: https://lore.kernel.org/rust-for-linux/20241221183024.3929500-1-gary@garyguo.net

Changes in v2:
- `Refcount::read` method is dropped
- `Refcount::dec_not_one` method is added and `Arc::into_unique_or_drop`
  is converted to use it.
- Link to v1: https://lore.kernel.org/rust-for-linux/20241004155247.2210469-1-gary@garyguo.net

Gary Guo (4):
  rust: implement `kernel::sync::Refcount`
  rust: convert `Arc` to use `Refcount`
  rust: block: convert `block::mq` to use `Refcount`
  MAINTAINERS: update atomic infrastructure entry to include Rust
Gary Guo (5):
  rust: implement `kernel::sync::Refcount`
  rust: make `Arc::into_unique_or_drop` associated function
  rust: convert `Arc` to use `Refcount`
  rust: block: convert `block::mq` to use `Refcount`
  MAINTAINERS: update atomic infrastructure entry to include Rust

 MAINTAINERS                        |   2 +
 rust/helpers/refcount.c            |  10 +++
 rust/kernel/block/mq/operations.rs |   7 +-
 rust/kernel/block/mq/request.rs    |  63 +++++------------
 rust/kernel/sync.rs                |   2 +
 rust/kernel/sync/arc.rs            |  55 +++++----------
 rust/kernel/sync/refcount.rs       | 108 +++++++++++++++++++++++++++++
 7 files changed, 161 insertions(+), 86 deletions(-)
 create mode 100644 rust/kernel/sync/refcount.rs


base-commit: e04c78d86a9699d136910cfc0bdcf01087e3267e
-- 
2.49.0

[PATCH v4 1/5] rust: implement `kernel::sync::Refcount`

[Gary Guo]

From: Gary Guo <gary@garyguo.net>

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.

Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Reviewed-by: Boqun Feng <boqun.feng@gmail.com>
Reviewed-by: Fiona Behrens <me@kloenk.dev>
Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/helpers/refcount.c      | 10 +++++
 rust/kernel/sync.rs          |  2 +
 rust/kernel/sync/refcount.rs | 86 ++++++++++++++++++++++++++++++++++++
 3 files changed, 98 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 36a719015583..09e81a1c2966 100644
--- a/rust/kernel/sync.rs
+++ b/rust/kernel/sync.rs
@@ -15,6 +15,7 @@
 mod locked_by;
 pub mod poll;
 pub mod rcu;
+mod refcount;
 
 pub use arc::{Arc, ArcBorrow, UniqueArc};
 pub use condvar::{new_condvar, CondVar, CondVarTimeoutResult};
@@ -22,6 +23,7 @@
 pub use lock::mutex::{new_mutex, Mutex, MutexGuard};
 pub use lock::spinlock::{new_spinlock, SpinLock, SpinLockGuard};
 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..a0fc22f6d645
--- /dev/null
+++ b/rust/kernel/sync/refcount.rs
@@ -0,0 +1,86 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Atomic reference counting.
+//!
+//! C header: [`include/linux/refcount.h`](srctree/include/linux/refcount.h)
+
+use crate::types::Opaque;
+
+/// 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/linux/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 if 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()) }
+    }
+}
+
+// SAFETY: `refcount_t` is thread-safe.
+unsafe impl Send for Refcount {}
+
+// SAFETY: `refcount_t` is thread-safe.
+unsafe impl Sync for Refcount {}
-- 
2.49.0

Re: [PATCH v4 1/5] rust: implement `kernel::sync::Refcount`

[Benno Lossin]

On Sun Jun 22, 2025 at 2:57 PM CEST, Gary Guo wrote:
> +impl Refcount {
> +    /// Construct a new [`Refcount`] from an initial value.
> +    #[inline]
> +    pub fn new(value: i32) -> Self {

Should we really allow users to set a negative value from the get-go?
Here a `u31` might come in real handy...

> +        // 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) {

Same here. We should of course provide a `saturate` function, but I
don't see a reason to set it to another negative value.

---
Cheers,
Benno

> +        // SAFETY: `self.as_ptr()` is valid.
> +        unsafe { bindings::refcount_set(self.as_ptr(), value) }
> +    }

Re: [PATCH v4 1/5] rust: implement `kernel::sync::Refcount`

[Gary Guo]

On Sun, 22 Jun 2025 23:05:09 +0200
"Benno Lossin" <lossin@kernel.org> wrote:

> On Sun Jun 22, 2025 at 2:57 PM CEST, Gary Guo wrote:
> > +impl Refcount {
> > +    /// Construct a new [`Refcount`] from an initial value.
> > +    #[inline]
> > +    pub fn new(value: i32) -> Self {  
> 
> Should we really allow users to set a negative value from the get-go?
> Here a `u31` might come in real handy...
> 
> > +        // 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) {  
> 
> Same here. We should of course provide a `saturate` function, but I
> don't see a reason to set it to another negative value.

Well, it's unlikely the caller would want to init/set the value to the
saturated range, but given that refcount is effectively a thin wrapper
of atomics, user can always obtain the atomics and modify the value.

Also, I did a quick grep of the C refcount API users:
* crypto/algapi.c uses -1 as a special value, and has invocation of
  refcount_set(..., -1)
* lib/stackdepot.c has a invocation to set the refcount into saturated
  range directly.
* There're a few cases where runtime values are set, so `build_assert`
  will not work for these use cases.

Ultimately I think it should be fine to expose `i32` to the user of
this API. Use of `Refcount` to manage resources typically require
reasoning on the user side, so unintended error would be caught that
way.

> 
> ---
> Cheers,
> Benno
> 
> > +        // SAFETY: `self.as_ptr()` is valid.
> > +        unsafe { bindings::refcount_set(self.as_ptr(), value) }
> > +    }  

Best,
Gary

[PATCH v4 2/5] rust: make `Arc::into_unique_or_drop` associated function

[Gary Guo]

From: Gary Guo <gary@garyguo.net>

Make `Arc::into_unique_or_drop` to become a mere associated function
instead of a method (i.e. removing the `self` receiver).

It's a general convention for Rust smart pointers to avoid having
methods defined on them, because if the pointee type has a method of the
same name, then it is shadowed. This is normally for avoiding semver
breakage, which isn't an issue for kernel codebase, but it's still
generally a good practice to follow this rule, so that `ptr.foo()` would
always be calling a method on the pointee type.

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

diff --git a/rust/kernel/sync/arc.rs b/rust/kernel/sync/arc.rs
index c7af0aa48a0a..09a471fea4da 100644
--- a/rust/kernel/sync/arc.rs
+++ b/rust/kernel/sync/arc.rs
@@ -320,7 +320,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());
@@ -336,18 +336,18 @@ 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.get();
 
         // 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
@@ -364,7 +364,7 @@ pub fn into_unique_or_drop(self) -> Option<Pin<UniqueArc<T>>> {
             // 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
-- 
2.49.0

Re: [PATCH v4 2/5] rust: make `Arc::into_unique_or_drop` associated function

[Benno Lossin]

On Sun Jun 22, 2025 at 2:57 PM CEST, Gary Guo wrote:
> From: Gary Guo <gary@garyguo.net>
>
> Make `Arc::into_unique_or_drop` to become a mere associated function
> instead of a method (i.e. removing the `self` receiver).
>
> It's a general convention for Rust smart pointers to avoid having
> methods defined on them, because if the pointee type has a method of the
> same name, then it is shadowed. This is normally for avoiding semver
> breakage, which isn't an issue for kernel codebase, but it's still
> generally a good practice to follow this rule, so that `ptr.foo()` would
> always be calling a method on the pointee type.
>
> Signed-off-by: Gary Guo <gary@garyguo.net>

Reviewed-by: Benno Lossin <lossin@kernel.org>

---
Cheers,
Benno

> ---
>  rust/kernel/sync/arc.rs | 12 ++++++------
>  1 file changed, 6 insertions(+), 6 deletions(-)

[PATCH v4 3/5] rust: convert `Arc` to use `Refcount`

[Gary Guo]

From: Gary Guo <gary@garyguo.net>

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      | 45 +++++++++++-------------------------
 rust/kernel/sync/refcount.rs |  8 +++++++
 2 files changed, 22 insertions(+), 31 deletions(-)

diff --git a/rust/kernel/sync/arc.rs b/rust/kernel/sync/arc.rs
index 09a471fea4da..fc63ded00f6e 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::InPlaceInit,
+    sync::Refcount,
     try_init,
-    types::{ForeignOwnable, Opaque},
+    types::ForeignOwnable,
 };
 use core::{
     alloc::Layout,
@@ -144,7 +144,7 @@ pub struct Arc<T: ?Sized> {
 #[pin_data]
 #[repr(C)]
 pub struct ArcInner<T: ?Sized> {
-    refcount: Opaque<bindings::refcount_t>,
+    refcount: Refcount,
     data: T,
 }
 
@@ -156,7 +156,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
@@ -228,8 +228,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,
         };
 
@@ -347,18 +346,13 @@ 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 this = ManuallyDrop::new(this);
         // SAFETY: We own a refcount, so the pointer is still valid.
-        let refcount = unsafe { this.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)) };
+        if refcount.dec_and_test() {
+            refcount.set(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
@@ -428,14 +422,10 @@ fn as_ref(&self) -> &T {
 
 impl<T: ?Sized> Clone for Arc<T> {
     fn clone(&self) -> Self {
-        // SAFETY: By the type invariant, there is necessarily a reference to the object, so it is
-        // safe to dereference it.
-        let refcount = unsafe { self.ptr.as_ref() }.refcount.get();
-
-        // 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(refcount) };
+        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) }
@@ -444,16 +434,10 @@ 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.
+        let is_zero = unsafe { self.ptr.as_ref().refcount.dec_and_test() };
         if is_zero {
             // The count reached zero, we must free the memory.
             //
@@ -747,8 +731,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 <- pin_init::uninit::<T, AllocError>(),
             }? AllocError),
             flags,
diff --git a/rust/kernel/sync/refcount.rs b/rust/kernel/sync/refcount.rs
index a0fc22f6d645..8e7b9b0c1979 100644
--- a/rust/kernel/sync/refcount.rs
+++ b/rust/kernel/sync/refcount.rs
@@ -71,6 +71,14 @@ pub fn dec(&self) {
     /// must come after.
     ///
     /// Returns true if the resulting refcount is 0, false otherwise.
+    ///
+    /// # Notes
+    ///
+    /// A common pattern of using `Refcount` is to free memory when the reference count reaches
+    /// zero. This means that the reference to `Refcount` could become invalid after calling this
+    /// function. This is fine as long as the reference to `Refcount` is no longer used when this
+    /// function returns `false`. It is not necessary to use raw pointers in this scenario, see
+    /// https://github.com/rust-lang/rust/issues/55005.
     #[inline]
     #[must_use = "use `dec` instead if you do not need to test if it is 0"]
     pub fn dec_and_test(&self) -> bool {
-- 
2.49.0

Re: [PATCH v4 3/5] rust: convert `Arc` to use `Refcount`

[Benno Lossin]

On Sun Jun 22, 2025 at 2:57 PM CEST, Gary Guo wrote:
> @@ -428,14 +422,10 @@ fn as_ref(&self) -> &T {
>  
>  impl<T: ?Sized> Clone for Arc<T> {
>      fn clone(&self) -> Self {
> -        // SAFETY: By the type invariant, there is necessarily a reference to the object, so it is
> -        // safe to dereference it.
> -        let refcount = unsafe { self.ptr.as_ref() }.refcount.get();
> -
> -        // 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(refcount) };
> +        unsafe { self.ptr.as_ref().refcount.inc() };

The `.refcount.inc()` can be outside of the `unsafe` block.

>  
>          // SAFETY: We just incremented the refcount. This increment is now owned by the new `Arc`.
>          unsafe { Self::from_inner(self.ptr) }
> @@ -444,16 +434,10 @@ 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.
> +        let is_zero = unsafe { self.ptr.as_ref().refcount.dec_and_test() };

Ditto.

>          if is_zero {
>              // The count reached zero, we must free the memory.
>              //
> @@ -747,8 +731,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 <- pin_init::uninit::<T, AllocError>(),
>              }? AllocError),
>              flags,
> diff --git a/rust/kernel/sync/refcount.rs b/rust/kernel/sync/refcount.rs
> index a0fc22f6d645..8e7b9b0c1979 100644
> --- a/rust/kernel/sync/refcount.rs
> +++ b/rust/kernel/sync/refcount.rs
> @@ -71,6 +71,14 @@ pub fn dec(&self) {
>      /// must come after.
>      ///
>      /// Returns true if the resulting refcount is 0, false otherwise.
> +    ///
> +    /// # Notes
> +    ///
> +    /// A common pattern of using `Refcount` is to free memory when the reference count reaches
> +    /// zero. This means that the reference to `Refcount` could become invalid after calling this
> +    /// function. This is fine as long as the reference to `Refcount` is no longer used when this
> +    /// function returns `false`. It is not necessary to use raw pointers in this scenario, see
> +    /// https://github.com/rust-lang/rust/issues/55005.

This should be in patch 1?

---
Cheers,
Benno

>      #[inline]
>      #[must_use = "use `dec` instead if you do not need to test if it is 0"]
>      pub fn dec_and_test(&self) -> bool {

[PATCH v4 4/5] rust: block: convert `block::mq` to use `Refcount`

[Gary Guo]

From: Gary Guo <gary@garyguo.net>

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.

Tested-by: David Gow <davidgow@google.com>
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    | 63 ++++++++----------------------
 rust/kernel/sync/refcount.rs       | 14 +++++++
 3 files changed, 34 insertions(+), 50 deletions(-)

diff --git a/rust/kernel/block/mq/operations.rs b/rust/kernel/block/mq/operations.rs
index 864ff379dc91..c399dcaa6740 100644
--- a/rust/kernel/block/mq/operations.rs
+++ b/rust/kernel/block/mq/operations.rs
@@ -10,9 +10,10 @@
     block::mq::Request,
     error::{from_result, Result},
     prelude::*,
+    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.
 ///
@@ -78,7 +79,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`.
@@ -187,7 +188,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 4a5b7ec914ef..683b5f32c620 100644
--- a/rust/kernel/block/mq/request.rs
+++ b/rust/kernel/block/mq/request.rs
@@ -8,13 +8,10 @@
     bindings,
     block::mq::Operations,
     error::Result,
+    sync::Refcount,
     types::{ARef, AlwaysRefCounted, Opaque},
 };
-use core::{
-    marker::PhantomData,
-    ptr::NonNull,
-    sync::atomic::{AtomicU64, Ordering},
-};
+use core::{marker::PhantomData, ptr::NonNull, sync::atomic::Ordering};
 
 /// A wrapper around a blk-mq [`struct request`]. This represents an IO request.
 ///
@@ -37,6 +34,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 the driver can still obtain new `ARef` even if there is no `ARef`s in existence by
+/// using `tag_to_rq`, hence the need to distinguish 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,8 +98,11 @@ 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(
+        // 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 let Err(_old) = this.wrapper_ref().refcount().as_atomic().compare_exchange(
             2,
             0,
             Ordering::Relaxed,
@@ -168,13 +171,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 +187,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 { &raw mut (*this).refcount }
@@ -200,47 +203,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 +221,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 8e7b9b0c1979..bea063ef3fe3 100644
--- a/rust/kernel/sync/refcount.rs
+++ b/rust/kernel/sync/refcount.rs
@@ -4,6 +4,8 @@
 //!
 //! C header: [`include/linux/refcount.h`](srctree/include/linux/refcount.h)
 
+use core::sync::atomic::AtomicI32;
+
 use crate::types::Opaque;
 
 /// Atomic reference counter.
@@ -30,6 +32,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().cast();
+        // 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 being considered saturated and "bad".
+        unsafe { &*ptr }
+    }
+
     /// Set a refcount's value.
     #[inline]
     pub fn set(&self, value: i32) {
-- 
2.49.0

[PATCH v4 5/5] MAINTAINERS: update atomic infrastructure entry to include Rust

[Gary Guo]

From: Gary Guo <gary@garyguo.net>

I would like to help review atomic related patches, especially Rust
related ones, hence add myself as an reviewer.

Suggested-by: Boqun Feng <boqun.feng@gmail.com>
Acked-by: Boqun Feng <boqun.feng@gmail.com>
Signed-off-by: Gary Guo <gary@garyguo.net>
---
 MAINTAINERS | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/MAINTAINERS b/MAINTAINERS
index 0c1d245bf7b8..85305709b8c6 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -3896,12 +3896,14 @@ M:	Will Deacon <will@kernel.org>
 M:	Peter Zijlstra <peterz@infradead.org>
 R:	Boqun Feng <boqun.feng@gmail.com>
 R:	Mark Rutland <mark.rutland@arm.com>
+R:	Gary Guo <gary@garyguo.net>
 L:	linux-kernel@vger.kernel.org
 S:	Maintained
 F:	Documentation/atomic_*.txt
 F:	arch/*/include/asm/atomic*.h
 F:	include/*/atomic*.h
 F:	include/linux/refcount.h
+F:	rust/kernel/sync/refcount.rs
 F:	scripts/atomic/
 
 ATTO EXPRESSSAS SAS/SATA RAID SCSI DRIVER
-- 
2.49.0