[PATCH v8 0/3] rust: Add local_irq abstraction, SpinLockIrq

[PATCH v8 0/3] rust: Add local_irq abstraction, SpinLockIrq

[Lyude Paul]

This adds a token for annotating contexts where IRQs may be disabled on
non-PREEMPT_RT kernels, a way to use these tokens with Lock types, and
introduces bindings for spin_lock_irqsave() and
spin_unlock_irqrestore().

This patch series depends on the NotThreadSafe type from Alice:

https://lore.kernel.org/rust-for-linux/20240808-alice-file-v9-1-2cb7b934e0e1@google.com/

(Please re-review, since there's been some changes here!)

Lyude Paul (3):
  rust: Introduce local_irq module
  rust: sync: Introduce lock::Backend::Context and
    lock::BackendWithContext
  rust: sync: Add SpinLockIrq

 rust/helpers/helpers.c            |   1 +
 rust/helpers/irq.c                |   8 ++
 rust/helpers/spinlock.c           |  14 +++
 rust/kernel/lib.rs                |   1 +
 rust/kernel/local_irq.rs          |  56 +++++++++++
 rust/kernel/sync.rs               |   2 +-
 rust/kernel/sync/lock.rs          | 118 ++++++++++++++++++++++-
 rust/kernel/sync/lock/mutex.rs    |   1 +
 rust/kernel/sync/lock/spinlock.rs | 150 ++++++++++++++++++++++++++++++
 9 files changed, 345 insertions(+), 6 deletions(-)
 create mode 100644 rust/helpers/irq.c
 create mode 100644 rust/kernel/local_irq.rs


base-commit: 6ce162a002657910104c7a07fb50017681bc476c
prerequisite-patch-id: 4e3cfb97c9dc94e99912e67f383497848ac6f81d
-- 
2.47.0

[PATCH v8 1/3] rust: Introduce local_irq module

[Lyude Paul]

This introduces a module for adding marker types to indicate the type of
interrupt context a function is called in. Note that nothing here enables
or disables interrupts on its own, this is simply for verifying correctness
at compile-time.

Signed-off-by: Lyude Paul <lyude@redhat.com>

---

V2:
* Actually make it so that we check whether or not we have interrupts
  disabled with debug assertions
* Fix issues in the documentation (added suggestions, missing periods, made
  sure that all rustdoc examples compile properly)
* Pass IrqDisabled by value, not reference
* Ensure that IrqDisabled is !Send and !Sync using
  PhantomData<(&'a (), *mut ())>
* Add all of the suggested derives from Benno Lossin

V3:
* Use `impl` for FnOnce bounds in with_irqs_disabled()
* Use higher-ranked trait bounds for the lifetime of with_irqs_disabled()
* Wording changes in the documentation for the module itself

V4:
* Use the actual unsafe constructor for IrqDisabled in
  with_irqs_disabled()
* Fix comment style in with_irqs_disabled example
* Check before calling local_irq_restore() in with_irqs_disabled that
  interrupts are still disabled. It would have been nice to do this from a
  Drop implementation like I hoped, but I realized rust doesn't allow that
  for types that implement Copy.
* Document that interrupts can't be re-enabled within the `cb` provided to
  `with_irqs_disabled`, and link to the github issue I just filed about
  this that describes the solution for this.

V5:
* Rebase against rust-next for the helpers split
* Fix typo (enabled -> disabled) - Dirk

V6:
* s/indicate/require/ in IrqDisabled docs
* Reword comment above with_irqs_disabled in code example requested by
  Benno
* Add paragraph to with_irqs_disabled docs requested by Benno
* Apply the comments from Boqun's review for V4 that I missed
* Document type invariants of `IrqDisabled`

V7:
* Change name of module to local_irq.rs
* Remove with_interrupts_disabled()
* Update documentation wording a bit to make mention of PREEMPT_RT

This patch depends on
https://lore.kernel.org/rust-for-linux/20240808-alice-file-v9-1-2cb7b934e0e1@google.com/

Signed-off-by: Lyude Paul <lyude@redhat.com>
---
 rust/helpers/helpers.c   |  1 +
 rust/helpers/irq.c       |  8 ++++++
 rust/kernel/lib.rs       |  1 +
 rust/kernel/local_irq.rs | 56 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 66 insertions(+)
 create mode 100644 rust/helpers/irq.c
 create mode 100644 rust/kernel/local_irq.rs

diff --git a/rust/helpers/helpers.c b/rust/helpers/helpers.c
index 20a0c69d5cc7b..fd70afe5069ca 100644
--- a/rust/helpers/helpers.c
+++ b/rust/helpers/helpers.c
@@ -12,6 +12,7 @@
 #include "build_assert.c"
 #include "build_bug.c"
 #include "err.c"
+#include "irq.c"
 #include "kunit.c"
 #include "mutex.c"
 #include "page.c"
diff --git a/rust/helpers/irq.c b/rust/helpers/irq.c
new file mode 100644
index 0000000000000..d129094cc1940
--- /dev/null
+++ b/rust/helpers/irq.c
@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <linux/irqflags.h>
+
+bool rust_helper_irqs_disabled(void)
+{
+	return irqs_disabled();
+}
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 620de74d128f9..b7e604bc968ce 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -40,6 +40,7 @@
 #[cfg(CONFIG_KUNIT)]
 pub mod kunit;
 pub mod list;
+pub mod local_irq;
 #[cfg(CONFIG_NET)]
 pub mod net;
 pub mod page;
diff --git a/rust/kernel/local_irq.rs b/rust/kernel/local_irq.rs
new file mode 100644
index 0000000000000..e9e82347392c7
--- /dev/null
+++ b/rust/kernel/local_irq.rs
@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Interrupt controls
+//!
+//! This module allows Rust code to annotate functions which can only be called in contexts where
+//! local interrupts on the CPU may be disabled.
+
+use crate::types::NotThreadSafe;
+use bindings;
+use core::marker::*;
+
+/// A token that is only available in contexts where interrupts are disabled on non-PREEMPT_RT
+/// kernels.
+///
+/// [`IrqDisabled`] is marker made available when local processor interrupts are disabled on
+/// non-PREEMPT_RT kernels. A function may require a [`IrqDisabled`] to ensure that functions may
+/// only be run with processor interrupts disabled on such kernels.
+///
+/// This is a marker type; it has no size, and is simply used as a compile-time guarantee that
+/// interrupts are disabled where required.
+///
+/// # Invariants
+///
+/// On kernels where `CONFIG_PREEMPT_RT=n` is set in the kernel config, local processor interrupts
+/// are disabled whenever an object of this type exists.
+#[derive(Copy, Clone, Debug, Ord, Eq, PartialOrd, PartialEq, Hash)]
+pub struct IrqDisabled<'a>(PhantomData<(&'a (), NotThreadSafe)>);
+
+impl IrqDisabled<'_> {
+    /// Create a new [`IrqDisabled`] token in an interrupt disabled context.
+    ///
+    /// This creates a [`IrqDisabled`] token, which can be passed to functions that must
+    /// be run without interrupts on kernels where `CONFIG_PREEMPT_RT=n`. If debug assertions are
+    /// enabled, this function will assert that interrupts match the expected state. Otherwise, it
+    /// has no size or cost at runtime.
+    ///
+    /// # Panics
+    ///
+    /// If debug assertions are enabled, then this function will assert on non-PREEMPT_RT kernels
+    /// that local processor interrupts are disabled upon creation.
+    ///
+    /// # Safety
+    ///
+    /// This function must only be called in contexts where it is known that on a non-PREEMPT_RT
+    /// kernel, local interrupts have been disabled for the current CPU. The user is making a
+    /// promise that they will remain disabled at least until this [`IrqDisabled`] is
+    /// dropped.
+    pub unsafe fn new() -> Self {
+        // SAFETY: FFI call with no special requirements
+        debug_assert!(unsafe { bindings::irqs_disabled() });
+
+        // INVARIANT: The safety requirements of this function ensure that IRQs are disabled on
+        // non-PREEMPT_RT kernels.
+        Self(PhantomData)
+    }
+}
-- 
2.47.0

Re: [PATCH v8 1/3] rust: Introduce local_irq module

[Boqun Feng]

On Fri, Oct 18, 2024 at 07:22:25PM -0400, Lyude Paul wrote:
> This introduces a module for adding marker types to indicate the type of
> interrupt context a function is called in. Note that nothing here enables
> or disables interrupts on its own, this is simply for verifying correctness
> at compile-time.
> 
> Signed-off-by: Lyude Paul <lyude@redhat.com>
> 
> ---
> 
> V2:
> * Actually make it so that we check whether or not we have interrupts
>   disabled with debug assertions
> * Fix issues in the documentation (added suggestions, missing periods, made
>   sure that all rustdoc examples compile properly)
> * Pass IrqDisabled by value, not reference
> * Ensure that IrqDisabled is !Send and !Sync using
>   PhantomData<(&'a (), *mut ())>
> * Add all of the suggested derives from Benno Lossin
> 
> V3:
> * Use `impl` for FnOnce bounds in with_irqs_disabled()
> * Use higher-ranked trait bounds for the lifetime of with_irqs_disabled()
> * Wording changes in the documentation for the module itself
> 
> V4:
> * Use the actual unsafe constructor for IrqDisabled in
>   with_irqs_disabled()
> * Fix comment style in with_irqs_disabled example
> * Check before calling local_irq_restore() in with_irqs_disabled that
>   interrupts are still disabled. It would have been nice to do this from a
>   Drop implementation like I hoped, but I realized rust doesn't allow that
>   for types that implement Copy.
> * Document that interrupts can't be re-enabled within the `cb` provided to
>   `with_irqs_disabled`, and link to the github issue I just filed about
>   this that describes the solution for this.
> 
> V5:
> * Rebase against rust-next for the helpers split
> * Fix typo (enabled -> disabled) - Dirk
> 
> V6:
> * s/indicate/require/ in IrqDisabled docs
> * Reword comment above with_irqs_disabled in code example requested by
>   Benno
> * Add paragraph to with_irqs_disabled docs requested by Benno
> * Apply the comments from Boqun's review for V4 that I missed
> * Document type invariants of `IrqDisabled`
> 
> V7:
> * Change name of module to local_irq.rs
> * Remove with_interrupts_disabled()
> * Update documentation wording a bit to make mention of PREEMPT_RT
> 

So with the current approach in your patch, there is no safe way to
enable interrrupts (or restore the interrupt state). Therefore I don't
think there is a safe version of CondVar::wait(), at least it won't be
easy to use at the current API. Right?

Whereas the POC [1] approach does support CondVar::wait() and a guard
interface, and has a relatively simple (and more Rusty) design. To
Thomas' point [2], I think we should look at it (interrupt disabling and
nested interrupt disabling) from Rust perspective and provide a cleaner
API in Rust.

And I think a guard interface (and we can even have a guard interface
for interrupt disabling itself) is what we *want* to provide to the
users, the fact that we cannot provide it currently is something we want
to workaround. Maybe we should focus on providing these better APIs with
acceptable small workarounds if necessary first.

Anyway, just make some of my points. Looking forward to hearing your
feedbacks!

Regards,
Boqun

[1]: https://lore.kernel.org/rust-for-linux/ce05b149-63ca-46c9-8eb3-67ff3dc5b2c5@gmail.com/
[2]: https://lore.kernel.org/rust-for-linux/87iktrahld.ffs@tglx/
[3]: https://lore.kernel.org/rust-for-linux/20241018055125.2784186-2-boqun.feng@gmail.com/

> This patch depends on
> https://lore.kernel.org/rust-for-linux/20240808-alice-file-v9-1-2cb7b934e0e1@google.com/
> 
> Signed-off-by: Lyude Paul <lyude@redhat.com>
> ---
>  rust/helpers/helpers.c   |  1 +
>  rust/helpers/irq.c       |  8 ++++++
>  rust/kernel/lib.rs       |  1 +
>  rust/kernel/local_irq.rs | 56 ++++++++++++++++++++++++++++++++++++++++
>  4 files changed, 66 insertions(+)
>  create mode 100644 rust/helpers/irq.c
>  create mode 100644 rust/kernel/local_irq.rs
> 
> diff --git a/rust/helpers/helpers.c b/rust/helpers/helpers.c
> index 20a0c69d5cc7b..fd70afe5069ca 100644
> --- a/rust/helpers/helpers.c
> +++ b/rust/helpers/helpers.c
> @@ -12,6 +12,7 @@
>  #include "build_assert.c"
>  #include "build_bug.c"
>  #include "err.c"
> +#include "irq.c"
>  #include "kunit.c"
>  #include "mutex.c"
>  #include "page.c"
> diff --git a/rust/helpers/irq.c b/rust/helpers/irq.c
> new file mode 100644
> index 0000000000000..d129094cc1940
> --- /dev/null
> +++ b/rust/helpers/irq.c
> @@ -0,0 +1,8 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +#include <linux/irqflags.h>
> +
> +bool rust_helper_irqs_disabled(void)
> +{
> +	return irqs_disabled();
> +}
> diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
> index 620de74d128f9..b7e604bc968ce 100644
> --- a/rust/kernel/lib.rs
> +++ b/rust/kernel/lib.rs
> @@ -40,6 +40,7 @@
>  #[cfg(CONFIG_KUNIT)]
>  pub mod kunit;
>  pub mod list;
> +pub mod local_irq;
>  #[cfg(CONFIG_NET)]
>  pub mod net;
>  pub mod page;
> diff --git a/rust/kernel/local_irq.rs b/rust/kernel/local_irq.rs
> new file mode 100644
> index 0000000000000..e9e82347392c7
> --- /dev/null
> +++ b/rust/kernel/local_irq.rs
> @@ -0,0 +1,56 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +//! Interrupt controls
> +//!
> +//! This module allows Rust code to annotate functions which can only be called in contexts where
> +//! local interrupts on the CPU may be disabled.
> +
> +use crate::types::NotThreadSafe;
> +use bindings;
> +use core::marker::*;
> +
> +/// A token that is only available in contexts where interrupts are disabled on non-PREEMPT_RT
> +/// kernels.
> +///
> +/// [`IrqDisabled`] is marker made available when local processor interrupts are disabled on
> +/// non-PREEMPT_RT kernels. A function may require a [`IrqDisabled`] to ensure that functions may
> +/// only be run with processor interrupts disabled on such kernels.
> +///
> +/// This is a marker type; it has no size, and is simply used as a compile-time guarantee that
> +/// interrupts are disabled where required.
> +///
> +/// # Invariants
> +///
> +/// On kernels where `CONFIG_PREEMPT_RT=n` is set in the kernel config, local processor interrupts
> +/// are disabled whenever an object of this type exists.
> +#[derive(Copy, Clone, Debug, Ord, Eq, PartialOrd, PartialEq, Hash)]
> +pub struct IrqDisabled<'a>(PhantomData<(&'a (), NotThreadSafe)>);
> +
> +impl IrqDisabled<'_> {
> +    /// Create a new [`IrqDisabled`] token in an interrupt disabled context.
> +    ///
> +    /// This creates a [`IrqDisabled`] token, which can be passed to functions that must
> +    /// be run without interrupts on kernels where `CONFIG_PREEMPT_RT=n`. If debug assertions are
> +    /// enabled, this function will assert that interrupts match the expected state. Otherwise, it
> +    /// has no size or cost at runtime.
> +    ///
> +    /// # Panics
> +    ///
> +    /// If debug assertions are enabled, then this function will assert on non-PREEMPT_RT kernels
> +    /// that local processor interrupts are disabled upon creation.
> +    ///
> +    /// # Safety
> +    ///
> +    /// This function must only be called in contexts where it is known that on a non-PREEMPT_RT
> +    /// kernel, local interrupts have been disabled for the current CPU. The user is making a
> +    /// promise that they will remain disabled at least until this [`IrqDisabled`] is
> +    /// dropped.
> +    pub unsafe fn new() -> Self {
> +        // SAFETY: FFI call with no special requirements
> +        debug_assert!(unsafe { bindings::irqs_disabled() });
> +
> +        // INVARIANT: The safety requirements of this function ensure that IRQs are disabled on
> +        // non-PREEMPT_RT kernels.
> +        Self(PhantomData)
> +    }
> +}
> -- 
> 2.47.0
> 

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

[Lyude Paul]

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

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

[Boqun Feng]

On Fri, Oct 18, 2024 at 07:22:26PM -0400, Lyude Paul wrote:
> 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,

I think this needs to be:

        cb: impl FnOnce(&mut Guard<'_, T, B>, B::Context<'_>) -> R,

i.e. using wildcard life for B::Context, which is equal to:

        cb: impl for<'b> FnOnce(&mut Guard<'b, T, B>, B::Context<'b>) -> R,

, which makes the lifetime of B::Context bound to the closure instead of
`lock_with_new()`. Otherwise, users can leak the `Context` outside:

    let irq_disabled_leak = lock1.lock_with_new(|_, irq_disabled| {
        irq_disabled
    });

playground:

	https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=520dc80444e6f3cc8c5782c7f0095cfe

> +    ) -> 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) };
> +

I think we have a soundness issue here, users can do:

(let's say we support static locks, which is a solid thing we want to
have)

static l1: SpinLockIrq<32> = ...;

    <in a function>

    let l2: &SpinLockIrq<i32> = ..;

    l2.lock_with_new(|guard2, context| {
        let mut guard1 = l1.lock_with(context);
	core::mem::swap(&mut guard1, guard2);
	drop(guard1); // actually unlock l2.
    }) // but when the cb returns, we dropped `l2` as well.

I have played this for a while, looks to me, only a static lock `l1` can
make the code pass the borrow checker, I'm not sure whether it is a
borrow checker implementation limitation, or this is by design. Because
lifetime of `guard1` should be `static` and lifetime of `guard2` should 
be `for<'a>, 'a`, seems they are interchangeable right now? A simplified
example:

	https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=8bcc1132734a7bd2420f766524af56a7

Regards,
Boqun

> +        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
> 

[PATCH v8 3/3] rust: sync: Add SpinLockIrq

[Lyude Paul]

A variant of SpinLock that is expected to be used in noirq contexts, and
thus requires that the user provide an kernel::irq::IrqDisabled to prove
they are in such a context upon lock acquisition. This is the rust
equivalent of spin_lock_irqsave()/spin_lock_irqrestore().

Signed-off-by: Lyude Paul <lyude@redhat.com>

---

V2:
* s/IrqSpinLock/SpinLockIrq/
* Implement `lock::Backend` now that we have `Context`
* Add missing periods
* Make sure rustdoc examples compile correctly
* Add documentation suggestions
V8:
* Drop a completely unrelated change that got staged into this patch by
  mistake

Signed-off-by: Lyude Paul <lyude@redhat.com>
---
 rust/helpers/spinlock.c           |  14 +++
 rust/kernel/sync.rs               |   2 +-
 rust/kernel/sync/lock/spinlock.rs | 149 ++++++++++++++++++++++++++++++
 3 files changed, 164 insertions(+), 1 deletion(-)

diff --git a/rust/helpers/spinlock.c b/rust/helpers/spinlock.c
index 775ed4d549aef..f4108d2d78648 100644
--- a/rust/helpers/spinlock.c
+++ b/rust/helpers/spinlock.c
@@ -27,3 +27,17 @@ int rust_helper_spin_trylock(spinlock_t *lock)
 {
 	return spin_trylock(lock);
 }
+
+size_t rust_helper_spin_lock_irqsave(spinlock_t *lock)
+{
+	size_t flags = 0;
+
+	spin_lock_irqsave(lock, flags);
+
+	return flags;
+}
+
+void rust_helper_spin_unlock_irqrestore(spinlock_t *lock, size_t flags)
+{
+	spin_unlock_irqrestore(lock, flags);
+}
diff --git a/rust/kernel/sync.rs b/rust/kernel/sync.rs
index 0ab20975a3b5d..b028ee325f2a6 100644
--- a/rust/kernel/sync.rs
+++ b/rust/kernel/sync.rs
@@ -15,7 +15,7 @@
 pub use arc::{Arc, ArcBorrow, UniqueArc};
 pub use condvar::{new_condvar, CondVar, CondVarTimeoutResult};
 pub use lock::mutex::{new_mutex, Mutex};
-pub use lock::spinlock::{new_spinlock, SpinLock};
+pub use lock::spinlock::{new_spinlock, new_spinlock_irq, SpinLock, SpinLockIrq};
 pub use locked_by::LockedBy;
 
 /// Represents a lockdep class. It's a wrapper around C's `lock_class_key`.
diff --git a/rust/kernel/sync/lock/spinlock.rs b/rust/kernel/sync/lock/spinlock.rs
index 9fbfd96ffba3e..b376920424dc6 100644
--- a/rust/kernel/sync/lock/spinlock.rs
+++ b/rust/kernel/sync/lock/spinlock.rs
@@ -3,6 +3,7 @@
 //! A kernel spinlock.
 //!
 //! This module allows Rust code to use the kernel's `spinlock_t`.
+use kernel::local_irq::*;
 
 /// Creates a [`SpinLock`] initialiser with the given name and a newly-created lock class.
 ///
@@ -127,3 +128,151 @@ unsafe fn try_lock(ptr: *mut Self::State) -> Option<Self::GuardState> {
         }
     }
 }
+
+/// Creates a [`SpinLockIrq`] initialiser with the given name and a newly-created lock class.
+///
+/// It uses the name if one is given, otherwise it generates one based on the file name and line
+/// number.
+#[macro_export]
+macro_rules! new_spinlock_irq {
+    ($inner:expr $(, $name:literal)? $(,)?) => {
+        $crate::sync::SpinLockIrq::new(
+            $inner, $crate::optional_name!($($name)?), $crate::static_lock_class!())
+    };
+}
+pub use new_spinlock_irq;
+
+/// A spinlock that may be acquired when interrupts are disabled.
+///
+/// A version of [`SpinLock`] that can only be used in contexts where interrupts for the local CPU
+/// are disabled. It requires that the user acquiring the lock provide proof that interrupts are
+/// disabled through [`IrqDisabled`].
+///
+/// For more info, see [`SpinLock`].
+///
+/// # Examples
+///
+/// The following example shows how to declare, allocate initialise and access a struct (`Example`)
+/// that contains two inner structs of type `Inner` that are protected by separate spinlocks.
+///
+/// ```
+/// use kernel::{
+///     sync::{new_spinlock_irq, SpinLockIrq},
+///     local_irq::IrqDisabled
+/// };
+///
+/// struct Inner {
+///     a: u32,
+///     b: u32,
+/// }
+///
+/// #[pin_data]
+/// struct Example {
+///     c: u32,
+///     #[pin]
+///     first: SpinLockIrq<Inner>,
+///     #[pin]
+///     second: SpinLockIrq<Inner>,
+/// }
+///
+/// impl Example {
+///     fn new() -> impl PinInit<Self> {
+///         pin_init!(Self {
+///             c: 10,
+///             first <- new_spinlock_irq!(Inner { a: 20, b: 30 }),
+///             second <- new_spinlock_irq!(Inner { a: 10, b: 20 }),
+///         })
+///     }
+/// }
+///
+/// // Allocate a boxed `Example`
+/// let example = KBox::pin_init(Example::new(), GFP_KERNEL)?;
+///
+/// // Accessing an `Inner` from a context where we don't have a `LocalInterruptsDisabled` token
+/// // already.
+/// let bb = example.first.lock_with_new(|first, irq| {
+///     assert_eq!(example.c, 10);
+///     assert_eq!(first.a, 20);
+///
+///     // Since we already have a `LocalInterruptsDisabled` token, we can reuse it to acquire the
+///     // second lock. This allows us to skip changing the local interrupt state unnecessarily on
+///     // non-PREEMPT_RT kernels.
+///     let second = example.second.lock_with(irq);
+///     assert_eq!(second.a, 10);
+///
+///     (first.b, second.b)
+/// });
+///
+/// assert_eq!(bb, (30, 20));
+/// # Ok::<(), Error>(())
+/// ```
+pub type SpinLockIrq<T> = super::Lock<T, SpinLockIrqBackend>;
+
+/// A kernel `spinlock_t` lock backend that is acquired in no-irq contexts.
+pub struct SpinLockIrqBackend;
+
+unsafe impl super::Backend for SpinLockIrqBackend {
+    type State = bindings::spinlock_t;
+    type GuardState = ();
+    type Context<'a> = IrqDisabled<'a>;
+
+    unsafe fn init(
+        ptr: *mut Self::State,
+        name: *const core::ffi::c_char,
+        key: *mut bindings::lock_class_key,
+    ) {
+        // SAFETY: The safety requirements ensure that `ptr` is valid for writes, and `name` and
+        // `key` are valid for read indefinitely.
+        unsafe { bindings::__spin_lock_init(ptr, name, key) }
+    }
+
+    unsafe fn lock(ptr: *mut Self::State) -> Self::GuardState {
+        // SAFETY: The safety requirements of this function ensure that `ptr` points to valid
+        // memory, and that it has been initialised before.
+        unsafe { bindings::spin_lock(ptr) }
+    }
+
+    unsafe fn unlock(ptr: *mut Self::State, _guard_state: &Self::GuardState) {
+        // SAFETY: The safety requirements of this function ensure that `ptr` is valid and that the
+        // caller is the owner of the spinlock.
+        unsafe { bindings::spin_unlock(ptr) }
+    }
+
+    unsafe fn try_lock(ptr: *mut Self::State) -> Option<Self::GuardState> {
+        // SAFETY: The `ptr` pointer is guaranteed to be valid and initialized before use.
+        let result = unsafe { bindings::spin_trylock(ptr) };
+
+        if result != 0 {
+            Some(())
+        } else {
+            None
+        }
+    }
+}
+
+impl super::BackendWithContext for SpinLockIrqBackend {
+    type ContextState = usize;
+
+    unsafe fn lock_with_context_saved<'a>(
+        ptr: *mut Self::State,
+    ) -> (Self::Context<'a>, Self::ContextState, Self::GuardState) {
+        // SAFETY: The safety requirements of this function ensure that `ptr` points to valid
+        // memory, and that it has been initialised before.
+        let flags = unsafe { bindings::spin_lock_irqsave(ptr) };
+
+        // SAFETY: We just disabled interrupts above
+        let context = unsafe { IrqDisabled::new() };
+
+        (context, flags, ())
+    }
+
+    unsafe fn unlock_with_context_restored(
+        ptr: *mut Self::State,
+        _guard_state: &Self::GuardState,
+        context_state: Self::ContextState,
+    ) {
+        // SAFETY: The safety requirements of this function ensure that `ptr` is valid and that the
+        // caller is the owner of the spinlock.
+        unsafe { bindings::spin_unlock_irqrestore(ptr, context_state) }
+    }
+}
-- 
2.47.0