[RFC 02/16] rust/sync: add Arc::drop_unless_unique()

[David Rheinsberg <david@readahead.eu>]

Introduce `Arc::drop_unless_unique()`, a rust version of
`refcount_dec_not_one()`.

It is very common to store an object with `struct kref` (or
`refcount_t`) in a locked lookup-tree, but without the tree actually
owning a reference. Instead, whenever the object is destroyed, it
should be removed from the lookup tree.

The problem is, this requires taking the lock whenever releasing a
reference. Otherwise, a racing lookup might see a ref-count of 0. But
this is needlessly expensive.

One solution provided by `struct kref` is `kref_put_mutex()`. This drops
a reference, unless it is the last reference. In that case it takes a
mutex and tries again. If it no longer is the last reference, it simply
unlocks and returns. But if it dropped the last reference, it calls the
destructor with the lock held.

This entire technique is built around `refcount_dec_not_one()`. Expose
this exact feature in Rust, so it can be used to get the same effect as
`kref_put_mutex()`.

Signed-off-by: David Rheinsberg <david@readahead.eu>
---
 rust/kernel/sync/arc.rs      | 21 +++++++++++++++++++++
 rust/kernel/sync/refcount.rs | 16 ++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/rust/kernel/sync/arc.rs b/rust/kernel/sync/arc.rs
index 921e19333b89..ead021a0dd5d 100644
--- a/rust/kernel/sync/arc.rs
+++ b/rust/kernel/sync/arc.rs
@@ -368,6 +368,27 @@ pub fn into_unique_or_drop(this: Self) -> Option<Pin<UniqueArc<T>>> {
             None
         }
     }
+
+    /// Drop this [`Arc`] unless it is the last reference.
+    ///
+    /// When this is the last reference to the `Arc`, it is returned as `Some`
+    /// unmodified. Otherwise, the `Arc` is dropped and [`None`] is returned.
+    ///
+    /// This function will never release the last reference to the object, and
+    /// as such never call its destructor.
+    pub fn drop_unless_unique(this: Self) -> Option<Self> {
+        let this = ManuallyDrop::new(this);
+
+        // SAFETY: We own a refcount, so the pointer is still valid.
+        if unsafe { this.ptr.as_ref() }.refcount.dec_not_one() {
+            // A single ref-count was dropped, but it was not the last. We
+            // manually drop `this` without invoking `Drop`.
+            None
+        } else {
+            // This was a no-op, return the reference to the caller.
+            Some(ManuallyDrop::into_inner(this))
+        }
+    }
 }
 
 // SAFETY: The pointer returned by `into_foreign` was originally allocated as an
diff --git a/rust/kernel/sync/refcount.rs b/rust/kernel/sync/refcount.rs
index 6c7ae8b05a0b..2a65b3e3f961 100644
--- a/rust/kernel/sync/refcount.rs
+++ b/rust/kernel/sync/refcount.rs
@@ -105,6 +105,22 @@ 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.
+    ///
+    /// It will `WARN` on underflow and fail to decrement when saturated.
+    ///
+    /// Provides release memory ordering when succeeding in decrementing the
+    /// refcount.
+    ///
+    /// Returns true if the decrement operation succeeded, false if the
+    /// decrement operation was skipped as the refcount is 1.
+    #[inline]
+    #[must_use = "refcount release is conditional and must be checked"]
+    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.
-- 
2.53.0