ax_runtime/thread/lifecycle/
extension.rs

1//! Runtime-owned scheduler extension composition and OS extension leases.
2
3use super::*;
4
5pub(in crate::thread) static RUNTIME_THREAD_EXTENSION_OPS: ThreadExtensionOps =
6    ThreadExtensionOps {
7        on_switch_in: runtime_thread_switch_in_hook,
8        on_switch_out: runtime_thread_switch_out_hook,
9        on_exit: runtime_thread_exit_hook,
10        on_deadline_overrun: runtime_thread_deadline_overrun_hook,
11        drop: runtime_thread_drop_hook,
12    };
13
14pub(in crate::thread) unsafe fn runtime_thread_extension(data: usize) -> ThreadExtension {
15    let os_extension = unsafe { runtime_thread_data_from_raw(data) }
16        .os_extension
17        .as_ref();
18    let scheduler_tick_cpu_time = os_extension.and_then(ThreadExtension::scheduler_tick_cpu_time);
19    let scheduler_tick_gate = os_extension.and_then(ThreadExtension::scheduler_tick_work_gate);
20    let forwards_running_policy = os_extension
21        .and_then(ThreadExtension::running_policy_applied_hook)
22        .is_some();
23    // SAFETY: the caller transfers one live `RuntimeThreadData` allocation
24    // whose final destruction right belongs to this outer extension.
25    let mut extension = unsafe { ThreadExtension::new(data, &RUNTIME_THREAD_EXTENSION_OPS) };
26    if let Some(accounting) = scheduler_tick_cpu_time {
27        extension = extension.with_scheduler_tick_cpu_time(accounting);
28    }
29    if let Some(gate) = scheduler_tick_gate {
30        // SAFETY: the outer callback retains `RuntimeThreadData` and forwards
31        // exactly one generation-authorized publication, with the IRQ-observed
32        // monotonic timestamp, to its inner extension.
33        extension =
34            unsafe { extension.with_scheduler_tick_work(gate, runtime_thread_scheduler_tick_hook) };
35    }
36    if forwards_running_policy {
37        // SAFETY: the outer runtime extension owns the inner extension and
38        // forwards the same running-owner base-policy observation without
39        // retaining either borrowed data value.
40        extension = unsafe {
41            extension.with_running_policy_applied_hook(runtime_thread_policy_applied_hook)
42        };
43    }
44    extension
45}
46
47unsafe extern "Rust" fn runtime_thread_switch_in_hook(
48    data: usize,
49    thread: ThreadId,
50    base_policy: SchedulePolicy,
51    charged_runtime_ns: u64,
52) {
53    let runtime = unsafe { runtime_thread_data_from_raw(data) };
54    if let Some(extension) = runtime.os_extension.as_ref() {
55        // SAFETY: `spawn_raw_with_extension` retains the OS extension until the
56        // outer runtime extension is reaped and forwards the same thread ID.
57        unsafe {
58            (extension.ops().on_switch_in)(
59                extension.data(),
60                thread,
61                base_policy,
62                charged_runtime_ns,
63            )
64        };
65    }
66}
67
68unsafe extern "Rust" fn runtime_thread_switch_out_hook(
69    data: usize,
70    thread: ThreadId,
71    reason: SwitchReason,
72) {
73    let runtime = unsafe { runtime_thread_data_from_raw(data) };
74    if let Some(extension) = runtime.os_extension.as_ref() {
75        // SAFETY: same composition contract as `runtime_thread_switch_in_hook`.
76        unsafe { (extension.ops().on_switch_out)(extension.data(), thread, reason) };
77    }
78}
79
80unsafe extern "Rust" fn runtime_thread_exit_hook(data: usize, thread: ThreadId) {
81    let runtime = unsafe { runtime_thread_data_from_raw(data) };
82    if let Some(extension) = runtime.os_extension.as_ref() {
83        // SAFETY: the TaskSystem invokes this in task context after committing exit.
84        unsafe { (extension.ops().on_exit)(extension.data(), thread) };
85    }
86    // Runtime threads normally publish completion before their final schedule,
87    // Linux-zombie style. Keep this idempotent fallback for externally marked
88    // exits and failed-spawn cleanup paths that never ran the trampoline.
89    publish_runtime_exit_completion(runtime);
90}
91
92pub(super) fn publish_runtime_exit_completion(runtime: &RuntimeThreadData) {
93    if runtime
94        .exit_completed
95        .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
96        .is_ok()
97    {
98        runtime.join_wait.notify_all();
99    }
100}
101
102unsafe extern "Rust" fn runtime_thread_deadline_overrun_hook(data: usize, thread: ThreadId) {
103    let runtime = unsafe { runtime_thread_data_from_raw(data) };
104    if let Some(extension) = runtime.os_extension.as_ref() {
105        // SAFETY: the scheduler defers this callback to an ordinary safe point.
106        unsafe { (extension.ops().on_deadline_overrun)(extension.data(), thread) };
107    }
108}
109
110unsafe extern "Rust" fn runtime_thread_policy_applied_hook(
111    data: usize,
112    thread: ThreadId,
113    base_policy: SchedulePolicy,
114    observed_ns: u64,
115) {
116    let runtime = unsafe { runtime_thread_data_from_raw(data) };
117    let Some(extension) = runtime.os_extension.as_ref() else {
118        panic!(
119            "runtime policy forwarding lost OS extension for thread {:#x}",
120            thread.as_u64()
121        );
122    };
123    if !unsafe { extension.forward_running_policy_applied(thread, base_policy, observed_ns) } {
124        panic!(
125            "runtime policy forwarding lost callback for thread {:#x}",
126            thread.as_u64()
127        );
128    }
129}
130
131unsafe extern "Rust" fn runtime_thread_scheduler_tick_hook(
132    data: usize,
133    thread: ThreadId,
134    observed_ns: u64,
135) -> SchedulerTickWorkDisposition {
136    let runtime = unsafe { runtime_thread_data_from_raw(data) };
137    let Some(extension) = runtime.os_extension.as_ref() else {
138        panic!(
139            "runtime scheduler-tick forwarding lost OS extension for thread {:#x}",
140            thread.as_u64()
141        );
142    };
143    unsafe { extension.forward_scheduler_tick_work(thread, observed_ns) }.unwrap_or_else(|| {
144        panic!(
145            "runtime scheduler-tick forwarding lost callback for thread {:#x}",
146            thread.as_u64()
147        )
148    })
149}
150
151unsafe extern "Rust" fn runtime_thread_drop_hook(data: usize) {
152    // SAFETY: the scheduler reaper invokes this exactly once for the pointer
153    // transferred through `RUNTIME_THREAD_EXTENSION_OPS`.
154    drop(unsafe { Box::from_raw(ptr::with_exposed_provenance_mut::<RuntimeThreadData>(data)) });
155}
156
157unsafe fn runtime_thread_data_from_raw(data: usize) -> &'static RuntimeThreadData {
158    // SAFETY: every outer callback receives the Box pointer installed by
159    // `spawn_raw_with_extension`, which remains valid until the drop callback.
160    unsafe { &*ptr::with_exposed_provenance::<RuntimeThreadData>(data) }
161}
162
163/// Borrows the OS extension composed inside a runtime-owned thread record.
164pub fn thread_os_extension(
165    thread: &ThreadHandle,
166) -> Result<Option<ThreadOsExtensionBorrow<'_>>, TaskError> {
167    let runtime = task_system()
168        .ok_or(TaskError::NotInitialized)?
169        .thread_extension(thread)?;
170    let RuntimeExtensionKind::Runtime = classify_runtime_extension(
171        runtime.as_ref().map(|extension| extension.ops()),
172        runtime.as_ref().map_or(0, |extension| extension.data()),
173    )?
174    else {
175        return Ok(None);
176    };
177    let Some(runtime) = runtime else {
178        unreachable!("classified runtime extension must be present")
179    };
180    // SAFETY: the checked ops identity belongs exclusively to RuntimeThreadData,
181    // and `runtime` borrows the strong caller handle for the whole result.
182    let data = unsafe { runtime_thread_data_from_raw(runtime.data()) };
183    Ok(data
184        .os_extension
185        .as_ref()
186        .map(|extension| ThreadOsExtensionBorrow {
187            data: extension.data(),
188            ops: extension.ops(),
189            _runtime: runtime,
190        }))
191}
192
193/// Leases the current thread's composed OS extension.
194pub fn current_os_extension() -> Result<Option<ThreadOsExtensionLease>, TaskError> {
195    let runtime = current_thread_extension()?;
196    let RuntimeExtensionKind::Runtime = classify_runtime_extension(
197        runtime.as_ref().map(|extension| extension.ops()),
198        runtime.as_ref().map_or(0, |extension| extension.data()),
199    )?
200    else {
201        return Ok(None);
202    };
203    let Some(runtime) = runtime else {
204        unreachable!("classified runtime extension must be present")
205    };
206    // SAFETY: the checked ops identity belongs exclusively to RuntimeThreadData,
207    // and the returned lease retains the outer scheduler extension lease.
208    let data = unsafe { runtime_thread_data_from_raw(runtime.data()) };
209    Ok(data
210        .os_extension
211        .as_ref()
212        .map(|extension| ThreadOsExtensionLease {
213            data: extension.data(),
214            ops: extension.ops(),
215            _runtime: runtime,
216        }))
217}
218
219#[derive(Clone, Copy, Debug, Eq, PartialEq)]
220pub(in crate::thread) enum RuntimeExtensionKind {
221    Missing,
222    Runtime,
223}
224
225pub(in crate::thread) fn classify_runtime_extension(
226    ops: Option<&ThreadExtensionOps>,
227    data: usize,
228) -> Result<RuntimeExtensionKind, TaskError> {
229    let Some(ops) = ops else {
230        return Ok(RuntimeExtensionKind::Missing);
231    };
232    if !core::ptr::eq(ops, &RUNTIME_THREAD_EXTENSION_OPS) {
233        return Err(TaskError::InvalidConfiguration);
234    }
235    if data == 0 || !data.is_multiple_of(core::mem::align_of::<RuntimeThreadData>()) {
236        return Err(TaskError::InvalidRuntimeHandle);
237    }
238    Ok(RuntimeExtensionKind::Runtime)
239}
ax_runtime/thread/lifecycle/extension.rs

ax_runtime/thread/lifecycle/
extension.rs
