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! {
@@ -125,6 +128,59 @@ impl KernelModule for RustExample {
             cv.free_waiters();
         }
 
+        // Test threads.
+        {
+            let mut a = 1;
+            // XXX use a completion or Barrier
+            let flag = Arc::try_new(AtomicBool::new(false))?;
+            let other = flag.clone();
+
+            let t1 = Thread::new(
+                move || {
+                    other.store(true, Ordering::Release);
+                    let b = Box::try_new(42)?;
+                    for _ in 0..20 {
+                        a = a + 1;
+                        println!("Hello Rust Thread {}", a + b.as_ref());
+                    }
+
+                    Ok(())
+                },
+                cstr!("Rust thread"),
+            )?;
+
+            t1.wake_up();
+
+            // Wait 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.
+        {
+            let mut a = 1;
+
+            let t1 = Thread::new(
+                move || {
+                    let b = Box::try_new(42)?;
+                    for _ in 0..20 {
+                        a = a + 1;
+                        println!("Hello Rust Thread {}", a + b.as_ref());
+                    }
+
+                    Ok(())
+                },
+                cstr!("Rust thread"),
+            )?;
+
+            // Without `wake_up`, `stop` will cause the thread exits -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/lib.rs

@@ -42,6 +42,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,200 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! A kernel thread (kthread)
+//!
+//! This modules 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};
+
+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 function pointer. No other user.
+#[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<()>>);
+
+    let ret = c();
+
+    if ret.is_ok() {
+        return 0;
+    } else {
+        // Already check `ret` is not `Ok`, won't panic here
+        return ret.unwrap_err().to_kernel_errno();
+    }
+}
+
+/// A kernel thread handle
+pub struct Thread {
+    /// Pointer to kernel thread
+    task: *mut bindings::task_struct,
+}
+
+impl Thread {
+    /// Create a new thread
+    ///
+    /// This function might sleep due to the memory allocation and waiting for completion in
+    /// `kthread_create_on_node`. Therefore cannot call this in a atomic context.
+    ///
+    /// # Examples:
+    /// use kernel::thread::Thread;
+    /// use alloc::boxed::Box;
+    ///
+    /// let mut a = 1;
+    ///
+    /// let t = Thread::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();
+    pub fn new<F>(f: F, name: CStr) -> KernelResult<Self>
+    where
+        F: FnOnce() -> KernelResult<()>,
+        F: Send + 'static,
+    {
+        // Allocate closure here, because this function maybe returns before `rust_thread_func`
+        // (the function that use the closure) get executed.
+        let bf: 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 `into_raw` from this side, and will `from_raw` at the other side to transfer the
+        // ownership of the boxed data.
+        let db = Box::into_raw(Box::try_new(bf)?) as *mut _;
+
+        let task;
+
+        // SAFETY:
+        //
+        // * `db` is a proper pointer (generated by `Box::into_raw()`), and if succeed, the new
+        // thread will get the ownership.
+        //
+        // * `kthread_create_on_node` will copy the content of name, so we don't need to make the
+        // name live longer.
+        unsafe {
+            task = bindings::kthread_create_on_node(
+                Some(rust_thread_func),
+                db,
+                bindings::NUMA_NO_NODE,
+                "%s".as_ptr() as _,
+                name.as_ptr(),
+            );
+        }
+
+        let result = ptr_to_result(task);
+
+        if result.is_err() {
+            // Creation fails, we need to get back the double boxed closure.
+            //
+            // SAFETY:
+            //
+            // `db` is a proper pointer generated by a `Box::into_raw()` from a box created by us,
+            // if the thread creation fails, no one will consume that pointer.
+            unsafe {
+                Box::from_raw(db);
+            }
+        } else {
+            // Increase the refcount of the task, so that it won't go away if it `do_exit`.
+            //
+            // SAFETY:
+            //
+            // `task` is a valid pointer to newly created thread
+            unsafe {
+                rust_helper_get_task_struct(task);
+            }
+        }
+
+        Ok(Thread { task: result? })
+    }
+
+    /// Wake up the kernel thread, note that a newly created kthread will not run until `wake_up`
+    /// is called.
+    ///
+    /// This function might sleep, don't call in atomic contexts.
+    pub fn wake_up(&self) {
+        // SAFETY:
+        //
+        // `task` is a valid pointer to a kernel thread structure, and it's not stopped because
+        // `stop` will consume the [`Thread`].
+        unsafe {
+            bindings::wake_up_process(self.task);
+        }
+    }
+
+    /// Stop the the kernel 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 `wake_up`
+    /// is called).
+    /// * Otherwise, wait for the closure to return or the thread `do_exit` itself.
+    ///
+    /// Consume the [`Thread`] so that it's not accessible. Return the result of the thread
+    /// closure (or the exit code in KernelResult format).
+    ///
+    /// This function might sleep, don't call in atomic contexts.
+    pub fn stop(self) -> KernelResult<()> {
+        let ret;
+        // SAFETY:
+        //
+        // `task` is a valid pointer to a kernel thread structure, and it's not stopped because
+        // `stop` will consume the [`Thread`].
+        unsafe { ret = bindings::kthread_stop(self.task) }
+
+        if ret == 0 {
+            Ok(())
+        } else {
+            Err(Error::from_kernel_errno(ret))
+        }
+    }
+}
+
+impl Drop for Thread {
+    fn drop(&mut self) {
+        // Decrease the refcount of the thread, the actuall thread may still be running after we
+        // `drop` the [`Thread`].
+        //
+        // SAFETY:
+        //
+        // At least one refcount is held by `Thread::new` and refcount of [`task_struct`] is
+        // implemented by atomics.
+        unsafe {
+            rust_helper_put_task_struct(self.task);
+        }
+    }
+}
+
+/// Try to give up the cpu and let another thread to run, mapping to kernel's `schedule` function.
+/// A similar concept to `std::thread::yield_now`.
+///
+/// This function might sleep, don't call in atomic contexts.
+pub fn schedule() {
+    // SAFETY:
+    //
+    // If we can schedule back from other thread, then this can be treated as no-ops. A special
+    // case are 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 don't allow this.
+    unsafe {
+        bindings::schedule();
+    }
+}