1 files changed, 154 insertions, 0 deletions
diff --git a/rust/kernel/kasync/executor.rs b/rust/kernel/kasync/executor.rs
new file mode 100644
index 000000000000..c8dc1cb26bf5
--- /dev/null
+++ b/rust/kernel/kasync/executor.rs
@@ -0,0 +1,154 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Kernel support for executing futures.
+
+use crate::{
+    sync::{Arc, ArcBorrow, LockClassKey},
+    types::PointerWrapper,
+    Result,
+};
+use core::{
+    future::Future,
+    task::{RawWaker, RawWakerVTable, Waker},
+};
+
+pub mod workqueue;
+
+/// Spawns a new task to run in the given executor.
+///
+/// It also automatically defines a new lockdep lock class for executors (e.g., workqueue) that
+/// require one.
+#[macro_export]
+macro_rules! spawn_task {
+    ($executor:expr, $task:expr) => {{
+        static CLASS: $crate::sync::LockClassKey = $crate::sync::LockClassKey::new();
+        $crate::kasync::executor::Executor::spawn($executor, &CLASS, $task)
+    }};
+}
+
+/// A task spawned in an executor.
+pub trait Task {
+    /// Synchronously stops the task.
+    ///
+    /// It ensures that the task won't run again and releases resources needed to run the task
+    /// (e.g., the closure is dropped). If the task is inflight, it waits for the task to block or
+    /// complete before cleaning up and returning.
+    ///
+    /// Callers must not call this from within the task itself as it will likely lead to a
+    /// deadlock.
+    fn sync_stop(self: Arc<Self>);
+}
+
+/// An environment for executing tasks.
+pub trait Executor: Sync + Send {
+    /// Starts executing a task defined by the given future.
+    ///
+    /// Callers are encouraged to use the [`spawn_task`] macro because it automatically defines a
+    /// new lock class key.
+    fn spawn(
+        self: ArcBorrow<'_, Self>,
+        lock_class_key: &'static LockClassKey,
+        future: impl Future + 'static + Send,
+    ) -> Result<Arc<dyn Task>>
+    where
+        Self: Sized;
+
+    /// Stops the executor.
+    ///
+    /// After it is called, attempts to spawn new tasks will result in an error and existing ones
+    /// won't be polled anymore.
+    fn stop(&self);
+}
+
+/// A waker that is wrapped in [`Arc`] for its reference counting.
+///
+/// Types that implement this trait can get a [`Waker`] by calling [`ref_waker`].
+pub trait ArcWake: Send + Sync {
+    /// Wakes a task up.
+    fn wake_by_ref(self: ArcBorrow<'_, Self>);
+
+    /// Wakes a task up and consumes a reference.
+    fn wake(self: Arc<Self>) {
+        self.as_arc_borrow().wake_by_ref();
+    }
+}
+
+/// Creates a [`Waker`] from a [`Arc<T>`], where `T` implements the [`ArcWake`] trait.
+pub fn ref_waker<T: 'static + ArcWake>(w: Arc<T>) -> Waker {
+    fn raw_waker<T: 'static + ArcWake>(w: Arc<T>) -> RawWaker {
+        let data = w.into_pointer();
+        RawWaker::new(
+            data.cast(),
+            &RawWakerVTable::new(clone::<T>, wake::<T>, wake_by_ref::<T>, drop::<T>),
+        )
+    }
+
+    unsafe fn clone<T: 'static + ArcWake>(ptr: *const ()) -> RawWaker {
+        // SAFETY: The data stored in the raw waker is the result of a call to `into_pointer`.
+        let w = unsafe { Arc::<T>::borrow(ptr.cast()) };
+        raw_waker(w.into())
+    }
+
+    unsafe fn wake<T: 'static + ArcWake>(ptr: *const ()) {
+        // SAFETY: The data stored in the raw waker is the result of a call to `into_pointer`.
+        let w = unsafe { Arc::<T>::from_pointer(ptr.cast()) };
+        w.wake();
+    }
+
+    unsafe fn wake_by_ref<T: 'static + ArcWake>(ptr: *const ()) {
+        // SAFETY: The data stored in the raw waker is the result of a call to `into_pointer`.
+        let w = unsafe { Arc::<T>::borrow(ptr.cast()) };
+        w.wake_by_ref();
+    }
+
+    unsafe fn drop<T: 'static + ArcWake>(ptr: *const ()) {
+        // SAFETY: The data stored in the raw waker is the result of a call to `into_pointer`.
+        unsafe { Arc::<T>::from_pointer(ptr.cast()) };
+    }
+
+    let raw = raw_waker(w);
+    // SAFETY: The vtable of the raw waker satisfy the behaviour requirements of a waker.
+    unsafe { Waker::from_raw(raw) }
+}
+
+/// A handle to an executor that automatically stops it on drop.
+pub struct AutoStopHandle<T: Executor + ?Sized> {
+    executor: Option<Arc<T>>,
+}
+
+impl<T: Executor + ?Sized> AutoStopHandle<T> {
+    /// Creates a new instance of an [`AutoStopHandle`].
+    pub fn new(executor: Arc<T>) -> Self {
+        Self {
+            executor: Some(executor),
+        }
+    }
+
+    /// Detaches from the auto-stop handle.
+    ///
+    /// That is, extracts the executor from the handle and doesn't stop it anymore.
+    pub fn detach(mut self) -> Arc<T> {
+        self.executor.take().unwrap()
+    }
+
+    /// Returns the executor associated with the auto-stop handle.
+    ///
+    /// This is so that callers can, for example, spawn new tasks.
+    pub fn executor(&self) -> ArcBorrow<'_, T> {
+        self.executor.as_ref().unwrap().as_arc_borrow()
+    }
+}
+
+impl<T: Executor + ?Sized> Drop for AutoStopHandle<T> {
+    fn drop(&mut self) {
+        if let Some(ex) = self.executor.take() {
+            ex.stop();
+        }
+    }
+}
+
+impl<T: 'static + Executor> From<AutoStopHandle<T>> for AutoStopHandle<dyn Executor> {
+    fn from(src: AutoStopHandle<T>) -> Self {
+        Self::new(src.detach())
+    }
+}