Re: [PATCH v6 6/7] rust: ww_mutex/exec: add high-level API

[Elle Rhumsaa <elle@weathered-steel.dev>]

On Wed, Sep 03, 2025 at 04:13:12PM +0300, Onur Özkan wrote:
> `ExecContext` is a helper built on top of ww_mutex
> that provides a retrying interface for lock acquisition.
> When `EDEADLK` is hit, it drops all held locks, resets
> the acquire context and retries the given (by the user)
> locking algorithm until it succeeds.
> 
> The API keeps track of acquired locks, cleans them up
> automatically and allows data access to the protected
> data through `with_locked()`. The `lock_all()` helper
> allows implementing multi-mutex algorithms in a simpler
> and less error-prone way while keeping the ww_mutex
> semantics.
> 
> Signed-off-by: Onur Özkan <work@onurozkan.dev>
> ---
>  rust/kernel/sync/lock/ww_mutex.rs      |   2 +
>  rust/kernel/sync/lock/ww_mutex/exec.rs | 176 +++++++++++++++++++++++++
>  2 files changed, 178 insertions(+)
>  create mode 100644 rust/kernel/sync/lock/ww_mutex/exec.rs
> 
> diff --git a/rust/kernel/sync/lock/ww_mutex.rs b/rust/kernel/sync/lock/ww_mutex.rs
> index b415d6deae9b..7de6578513e5 100644
> --- a/rust/kernel/sync/lock/ww_mutex.rs
> +++ b/rust/kernel/sync/lock/ww_mutex.rs
> @@ -16,6 +16,8 @@
>  use core::cell::UnsafeCell;
>  use core::marker::PhantomData;
> 
> +pub mod exec;
> +
>  /// Create static [`WwClass`] instances.
>  ///
>  /// # Examples
> diff --git a/rust/kernel/sync/lock/ww_mutex/exec.rs b/rust/kernel/sync/lock/ww_mutex/exec.rs
> new file mode 100644
> index 000000000000..2f1fc540f0b8
> --- /dev/null
> +++ b/rust/kernel/sync/lock/ww_mutex/exec.rs
> @@ -0,0 +1,176 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +//! A high-level [`WwMutex`] execution helper.
> +//!
> +//! Provides a retrying lock mechanism on top of [`WwMutex`] and [`WwAcquireCtx`].
> +//! It detects [`EDEADLK`] and handles it by rolling back and retrying the
> +//! user-supplied locking algorithm until success.
> +
> +use crate::prelude::*;
> +use crate::sync::lock::ww_mutex::{WwAcquireCtx, WwClass, WwMutex, WwMutexGuard};
> +use core::ptr;
> +
> +/// High-level execution type for ww_mutex.
> +///
> +/// Tracks a series of locks acquired under a common [`WwAcquireCtx`].
> +/// It ensures proper cleanup and retry mechanism on deadlocks and provides
> +/// type-safe access to locked data via [`with_locked`].
> +///
> +/// Typical usage is through [`lock_all`], which retries a user-supplied
> +/// locking algorithm until it succeeds without deadlock.
> +pub struct ExecContext<'a> {
> +    class: &'a WwClass,
> +    acquire: Pin<KBox<WwAcquireCtx<'a>>>,
> +    taken: KVec<WwMutexGuard<'a, ()>>,
> +}
> +
> +impl<'a> Drop for ExecContext<'a> {
> +    fn drop(&mut self) {
> +        self.release_all_locks();
> +    }
> +}
> +
> +impl<'a> ExecContext<'a> {
> +    /// Creates a new [`ExecContext`] for the given lock class.
> +    ///
> +    /// All locks taken through this context must belong to the same class.
> +    ///
> +    /// TODO: Add some safety mechanism to ensure classes are not different.
> +    pub fn new(class: &'a WwClass) -> Result<Self> {
> +        Ok(Self {
> +            class,
> +            acquire: KBox::pin_init(WwAcquireCtx::new(class), GFP_KERNEL)?,
> +            taken: KVec::new(),
> +        })
> +    }
> +
> +    /// Attempts to lock a [`WwMutex`] and records the guard.
> +    ///
> +    /// Returns [`EDEADLK`] if lock ordering would cause a deadlock.
> +    pub fn lock<T>(&mut self, mutex: &'a WwMutex<'a, T>) -> Result<()> {
> +        let guard = self.acquire.lock(mutex)?;
> +        // SAFETY: Type is erased for storage. Actual access uses `with_locked`
> +        // which safely casts back.
> +        let erased: WwMutexGuard<'a, ()> = unsafe { core::mem::transmute(guard) };
> +        self.taken.push(erased, GFP_KERNEL)?;
> +
> +        Ok(())
> +    }
> +
> +    /// Runs `locking_algorithm` until success with retrying on deadlock.
> +    ///
> +    /// `locking_algorithm` should attempt to acquire all needed locks.
> +    /// If [`EDEADLK`] is detected, this function will roll back, reset
> +    /// the context and retry automatically.
> +    ///
> +    /// Once all locks are acquired successfully, `on_all_locks_taken` is
> +    /// invoked for exclusive access to the locked values. Afterwards, all
> +    /// locks are released.
> +    ///
> +    /// # Example
> +    ///
> +    /// ```
> +    /// use kernel::alloc::KBox;
> +    /// use kernel::c_str;
> +    /// use kernel::prelude::*;
> +    /// use kernel::sync::Arc;
> +    /// use kernel::sync::lock::ww_mutex;
> +    /// use pin_init::stack_pin_init;
> +    ///
> +    /// stack_pin_init!(let class = ww_mutex::WwClass::new_wound_wait(c_str!("lock_all_example")));
> +    ///
> +    /// let mutex1 = Arc::pin_init(ww_mutex::WwMutex::new(0, &class), GFP_KERNEL)?;
> +    /// let mutex2 = Arc::pin_init(ww_mutex::WwMutex::new(0, &class), GFP_KERNEL)?;
> +    /// let mut ctx = KBox::pin_init(ww_mutex::exec::ExecContext::new(&class)?, GFP_KERNEL)?;
> +    ///
> +    /// ctx.lock_all(
> +    ///     |ctx| {
> +    ///         // Try to lock both mutexes.
> +    ///         ctx.lock(&mutex1)?;
> +    ///         ctx.lock(&mutex2)?;
> +    ///
> +    ///         Ok(())
> +    ///     },
> +    ///     |ctx| {
> +    ///         // Safely mutate both values while holding the locks.
> +    ///         ctx.with_locked(&mutex1, |v| *v += 1)?;
> +    ///         ctx.with_locked(&mutex2, |v| *v += 1)?;
> +    ///
> +    ///         Ok(())
> +    ///     },
> +    /// )?;
> +    ///
> +    /// # Ok::<(), Error>(())
> +    /// ```
> +    pub fn lock_all<T, Y, Z>(
> +        &mut self,
> +        mut locking_algorithm: T,
> +        mut on_all_locks_taken: Y,
> +    ) -> Result<Z>
> +    where
> +        T: FnMut(&mut ExecContext<'a>) -> Result<()>,
> +        Y: FnMut(&mut ExecContext<'a>) -> Result<Z>,
> +    {
> +        loop {
> +            match locking_algorithm(self) {
> +                Ok(()) => {
> +                    // All locks in `locking_algorithm` succeeded.
> +                    // The user can now safely use them in `on_all_locks_taken`.
> +                    let res = on_all_locks_taken(self);
> +                    self.release_all_locks();
> +
> +                    return res;
> +                }
> +                Err(e) if e == EDEADLK => {
> +                    // Deadlock detected, retry from scratch.
> +                    self.cleanup_on_deadlock()?;
> +                    continue;
> +                }
> +                Err(e) => {
> +                    return Err(e);
> +                }
> +            }
> +        }
> +    }
> +
> +    /// Executes `f` with a mutable reference to the data behind `mutex`.
> +    ///
> +    /// Fails with [`EINVAL`] if the mutex was not locked in this context.
> +    pub fn with_locked<T, Y>(
> +        &mut self,
> +        mutex: &'a WwMutex<'a, T>,
> +        f: impl FnOnce(&mut T) -> Y,
> +    ) -> Result<Y> {
> +        // Find the matching guard.
> +        for guard in &mut self.taken {
> +            if mutex.as_ptr() == guard.mutex.as_ptr() {
> +                // SAFETY: We know this guard belongs to `mutex` and holds the lock.
> +                let typed = unsafe { &mut *ptr::from_mut(guard).cast::<WwMutexGuard<'a, T>>() };
> +                return Ok(f(&mut **typed));
> +            }
> +        }
> +
> +        // `mutex` isn't locked in this `ExecContext`.
> +        Err(EINVAL)
> +    }

This can be refactored using a more idiomatic functional style:

diff --git a/rust/kernel/sync/lock/ww_mutex/exec.rs b/rust/kernel/sync/lock/ww_mutex/exec.rs
index 543c5218232a..8fdfb11104b9 100644
--- a/rust/kernel/sync/lock/ww_mutex/exec.rs
+++ b/rust/kernel/sync/lock/ww_mutex/exec.rs
@@ -142,16 +142,17 @@ pub fn with_locked<T, Y>(
         f: impl FnOnce(&mut T) -> Y,
     ) -> Result<Y> {
         // Find the matching guard.
-        for guard in &mut self.taken {
-            if mutex.as_ptr() == guard.mutex.as_ptr() {
+        self.taken
+            .iter_mut()
+            .find_map(|guard| if mutex.as_ptr() == guard.mutex.as_ptr() {
                 // SAFETY: We know this guard belongs to `mutex` and holds the lock.
-                let typed = unsafe { &mut *ptr::from_mut(guard).cast::<WwMutexGuard<'a, T>>() };
-                return Ok(f(&mut **typed));
-            }
-        }
-
-        // `mutex` isn't locked in this `ExecContext`.
-        Err(EINVAL)
+                Some(unsafe { &mut *ptr::from_mut(guard).cast::<WwMutexGuard<'a, T>>() })
+            } else {
+                None
+            })
+            .map(|typed| f(&mut **typed))
+            // `mutex` isn't locked in this `ExecContext`.
+            .ok_or(EINVAL)
     }

     /// Releases all currently held locks in this context.

> +
> +    /// Releases all currently held locks in this context.
> +    ///
> +    /// It is intended to be used for internal implementation only.
> +    fn release_all_locks(&mut self) {
> +        self.taken.clear();
> +    }
> +
> +    /// Resets this context after a deadlock detection.
> +    ///
> +    /// Drops all held locks and reinitializes the [`WwAcquireCtx`].
> +    ///
> +    /// It is intended to be used for internal implementation only.
> +    fn cleanup_on_deadlock(&mut self) -> Result {
> +        self.release_all_locks();
> +        // Re-init fresh `WwAcquireCtx`.
> +        self.acquire = KBox::pin_init(WwAcquireCtx::new(self.class), GFP_KERNEL)?;
> +
> +        Ok(())
> +    }
> +}
> --
> 2.50.0
> 
> 

Reviewed-by: Elle Rhumsaa <elle@weathered-steel.dev>