Skip to main content

ax_runtime/thread/lifecycle/
completion.rs

1//! Runtime thread entry, exit, completion, wait, and reap lifecycle.
2
3use super::*;
4
5/// Stores the exit code, marks the current thread exited, and switches away.
6pub fn exit_current(exit_code: i32) -> ! {
7    let current = current_thread_id()
8        .unwrap_or_else(|error| panic!("failed to identify exiting runtime thread: {error}"));
9    let primary = primary_bootstrap_thread()
10        .unwrap_or_else(|| panic!("primary bootstrap thread identity is not initialized"));
11    if primary == current {
12        debug!("main task exited: exit_code={exit_code}");
13        crate::terminate();
14    }
15
16    let exit_permit = ax_task::thread::current::prepare_current_exit()
17        .unwrap_or_else(|error| panic!("failed to prepare scheduler thread exit: {error}"));
18    publish_current_runtime_exit(exit_code)
19        .unwrap_or_else(|error| panic!("failed to publish thread exit: {error}"));
20    ax_task::thread::current::commit_current_exit(exit_permit)
21}
22
23/// Waits for a thread to finish executing without consuming its owning handle.
24///
25/// This split wait operation lets handle registries keep their raw-pointer or
26/// map entry valid while the target still runs. Completion is published by the
27/// exiting thread after its entry function and exit code are final, before the
28/// non-returning scheduler exit. Physical off-CPU completion and final resource
29/// reclamation are separate phases.
30pub fn wait_thread(handle: &ThreadHandle) -> Result<i32, TaskError> {
31    if current_thread_id()? == handle.id() {
32        return Err(TaskError::InvalidConfiguration);
33    }
34    let data = runtime_thread_data(handle)?;
35    data.join_wait
36        .try_wait_until(|| data.exit_completed.load(Ordering::Acquire))?;
37    Ok(data.exit_code.load(Ordering::Acquire))
38}
39
40/// Waits for an exited thread and returns its exit code.
41///
42/// Resource teardown is attempted synchronously once. A late IRQ wake or other
43/// stable header reference may legitimately defer final reclamation, so join
44/// releases its owning handle to the bounded task-system reaper instead of
45/// spinning until unrelated references disappear.
46pub fn join_thread(handle: ThreadHandle) -> Result<i32, TaskError> {
47    let exit_code = wait_thread(&handle)?;
48    match task_system()
49        .ok_or(TaskError::NotInitialized)?
50        .reap_thread_handle(handle)
51    {
52        Ok(()) => {}
53        Err(error) => {
54            let task_error = error.task_error();
55            if !matches!(task_error, TaskError::ThreadBusy | TaskError::NotExited) {
56                return Err(task_error);
57            }
58            drop(error.into_retry_handle());
59        }
60    }
61    Ok(exit_code)
62}
63
64pub(in crate::thread) unsafe extern "C" fn runtime_thread_entry() -> ! {
65    finish_initial_scheduler_switch();
66    let extension = ax_task::thread::current::current_thread_extension()
67        .unwrap_or_else(|error| panic!("kernel thread has no scheduler extension: {error}"))
68        .unwrap_or_else(|| panic!("kernel thread entry is missing runtime data"));
69    let data_raw = extension_data_after_releasing_lease(extension, &RUNTIME_THREAD_EXTENSION_OPS)
70        .unwrap_or_else(|error| panic!("kernel thread extension type is invalid: {error}"));
71    // SAFETY: the ops identity above proves the data pointer was created from
72    // `Box<RuntimeThreadData>`. The registry record keeps it live through exit;
73    // the temporary lease must not survive the non-unwinding exit path.
74    let data = unsafe { &*ptr::with_exposed_provenance::<RuntimeThreadData>(data_raw) };
75    if !data.start.wait_for_activation() {
76        exit_current(0);
77    }
78    let entry = data
79        .entry
80        .lock_irqsave()
81        .take()
82        .unwrap_or_else(|| panic!("kernel thread entry was already consumed"));
83    entry();
84    exit_current(0)
85}
86
87pub(in crate::thread) fn extension_data_after_releasing_lease(
88    extension: ax_task::thread::ThreadExtensionLease,
89    expected_ops: &'static ThreadExtensionOps,
90) -> Result<usize, TaskError> {
91    if !core::ptr::eq(extension.ops(), expected_ops) {
92        return Err(TaskError::InvalidConfiguration);
93    }
94    let extension = unsafe {
95        // SAFETY: the runtime calls this only from the leased running thread's
96        // entry trampoline, and its registry record remains live through exit.
97        extension.release_for_current_thread_entry()
98    };
99    Ok(extension.data())
100}
101
102pub(in crate::thread) fn finish_initial_scheduler_switch() {
103    // SAFETY: both architecture entry trampolines invoke this exactly once as
104    // their first operation after inheriting the scheduler IRQ-guard baton.
105    unsafe { ax_task::runtime::switch::finish_initial_context_switch() }
106        .unwrap_or_else(|error| panic!("failed to complete initial context switch: {error}"));
107}
108
109pub(in crate::thread) unsafe fn release_transferred_extension(extension: Option<ThreadExtension>) {
110    drop(extension);
111}
112
113pub(in crate::thread) fn runtime_thread_data(
114    thread: &ThreadHandle,
115) -> Result<&RuntimeThreadData, TaskError> {
116    let extension = task_system()
117        .ok_or(TaskError::NotInitialized)?
118        .thread_extension(thread)?
119        .ok_or(TaskError::InvalidConfiguration)?;
120    if !core::ptr::eq(extension.ops(), &RUNTIME_THREAD_EXTENSION_OPS) {
121        return Err(TaskError::InvalidConfiguration);
122    }
123    // SAFETY: the checked ops identity belongs exclusively to RuntimeThreadData,
124    // and the returned reference is bounded by the strong caller handle.
125    Ok(unsafe { &*ptr::with_exposed_provenance::<RuntimeThreadData>(extension.data()) })
126}
127
128fn publish_current_runtime_exit(exit_code: i32) -> Result<(), TaskError> {
129    let thread = current_thread_handle()?;
130    let data = runtime_thread_data(&thread)?;
131    data.exit_code.store(exit_code, Ordering::Release);
132    super::extension::publish_runtime_exit_completion(data);
133    Ok(())
134}