rust: thread: Add Thread support

Signed-off-by: Boqun Feng <[email protected]>

Author: fbq

drivers/char/rust_example.rs

@@ -7,13 +7,16 @@
 #![feature(test)]
 
 use alloc::boxed::Box;
+use alloc::sync::Arc;
 use core::pin::Pin;
+use core::sync::atomic::{AtomicBool, Ordering};
 use kernel::prelude::*;
 use kernel::{
     chrdev, condvar_init, cstr,
     file_operations::FileOperations,
     miscdev, mutex_init, spinlock_init,
     sync::{CondVar, Mutex, SpinLock},
+    thread::{schedule, Thread},
 };
 
 module! {
@@ -133,6 +136,53 @@ impl KernelModule for RustExample {
             cv.free_waiters();
         }
 
+        // Test threads.
+        {
+            let mut a = 1;
+            // FIXME: use a completion or a barrier.
+            let flag = Arc::try_new(AtomicBool::new(false))?;
+            let other = flag.clone();
+
+            let t1 = Thread::try_new(cstr!("rust-thread"), move || {
+                other.store(true, Ordering::Release);
+                let b = Box::try_new(42)?;
+                for _ in 0..20 {
+                    a += 1;
+                    println!("Hello Rust Thread {}", a + b.as_ref());
+                }
+
+                Ok(())
+            })?;
+
+            t1.wake_up();
+
+            // Waits to observe the thread run.
+            while !flag.load(Ordering::Acquire) {
+                schedule();
+            }
+
+            // `t1` should exit normally.
+            t1.stop().expect("Rust thread should exit normally");
+        }
+
+        // Test threads (not up for running).
+        {
+            let mut a = 1;
+
+            let t1 = Thread::try_new(cstr!("rust-thread"), move || {
+                let b = Box::try_new(42)?;
+                for _ in 0..20 {
+                    a += 1;
+                    println!("Hello Rust Thread {}", a + b.as_ref());
+                }
+
+                Ok(())
+            })?;
+
+            // Without `wake_up`, `stop` will cause the thread to exits with `-EINTR`.
+            t1.stop().expect_err("Rust thread should exit abnormally");
+        }
+
         // Including this large variable on the stack will trigger
         // stack probing on the supported archs.
         // This will verify that stack probing does not lead to

rust/helpers.c

@@ -4,6 +4,7 @@
 #include <linux/build_bug.h>
 #include <linux/uaccess.h>
 #include <linux/sched/signal.h>
+#include <linux/sched/task.h>
 
 void rust_helper_BUG(void)
 {
@@ -60,6 +61,18 @@ int rust_helper_signal_pending(void)
 }
 EXPORT_SYMBOL(rust_helper_signal_pending);
 
+void rust_helper_get_task_struct(struct task_struct *task)
+{
+	(void)get_task_struct(task);
+}
+EXPORT_SYMBOL(rust_helper_get_task_struct);
+
+void rust_helper_put_task_struct(struct task_struct *task)
+{
+	put_task_struct(task);
+}
+EXPORT_SYMBOL(rust_helper_put_task_struct);
+
 // See https://github.com/rust-lang/rust-bindgen/issues/1671
 static_assert(__builtin_types_compatible_p(size_t, uintptr_t),
 	"size_t must match uintptr_t, what architecture is this??");

rust/kernel/bindings_helper.h

@@ -10,6 +10,8 @@
 #include <linux/version.h>
 #include <linux/miscdevice.h>
 #include <linux/poll.h>
+#include <linux/kthread.h>
+#include <linux/err.h>
 
 // `bindgen` gets confused at certain things
 const gfp_t BINDINGS_GFP_KERNEL = GFP_KERNEL;

rust/kernel/error.rs

@@ -86,13 +86,13 @@ impl From<AllocError> for Error {
     }
 }
 
-/// Convert a kernel pointer to [`KernelResult`]
+/// Converts a kernel pointer to [`KernelResult`].
 ///
 /// # Pointer value range
 ///
-/// (According to include/linux/err.h)
+/// According to `include/linux/err.h`,
 /// [0, .., `core::usize::MAX - bindings::MAX_ERRNO`) is the range for normal values of pointer,
-/// [`core::unsize::MAX - bindings::MAX_ERRNO`,..,`core::usize::MAX] is the range for error value
+/// [`core::unsize::MAX - bindings::MAX_ERRNO`,..,`core::usize::MAX] is the range for error values
 /// stored in pointer.
 pub fn ptr_to_result<T>(ptr: *mut T) -> Result<*mut T, Error> {
     let value = ptr as usize;

rust/kernel/lib.rs

@@ -45,6 +45,7 @@ pub mod printk;
 pub mod random;
 mod static_assert;
 pub mod sync;
+pub mod thread;
 
 #[cfg(CONFIG_SYSCTL)]
 pub mod sysctl;

rust/kernel/thread.rs

@@ -0,0 +1,245 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! A kernel thread (kthread).
+//!
+//! This module allows Rust code to create/wakup/stop a kernel thread.
+
+use crate::c_types;
+use crate::error::{ptr_to_result, Error, KernelResult};
+use crate::{bindings, cstr, CStr};
+
+use alloc::boxed::Box;
+use core::ops::FnOnce;
+
+extern "C" {
+    #[allow(improper_ctypes)]
+    fn rust_helper_get_task_struct(task: *mut bindings::task_struct);
+    #[allow(improper_ctypes)]
+    fn rust_helper_put_task_struct(task: *mut bindings::task_struct);
+}
+
+/// Function passed to `kthread_create_on_node` as the thread function pointer.
+#[no_mangle]
+unsafe extern "C" fn rust_thread_func(data: *mut c_types::c_void) -> c_types::c_int {
+    // `Box::from_raw()` to get the ownership of the closure.
+    let c = Box::from_raw(data as *mut Box<dyn FnOnce() -> KernelResult<()>>);
+
+    match c() {
+        Ok(_) => 0,
+        Err(e) => e.to_kernel_errno(),
+    }
+}
+
+/// A kernel thread handle.
+pub struct Thread {
+    /// Pointer to the kernel thread.
+    task: *mut bindings::task_struct,
+}
+
+impl Thread {
+    /// Creates a new thread using a C-style function pointer.
+    ///
+    /// No extra memory allocation for thread creation. Use this when closure
+    /// allocation overhead is unacceptable or there is already a C style
+    /// thread function. Otherwise, please consider using [`Thread::try_new`].
+    ///
+    /// # Safety
+    ///
+    /// This function actually doesn't dereference `arg` or call `f`, so even if
+    /// the users pass incorrect parameters this function won't run into
+    /// trouble. But if the users provide incorrect `arg` or `f` the new
+    /// thread will corrupt memory or do other unsafe behaviors, so
+    /// make it `unsafe`. Otherwise, it is [`Thread::wake_up`] that should be
+    /// made as `unsafe` because it will trigger the call of the unsafe function
+    /// `f`, this is undesirable given that [`Thread::wake_up`] on a
+    /// [`Thread::try_new`] created thread is safe.
+    ///
+    /// The safety requirements of calling this function are:
+    ///
+    /// - Make sure `arg` is a proper pointer that points to a valid memory
+    ///   location when the new thread begins to run.
+    ///
+    /// - Make sure `f` is a proper function pointer and `f` handles `arg`
+    ///   correctly.
+    ///
+    /// # Context
+    ///
+    /// This function might sleep due to the memory allocation and waiting for
+    /// completion in `kthread_create_on_node`. Therefore cannot call this
+    /// in atomic contexts (i.e. preemption-off contexts).
+    pub unsafe fn try_new_c_style(
+        name: CStr,
+        f: unsafe extern "C" fn(*mut c_types::c_void) -> c_types::c_int,
+        arg: *mut c_types::c_void,
+    ) -> KernelResult<Self> {
+        let task;
+
+        // SAFETY: `kthread_create_on_node` will copy the content of `name`,
+        // so we don't need to make the `name` live longer.
+        task = ptr_to_result(bindings::kthread_create_on_node(
+            Some(f),
+            arg,
+            bindings::NUMA_NO_NODE,
+            cstr!("%s").as_ptr() as _,
+            name.as_ptr(),
+        ))?;
+
+        // Increases the refcount of the task, so that it won't go away if it `do_exit`.
+        //
+        // SAFETY: `task` is a proper pointer pointing to a newly created thread.
+        rust_helper_get_task_struct(task);
+
+        Ok(Thread { task })
+    }
+
+    /// Creates a new thread.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::thread::Thread;
+    /// use alloc::boxed::Box;
+    ///
+    /// let mut a = 1;
+    ///
+    /// let t = Thread::try_new(
+    ///   move || {
+    ///     let b = Box::try_new(42)?;
+    ///
+    ///     for _ in 0..10 {
+    ///         a = a + 1;
+    ///         println!("Hello Rust Thread {}", a + b.as_ref());
+    ///     }
+    ///     Ok(())
+    ///   },
+    ///   cstr!("rust-thread")
+    /// )?;
+    ///
+    /// t.wake_up();
+    /// ```
+    ///
+    /// # Context
+    ///
+    /// This function might sleep due to the memory allocation and waiting for
+    /// the completion in `kthread_create_on_node`. Therefore do not call this
+    /// in atomic contexts (i.e. preemption-off contexts).
+    pub fn try_new<F>(name: CStr, f: F) -> KernelResult<Self>
+    where
+        F: FnOnce() -> KernelResult<()>,
+        F: Send + 'static,
+    {
+        // Allocate closure here, because this function maybe returns before
+        // `rust_thread_func` (the function that uses the closure) get executed.
+        let boxed_fn: Box<dyn FnOnce() -> KernelResult<()> + 'static> = Box::try_new(f)?;
+
+        // Double boxing here because `dyn FnOnce` is a fat pointer, and we can only
+        // pass a `usize` as the `data` for `kthread_create_on_node`.
+        //
+        // We `Box::into_raw` from this side, and will `Box::from_raw` at the other
+        // side to transfer the ownership of the boxed data.
+        let double_box_ptr = Box::into_raw(Box::try_new(boxed_fn)?) as *mut _;
+
+        // SAFETY: a) `double_box_ptr` is a proper pointer (generated by `Box::into_raw`),
+        // and if succeed, the new thread will get the ownership. And b) `rust_thread_func`
+        // is provided by us and correctly handles the dereference of the `double_box_ptr`
+        // (via `Box::from_raw`).
+        let result = unsafe { Self::try_new_c_style(name, rust_thread_func, double_box_ptr) };
+
+        if let Err(e) = result {
+            // Creation fails, we need to get back the double boxed closure.
+            //
+            // SAFETY: `double_box_ptr` is a proper pointer generated by a `Box::into_raw()`
+            // from a box created by us. If the thread creation fails, no one will reference
+            // that pointer.
+            unsafe {
+                Box::from_raw(double_box_ptr);
+            }
+
+            Err(e)
+        } else {
+            result
+        }
+    }
+
+    /// Wakes up the thread.
+    ///
+    /// Note that a newly created thread (e.g. via [`Thread::try_new`]) will not
+    /// run until a [`Thread::wake_up`] is called.
+    ///
+    /// # Context
+    ///
+    /// This function might sleep, don't call it in atomic contexts.
+    pub fn wake_up(&self) {
+        // SAFETY: `task` is a valid pointer to a kernel thread structure, the refcount
+        // of which is increased in `try_new*`, so it won't point to a freed
+        // `task_struct`. And it's not stopped because `stop` will consume the
+        // [`Thread`].
+        unsafe {
+            bindings::wake_up_process(self.task);
+        }
+    }
+
+    /// Stops the thread.
+    ///
+    /// - If the thread hasn't been waken up after creation, the thread closure
+    ///   won't be called, and will return `-EINTR`. Note that a thread may not
+    ///   be waken up even after [`Thread::wake_up`] is called.
+    ///
+    /// - Otherwise, waits for the closure to return or the thread `do_exit`
+    ///   itself.
+    ///
+    /// Consumes the [`Thread`] so that it's not accessible. In case of error,
+    /// returns the exit code of the thread.
+    ///
+    /// # Context
+    ///
+    /// This function might sleep, don't call it in atomic contexts.
+    pub fn stop(self) -> KernelResult<()> {
+        // SAFETY: `task` is a valid pointer to a kernel thread structure, the
+        // refcount of which is increased in `try_new*`, so it won't point to
+        // a freed `task_struct`. And it's not stopped because `stop` will
+        // consume the [`Thread`].
+        let ret = unsafe { bindings::kthread_stop(self.task) };
+
+        if ret == 0 {
+            Ok(())
+        } else {
+            Err(Error::from_kernel_errno(ret))
+        }
+    }
+}
+
+impl Drop for Thread {
+    fn drop(&mut self) {
+        // Decreases the refcount of the thread, the thread may still be running after
+        // we `drop` the `Thread`.
+        //
+        // SAFETY: At least one refcount is held by `Thread::try_new*` and
+        // the refcount of `task_struct` is implemented by atomics.
+        unsafe {
+            rust_helper_put_task_struct(self.task);
+        }
+    }
+}
+
+/// Tries to give up the CPU and lets another thread run.
+///
+/// This maps to kernel's `schedule` function, which is similar to
+/// [`std::thread::yield_now`].
+///
+/// # Context
+///
+/// This function might sleep, don't call in atomic contexts.
+///
+/// [`std::thread::yield_now`]: https://doc.rust-lang.org/std/thread/fn.yield_now.html
+#[doc(alias = "yield_now")]
+pub fn schedule() {
+    // SAFETY: ...
+    // If we can schedule back from other thread, then this can be treated as a
+    // no-op. A special case occurs when a thread sets its state to `TASK_DEAD`,
+    // and then `schedule` will not come. Currently we don't have a way to do
+    // this safely in Rust, and in the future, we probably still won't allow it.
+    unsafe {
+        bindings::schedule();
+    }
+}