1use alloc::{boxed::Box, string::String};
4use core::{pin::Pin, ptr};
5
6use ax_hal::percpu::CpuPin;
7use ax_lazyinit::LazyInit;
8use ax_task::{
9 runtime::{
10 RuntimeHandleResult, RuntimeStatus, TaskSystem, TaskSystemHandle,
11 config::TaskSystemConfig,
12 cpu::{
13 CpuLocal, CpuRemote, CpuRemoteHandle, CurrentCpuLocalHandle, CurrentCpuOwnerHandles,
14 IrqGuardToken, RuntimeCpuId,
15 },
16 resource::{
17 AddressSpaceDestroyOutcome, AddressSpaceHandle, AddressSpaceMembarrierState,
18 AddressSpaceReclaimArmOutcome, ExecutionContextHandle, KernelContextRequest,
19 MembarrierRegistration, MembarrierRegistrationPhase, RuntimeMembarrierAction,
20 StackHandle, StackRequest, ThreadResources, TlsHandle, UserContextRequest,
21 },
22 switch::{
23 ContextThreadBinding, CurrentThreadPublication, RuntimeSchedulerFrameEnterResult,
24 RuntimeSwitchPlan, SchedSwitchRecord, ThreadIdentityV1,
25 },
26 task_runtime::impl_trait as impl_task_runtime,
27 },
28 sched::{CpuId, CpuSet, FairMode, Nice, SchedulePolicy},
29 thread::{TaskError, ThreadId, ThreadSpec, current::current_thread_id},
30};
31
32mod address_space;
33mod mm_activation;
34pub use mm_activation::{
35 AddressSpaceSwitchProof, SchedulerAddressSpaceActivation, SchedulerAddressSpaceOwner,
36 UserAddressSpaceOwner,
37};
38mod bootstrap;
39pub(crate) mod context;
40
41#[cfg(feature = "fault-injection")]
42pub mod creation_probe;
43mod resources;
44pub(crate) mod runtime_impl;
45pub(crate) mod scheduler_events;
46mod spawn;
47mod thread_resources;
48#[cfg(feature = "uspace")]
49mod user_entry;
50
51pub use address_space::{
52 AddressSpaceCpuState, TaskAddressSpace, detach_current_address_space,
53 switch_current_address_space,
54};
55use address_space::{
56 arm_runtime_address_space_reclaim, destroy_runtime_address_space,
57 release_current_active_address_space, runtime_address_space_membarrier_state,
58 update_runtime_address_space_membarrier_state,
59};
60#[cfg(feature = "uspace")]
61use bootstrap::current_cpu_remote;
62#[cfg(kernel_tls)]
63pub(crate) use bootstrap::initialize_early_bootstrap_tls;
64#[cfg(test)]
65use bootstrap::{IdleEntryAction, idle_entry_action};
66pub(crate) use bootstrap::{
67 PublishedCpuOnline, initialize_primary, publish_current_cpu_online,
68 start_current_ktimer_service, start_deferred_task_work_service,
69};
70use bootstrap::{
71 cpu_remote, current_cpu_owner_handles, idle_context_entry, primary_bootstrap_thread,
72 scheduler_current_cpu_remote_handle, task_system, with_current_cpu_pin,
73};
74#[cfg(feature = "smp")]
75pub(crate) use bootstrap::{initialize_secondary, run_idle};
76use context::{
77 bind_bootstrap_runtime_context, bind_runtime_context_thread, create_bootstrap_context,
78 create_runtime_context, create_user_runtime_context, destroy_runtime_context,
79 finish_runtime_context_switch_tail, scheduler_current_thread_identity,
80 scheduler_current_thread_publication, switch_runtime_context,
81};
82
83pub(crate) fn runtime_task_system_handle() -> TaskSystemHandle {
84 task_system().map_or(TaskSystemHandle::NONE, |system| {
85 unsafe { TaskSystemHandle::from_raw((system as *const TaskSystem).expose_provenance()) }
88 })
89}
90
91pub(crate) fn scheduler_frame_capabilities(cpu_pin: &CpuPin) -> RuntimeSchedulerFrameEnterResult {
92 unsafe {
97 RuntimeSchedulerFrameEnterResult::success(
98 runtime_task_system_handle(),
99 current_cpu_owner_handles(cpu_pin),
100 )
101 }
102}
103
104#[cfg(feature = "uspace")]
105pub(crate) fn current_cpu_needs_reschedule_pinned(cpu_pin: &CpuPin) -> Result<bool, TaskError> {
106 Ok(current_cpu_remote(cpu_pin)
107 .ok_or(TaskError::NotInitialized)?
108 .needs_reschedule())
109}
110
111#[cfg(kernel_tls)]
112use resources::runtime_tls_pointer;
113use resources::{
114 allocate_runtime_stack, allocate_runtime_tls, deallocate_runtime_stack, deallocate_runtime_tls,
115};
116pub(crate) use scheduler_events::{on_clock_event, publish_scheduler_tick};
117#[cfg(feature = "qperf-metrics")]
118pub(crate) use scheduler_events::{
119 record_irq_return_scheduler_continuation, record_irq_return_scheduler_window,
120};
121
122#[cfg(axtest)]
124pub fn kernel_thread_retains_active_mm_membarrier_state_for_test() -> bool {
125 static IDENTITY_ANCHOR: u8 = 0;
126
127 let identity_raw = (&IDENTITY_ANCHOR as *const u8).expose_provenance();
128 let identity =
131 unsafe { ax_task::runtime::resource::AddressSpaceMembarrierId::from_raw(identity_raw) };
132 let active_mm_state =
135 unsafe { ax_task::runtime::resource::AddressSpaceMembarrierState::new(identity, 0) };
136
137 ax_task::runtime::resource::scheduled_membarrier_state_for_test(
138 active_mm_state,
139 ax_task::runtime::resource::AddressSpaceMembarrierState::NONE,
140 ) == active_mm_state
141}
142pub fn reset_current_user_fp_state() -> Result<(), TaskError> {
144 context::reset_current_user_fp_state()
145}
146
147#[cfg(all(target_arch = "x86_64", feature = "fp-simd", feature = "uspace"))]
149pub fn capture_current_user_fp_state() -> Result<ax_hal::cpu::registers::UserXstate, TaskError> {
150 context::capture_current_user_fp_state()
151}
152
153#[cfg(all(target_arch = "x86_64", feature = "fp-simd", feature = "uspace"))]
155pub fn replace_current_user_fp_state(
156 state: ax_hal::cpu::registers::UserXstate,
157) -> Result<(), TaskError> {
158 context::replace_current_user_fp_state(state)
159}
160#[cfg(all(feature = "qperf-metrics", any(feature = "ipi", feature = "wake-ipi")))]
161pub(crate) use scheduler_events::{record_scheduler_ipi_consume, record_scheduler_ipi_send};
162pub use spawn::{UserContextOptions, builder, exit_current, prepare_user_thread};
163
164fn finish_initial_scheduler_switch() {
165 unsafe { ax_task::runtime::switch::finish_initial_context_switch() }
167 .expect("initial context switch must finish");
168}
169#[cfg(all(test, kernel_tls))]
170use thread_resources::assemble_bootstrap_resources;
171use thread_resources::{create_bootstrap_resources, create_idle_resources, create_user_resources};
172#[cfg(feature = "uspace")]
173pub use user_entry::UserExecutionContext;
174
175const PAGE_SIZE: usize = 4096;
176
177#[cfg(not(feature = "fs"))]
178const DEFAULT_TASK_STACK_SIZE: usize = 256 * 1024;
179
180const fn runtime_status_error(status: RuntimeStatus) -> TaskError {
181 TaskError::RuntimeFailure(status as u32)
182}
183
184const fn runtime_task_stack_size() -> usize {
185 #[cfg(feature = "fs")]
186 {
187 crate::build_info::TASK_STACK_SIZE
188 }
189 #[cfg(not(feature = "fs"))]
190 {
191 DEFAULT_TASK_STACK_SIZE
192 }
193}
194
195pub const fn default_task_stack_size() -> usize {
197 runtime_task_stack_size()
198}
199
200#[cfg(test)]
201mod tests;
202
203mod allocation;