[PATCH v8 2/3] rust: sync: Introduce lock::Backend::Context and lock::BackendWithContext

[Lyude Paul <lyude@redhat.com>]

Now that we've introduced an `IrqDisabled` token for marking contexts in
which IRQs are disabled, we need a way to be able to pass it to locks that
require that IRQs are disabled. In order to continue using the
`lock::Backend` type instead of inventing our own thing, we accomplish this
by adding the associated Context type, along with a `lock_with()` function
that can accept a Context when acquiring a lock. To allow current users of
context-less locks to keep using the normal `lock()` method, we take an
example from Wedson Almeida Filho's work and add a `where T<'a>: Default`
bound to `lock()` so that it can only be called on lock types where the
context is simply a placeholder value, then re-implement it through the new
`lock_with()` function.

We additionally add a BackendWithContext trait, to handle calling the
various locking primatives necessary for these types - along with providing
a `lock_with_new()` function for using those primitives and creating the
relevant context tokens.

Signed-off-by: Lyude Paul <lyude@redhat.com>
Suggested-by: Benno Lossin <benno.lossin@proton.me>

---

V3:
* Use explicit lifetimes in lock_with() to ensure self and _context have
  the same lifetime (Benno)
* Use () for locks that don't need a Context instead of PhantomData (Benno)
V4:
* Fix typo (Dirk)
V7:
* Introduce BackendWithContext and lock_with_new following tglx's feedback
* Name functions lock_with_context_saved and unlock_with_context_restored

Signed-off-by: Lyude Paul <lyude@redhat.com>
---
 rust/kernel/sync/lock.rs          | 118 ++++++++++++++++++++++++++++--
 rust/kernel/sync/lock/mutex.rs    |   1 +
 rust/kernel/sync/lock/spinlock.rs |   1 +
 3 files changed, 115 insertions(+), 5 deletions(-)

diff --git a/rust/kernel/sync/lock.rs b/rust/kernel/sync/lock.rs
index 7b4859b05d3a7..8949a69dd53c5 100644
--- a/rust/kernel/sync/lock.rs
+++ b/rust/kernel/sync/lock.rs
@@ -12,7 +12,7 @@
     str::CStr,
     types::{NotThreadSafe, Opaque, ScopeGuard},
 };
-use core::{cell::UnsafeCell, marker::PhantomPinned};
+use core::{cell::UnsafeCell, marker::PhantomPinned, mem::ManuallyDrop};
 use macros::pin_data;
 
 pub mod mutex;
@@ -43,6 +43,11 @@ pub unsafe trait Backend {
     /// [`unlock`]: Backend::unlock
     type GuardState;
 
+    /// The context which must be provided to acquire the lock.
+    ///
+    /// Can be `()`, or another type if [`BackendWithContext`] is implemented.
+    type Context<'a>;
+
     /// Initialises the lock.
     ///
     /// # Safety
@@ -89,6 +94,54 @@ unsafe fn relock(ptr: *mut Self::State, guard_state: &mut Self::GuardState) {
     }
 }
 
+/// An additional trait for [`Backend`] implementations with a non-`()` [`Context`].
+///
+/// Restricts the context in which a [`Lock`] may be locked. It can initially be created using
+/// [`Lock::lock_with_new`], and can be reused to acquire additional [`Lock`] objects using
+/// [`Lock::lock_with`].
+///
+/// An example of a locking context would be a context in which local CPU interrupts are disabled,
+/// where we must restrict the ability to acquire the [`Lock`] so that it can only be acquired
+/// within that context.
+///
+/// [`Context`]: `Backend::Context`
+pub trait BackendWithContext: Backend {
+    /// The type that represents the state of this [`Context`].
+    ///
+    /// [`Context`]: Backend::Context
+    type ContextState;
+
+    /// Fulfills the invariants of [`State`] and acquires the lock, making the caller its owner.
+    ///
+    /// This returns any state data ([`Context::State`]) needed upon unlock.
+    ///
+    /// # Safety
+    ///
+    /// * Same as [`Backend::lock`].
+    ///
+    /// [`State`]: Context::State
+    unsafe fn lock_with_context_saved<'a>(
+        ptr: *mut Self::State,
+    ) -> (Self::Context<'a>, Self::ContextState, Self::GuardState);
+
+    /// Performs the final unlock within [`Context`].
+    ///
+    /// Passes the [`Context::State`] returned from [`first_lock`].
+    ///
+    /// # Safety
+    ///
+    /// * This function may only be called after [`first_lock`].
+    /// * `state` must be the value returned from [`first_lock`].
+    ///
+    /// [`first_lock`]: Backend::first_lock
+    /// [`Context`]: Backend::Context
+    unsafe fn unlock_with_context_restored(
+        ptr: *mut Self::State,
+        guard_state: &Self::GuardState,
+        context_state: Self::ContextState,
+    );
+}
+
 /// A mutual exclusion primitive.
 ///
 /// Exposes one of the kernel locking primitives. Which one is exposed depends on the lock
@@ -132,8 +185,9 @@ pub fn new(t: T, name: &'static CStr, key: &'static LockClassKey) -> impl PinIni
 }
 
 impl<T: ?Sized, B: Backend> Lock<T, B> {
-    /// Acquires the lock and gives the caller access to the data protected by it.
-    pub fn lock(&self) -> Guard<'_, T, B> {
+    /// Acquires the lock with the given context and gives the caller access to the data protected
+    /// by it.
+    pub fn lock_with<'a>(&'a self, _context: B::Context<'a>) -> Guard<'a, T, B> {
         // SAFETY: The constructor of the type calls `init`, so the existence of the object proves
         // that `init` was called.
         let state = unsafe { B::lock(self.state.get()) };
@@ -141,14 +195,68 @@ pub fn lock(&self) -> Guard<'_, T, B> {
         unsafe { Guard::new(self, state) }
     }
 
-    /// Tries to acquire the lock.
+    /// Acquires the lock and gives the caller access to the data protected by it.
+    #[inline]
+    pub fn lock<'a>(&'a self) -> Guard<'a, T, B>
+    where
+        B::Context<'a>: Default,
+    {
+        self.lock_with(Default::default())
+    }
+
+    /// Tries to acquire the lock with the given context.
     ///
     /// Returns a guard that can be used to access the data protected by the lock if successful.
-    pub fn try_lock(&self) -> Option<Guard<'_, T, B>> {
+    pub fn try_lock_with<'a>(&'a self, _context: B::Context<'a>) -> Option<Guard<'a, T, B>> {
         // SAFETY: The constructor of the type calls `init`, so the existence of the object proves
         // that `init` was called.
         unsafe { B::try_lock(self.state.get()).map(|state| Guard::new(self, state)) }
     }
+
+    /// Tries to acquire the lock.
+    ///
+    /// Returns a guard that can be used to access the data protected by the lock if successful.
+    #[inline]
+    pub fn try_lock<'a>(&'a self) -> Option<Guard<'a, T, B>>
+    where
+        B::Context<'a>: Default,
+    {
+        self.try_lock_with(Default::default())
+    }
+}
+
+impl<T: ?Sized, B: BackendWithContext> Lock<T, B> {
+    /// Acquire the lock with a new [`Context`].
+    ///
+    /// Creates a new instance of [`Context`], and then calls `cb` with said [`Context`] and a
+    /// [`Guard`] for `self`. The [`Context`] will be dropped once `cb` finishes, and it may be used
+    /// within `cb` to acquire additional locks.
+    ///
+    /// [`Context`]: Backend::Context
+    pub fn lock_with_new<'a, R>(
+        &self,
+        cb: impl FnOnce(&mut Guard<'_, T, B>, B::Context<'a>) -> R,
+    ) -> R {
+        let ptr = self.state.get();
+
+        // SAFETY: The constructor of the type calls `init`, so the existence of the object proves
+        // that `init` was called.
+        let (context, context_state, guard_state) = unsafe { B::lock_with_context_saved(ptr) };
+
+        // We don't want Guard's destructor to get called, since we'll drop the lock manually with
+        // B::unlock_with_context_restored later. So we store it in a ManuallyDrop and pass it to cb
+        // via reference.
+        // SAFETY: We acquired the lock when we called [`B::lock_with_context_saved`] above.
+        let mut guard = ManuallyDrop::new(unsafe { Guard::new(self, guard_state) });
+
+        let result = cb(&mut guard, context);
+
+        // SAFETY: We called `B::lock_with_context_saved` above, `context_state` was returned from
+        // there.
+        unsafe { B::unlock_with_context_restored(ptr, &guard.state, context_state) };
+
+        result
+    }
 }
 
 /// A lock guard.
diff --git a/rust/kernel/sync/lock/mutex.rs b/rust/kernel/sync/lock/mutex.rs
index 9ce43ccb45158..9a873cb5b438b 100644
--- a/rust/kernel/sync/lock/mutex.rs
+++ b/rust/kernel/sync/lock/mutex.rs
@@ -93,6 +93,7 @@ macro_rules! new_mutex {
 unsafe impl super::Backend for MutexBackend {
     type State = bindings::mutex;
     type GuardState = ();
+    type Context<'a> = ();
 
     unsafe fn init(
         ptr: *mut Self::State,
diff --git a/rust/kernel/sync/lock/spinlock.rs b/rust/kernel/sync/lock/spinlock.rs
index 040dc16975a68..9fbfd96ffba3e 100644
--- a/rust/kernel/sync/lock/spinlock.rs
+++ b/rust/kernel/sync/lock/spinlock.rs
@@ -92,6 +92,7 @@ macro_rules! new_spinlock {
 unsafe impl super::Backend for SpinLockBackend {
     type State = bindings::spinlock_t;
     type GuardState = ();
+    type Context<'a> = ();
 
     unsafe fn init(
         ptr: *mut Self::State,
-- 
2.47.0