[PATCH v2 13/14] rust: hrtimer: add `schedule_function` to schedule closures

[Andreas Hindborg <a.hindborg@kernel.org>]

Add a way to quickly schedule functions for execution after a certain
duration has elapsed.

Signed-off-by: Andreas Hindborg <a.hindborg@kernel.org>
---
 rust/kernel/hrtimer.rs         | 49 +++++++++++++++++++++++
 rust/kernel/hrtimer/closure.rs | 72 ++++++++++++++++++++++++++++++++++
 2 files changed, 121 insertions(+)
 create mode 100644 rust/kernel/hrtimer/closure.rs

diff --git a/rust/kernel/hrtimer.rs b/rust/kernel/hrtimer.rs
index 1750016b2b22..2fd3e68368da 100644
--- a/rust/kernel/hrtimer.rs
+++ b/rust/kernel/hrtimer.rs
@@ -132,6 +132,52 @@
 //! pr_info!("Flag raised\n");
 //! # Ok::<(), kernel::error::Error>(())
 //! ```
+//!
+//! Using a helper:
+//! ```
+//! use kernel::{
+//!     hrtimer::schedule_function,
+//!     impl_has_timer, new_condvar, new_mutex,
+//!     prelude::*,
+//!     stack_try_pin_init,
+//!     sync::{Arc, CondVar, Mutex},
+//!     time::Ktime,
+//! };
+//!
+//! #[pin_data]
+//! struct Data {
+//!     #[pin]
+//!     flag: Mutex<bool>,
+//!     #[pin]
+//!     cond: CondVar,
+//! }
+//!
+//! impl Data {
+//!     fn new() -> impl PinInit<Self, kernel::error::Error> {
+//!         try_pin_init!(Self {
+//!             flag <- new_mutex!(false),
+//!             cond <- new_condvar!(),
+//!         })
+//!     }
+//! }
+//!
+//! let data = Arc::pin_init(Data::new(), GFP_KERNEL)?;
+//! let data2 = data.clone();
+//!
+//! let handle = schedule_function(Ktime::from_ns(200_000_000), move || {
+//!     pr_info!("Hello from the future");
+//!     *data2.flag.lock() = true;
+//!     data2.cond.notify_all();
+//! });
+//!
+//! let mut guard = data.flag.lock();
+//! while !*guard {
+//!     data.cond.wait(&mut guard);
+//! }
+//!
+//! pr_info!("Flag raised\n");
+//! # Ok::<(), kernel::error::Error>(())
+//! ```
 
 use crate::{init::PinInit, prelude::*, time::Ktime, types::Opaque};
 use core::marker::PhantomData;
@@ -497,5 +543,8 @@ unsafe fn raw_get_timer(ptr: *const Self) ->
 mod tbox;
 
 mod arc;
+mod closure;
 mod pin;
 mod pin_mut;
+
+pub use closure::schedule_function;
diff --git a/rust/kernel/hrtimer/closure.rs b/rust/kernel/hrtimer/closure.rs
new file mode 100644
index 000000000000..96286a12e87a
--- /dev/null
+++ b/rust/kernel/hrtimer/closure.rs
@@ -0,0 +1,72 @@
+// SPDX-License-Identifier: GPL-2.0
+
+use super::{pin_init, tbox::BoxTimerHandle, Timer, TimerCallback, TimerPointer, TimerRestart};
+use crate::{
+    alloc::{flags, Flags},
+    impl_has_timer, new_mutex,
+    prelude::*,
+    sync::Mutex,
+    time::Ktime,
+};
+use macros::pin_data;
+
+#[pin_data]
+pub struct ClosureTimer<T> {
+    #[pin]
+    timer: Timer<ClosureTimer<T>>,
+    #[pin]
+    callback: Mutex<Option<T>>,
+}
+
+impl_has_timer! {
+    impl{T} HasTimer<Self> for ClosureTimer<T> { self.timer }
+}
+
+impl<T> TimerCallback for ClosureTimer<T>
+where
+    T: FnOnce() + 'static,
+{
+    type CallbackTarget<'a> = Pin<Box<ClosureTimer<T>>>;
+    type CallbackPointer<'a> = &'a ClosureTimer<T>;
+
+    fn run(this: Self::CallbackPointer<'_>) -> TimerRestart
+    where
+        Self: Sized,
+    {
+        if let Some(callback) = this.callback.lock().take() {
+            callback();
+        }
+        TimerRestart::NoRestart
+    }
+}
+
+impl<T> ClosureTimer<T>
+where
+    T: FnOnce() + 'static,
+    T: Send,
+    T: Sync,
+{
+    fn new(f: T, flags: Flags) -> Result<Pin<Box<Self>>> {
+        Box::pin_init(
+            pin_init!(
+                Self {
+                    timer <- Timer::new(),
+                    callback <- new_mutex!(Some(f)),
+                }
+            ),
+            flags,
+        )
+    }
+}
+
+/// Schedule `f` for execution after `expires` time.
+pub fn schedule_function<T>(expires: Ktime, f: T) -> Result<BoxTimerHandle<ClosureTimer<T>>>
+where
+    T: FnOnce() + 'static,
+    T: Send,
+    T: Sync,
+{
+    let timer = ClosureTimer::<T>::new(f, flags::GFP_KERNEL)?;
+    let handle = <Pin<Box<ClosureTimer<T>>> as TimerPointer>::schedule(timer, expires);
+    Ok(handle)
+}
-- 
2.46.0