[PATCH v6 4/6] rust: xarray: Add support for index reservations

[Philipp Stanner <phasta@kernel.org>]

One often cannot allocate in the kernel with the desired flags, most
notably in atomic context. Pre-allocating the memory is the preferred
solution in such situations.

Add support for xa_reserve() in the Rust abstractions of xarray. Create
a Reservation object similar to a lock-guard, that can be dropped once
the reservation is no longer needed or once the index was stored to.

Signed-off-by: Philipp Stanner <phasta@kernel.org>
---
Please regard this more as an RFC.

I need pre-allocating in XArray for DmaFence. How exactly we achieve
this is open for discussion.


P.
---
 rust/helpers/xarray.c | 10 ++++++++++
 rust/kernel/xarray.rs | 45 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/rust/helpers/xarray.c b/rust/helpers/xarray.c
index 08979b304341..1504ab242949 100644
--- a/rust/helpers/xarray.c
+++ b/rust/helpers/xarray.c
@@ -26,3 +26,13 @@ __rust_helper void rust_helper_xa_unlock(struct xarray *xa)
 {
 	return xa_unlock(xa);
 }
+
+__rust_helper int rust_helper_xa_reserve(struct xarray *xa, unsigned long index, gfp_t flags)
+{
+	return xa_reserve(xa, index, flags);
+}
+
+__rust_helper void rust_helper_xa_release(struct xarray *xa, unsigned long index)
+{
+	xa_release(xa, index);
+}
diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
index 987c9c0c2198..fa5d47935f0a 100644
--- a/rust/kernel/xarray.rs
+++ b/rust/kernel/xarray.rs
@@ -8,7 +8,10 @@
     alloc,
     bindings,
     build_assert::build_assert,
-    error::{Error, Result},
+    error::{
+        Error,
+        Result, //
+    },
     ffi::c_void,
     types::{
         ForeignOwnable,
@@ -147,6 +150,46 @@ pub fn lock(&self) -> Guard<'_, T> {
             _not_send: NotThreadSafe,
         }
     }
+
+    /// Reserve memory in the XArray at the appropriate index. If this call
+    /// succeeds, later storing at the index will not require an allocation.
+    ///
+    /// Loading from reserved entries will return `None`.
+    pub fn reserve(&self, index: usize, gfp: alloc::Flags) -> Result<Reservation<'_, T>> {
+        // SAFETY:
+        // - `self.xa.xa` is always valid by the type invariant.
+        // - The caller holds the lock.
+        //
+        // INVARIANT: `new` came from `T::into_foreign`.
+        let ret = unsafe { bindings::xa_reserve(self.xa.get(), index, gfp.as_raw()) };
+        match ret {
+            0 => Ok(Reservation { xa: self, index }),
+            _ => Err(Error::from_errno(ret)),
+        }
+    }
+}
+
+/// An index reservation object.
+///
+/// As long as you keep this object alive, your reservation will stay valid. If
+/// someone has stored to this index in the meantime, dropping your reservation
+/// does nothing, which is why you can always drop this reservation object once
+/// you performed your store operation.
+///
+/// Refer to the main C xarray documentation for more details.
+pub struct Reservation<'a, T: ForeignOwnable> {
+    xa: &'a XArray<T>,
+    index: usize,
+}
+
+impl<T: ForeignOwnable> Drop for Reservation<'_, T> {
+    fn drop(&mut self) {
+        // SAFETY:
+        // - `self.xa` is always valid by the type invariant.
+        // - If `self.id` is not used or has been stored to by a racing party,
+        //   `xa_release()` will do nothing.
+        unsafe { bindings::xa_release(self.xa.xa.get(), self.index) };
+    }
 }
 
 /// A lock guard.
-- 
2.55.0