Skip to content

Commit ac80c77

Browse files
committed
rust: introduce Task::spawn
This is a simple function that allows us to create and automatically run kernel threads. Signed-off-by: Wedson Almeida Filho <[email protected]>
1 parent f56c5f3 commit ac80c77

File tree

1 file changed

+90
-2
lines changed

1 file changed

+90
-2
lines changed

rust/kernel/task.rs

+90-2
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,12 @@
44
//!
55
//! C header: [`include/linux/sched.h`](../../../../include/linux/sched.h).
66
7-
use crate::{bindings, ARef, AlwaysRefCounted};
8-
use core::{cell::UnsafeCell, marker::PhantomData, ops::Deref, ptr};
7+
use crate::{
8+
bindings, c_str, c_types, error::from_kernel_err_ptr, error::from_kernel_result,
9+
types::PointerWrapper, ARef, AlwaysRefCounted, Result, ScopeGuard,
10+
};
11+
use alloc::boxed::Box;
12+
use core::{cell::UnsafeCell, fmt, marker::PhantomData, ops::Deref, ptr};
913

1014
/// Wraps the kernel's `struct task_struct`.
1115
///
@@ -101,6 +105,90 @@ impl Task {
101105
// SAFETY: By the type invariant, we know that `self.0` is valid.
102106
unsafe { bindings::signal_pending(self.0.get()) != 0 }
103107
}
108+
109+
/// Starts a new kernel thread and runs it.
110+
///
111+
/// # Examples
112+
///
113+
/// Launches 10 threads and waits for them to complete.
114+
///
115+
/// ```
116+
/// use kernel::task::Task;
117+
/// use kernel::sync::{CondVar, Mutex};
118+
/// use core::sync::atomic::{AtomicU32, Ordering};
119+
///
120+
/// kernel::init_static_sync! {
121+
/// static COUNT: Mutex<u32> = 0;
122+
/// static COUNT_IS_ZERO: CondVar;
123+
/// }
124+
///
125+
/// fn threadfn() {
126+
/// pr_info!("Running from thread {}\n", Task::current().pid());
127+
/// let mut guard = COUNT.lock();
128+
/// *guard -= 1;
129+
/// if *guard == 0 {
130+
/// COUNT_IS_ZERO.notify_all();
131+
/// }
132+
/// }
133+
///
134+
/// // Set count to 10 and spawn 10 threads.
135+
/// *COUNT.lock() = 10;
136+
/// for i in 0..10 {
137+
/// Task::spawn(fmt!("test{i}"), threadfn).unwrap();
138+
/// }
139+
///
140+
/// // Wait for count to drop to zero.
141+
/// let mut guard = COUNT.lock();
142+
/// while (*guard != 0) {
143+
/// COUNT_IS_ZERO.wait(&mut guard);
144+
/// }
145+
/// ```
146+
pub fn spawn<T: FnOnce() + Send + 'static>(
147+
name: fmt::Arguments<'_>,
148+
func: T,
149+
) -> Result<ARef<Task>> {
150+
unsafe extern "C" fn threadfn<T: FnOnce() + Send + 'static>(
151+
arg: *mut c_types::c_void,
152+
) -> c_types::c_int {
153+
from_kernel_result! {
154+
// SAFETY: The thread argument is always a `Box<T>` because it is only called
155+
// via the thread creation below.
156+
let bfunc = unsafe { Box::<T>::from_pointer(arg) };
157+
bfunc();
158+
Ok(0)
159+
}
160+
}
161+
162+
let arg = Box::try_new(func)?.into_pointer();
163+
164+
// SAFETY: `arg` was just created with a call to `into_pointer` above.
165+
let guard = ScopeGuard::new(|| unsafe {
166+
Box::<T>::from_pointer(arg);
167+
});
168+
169+
// SAFETY: The function pointer is always valid (as long as the module remains loaded).
170+
// Ownership of `raw` is transferred to the new thread (if one is actually created), so it
171+
// remains valid. Lastly, the C format string is a constant that require formatting as the
172+
// one and only extra argument.
173+
let ktask = from_kernel_err_ptr(unsafe {
174+
bindings::kthread_create_on_node(
175+
Some(threadfn::<T>),
176+
arg as _,
177+
bindings::NUMA_NO_NODE,
178+
c_str!("%pA").as_char_ptr(),
179+
&name as *const _ as *const c_types::c_void,
180+
)
181+
})?;
182+
183+
// SAFETY: Since the kthread creation succeeded and we haven't run it yet, we know the task
184+
// is valid.
185+
let task = unsafe { &*(ktask as *const Task) }.into();
186+
187+
// SAFETY: Since the kthread creation succeeded, we know `ktask` is valid.
188+
unsafe { bindings::wake_up_process(ktask) };
189+
guard.dismiss();
190+
Ok(task)
191+
}
104192
}
105193

106194
// SAFETY: The type invariants guarantee that `Task` is always ref-counted.

0 commit comments

Comments
 (0)