Skip to main content

ax_task/thread/
current.rs

1//! Operations bound to the calling scheduler thread.
2
3use alloc::sync::Arc;
4use core::marker::PhantomData;
5
6pub use crate::{
7    runtime::switch::dispatch::{
8        ExitPermit, commit_current_exit, exit_current_thread, prepare_current_exit,
9        yield_current_cpu,
10    },
11    sync::wait_queue::{sleep, sleep_until},
12    thread::current::park::{
13        CurrentParkDisposition, CurrentParkResume, CurrentParkStart, PreparedCurrentPark,
14        begin_current_park,
15    },
16};
17use crate::{
18    runtime::{
19        context::{
20            RuntimeSchedulerFrameGuard, runtime_current_cpu_mut, runtime_task_system,
21            validate_schedule_context,
22        },
23        switch::{RuntimeScheduleOrigin, RuntimeSchedulerEntry, dispatch::execute_switch_plan},
24        task_runtime,
25    },
26    sched::CpuSet,
27    thread::{
28        CurrentThreadToken, TaskError, ThreadCore, ThreadExtensionLease, ThreadHandle, ThreadId,
29    },
30};
31
32/// Returns a strong handle for the calling scheduler thread.
33///
34/// # Errors
35///
36/// Returns [`TaskError::NotInitialized`] before runtime CPU publication,
37/// [`TaskError::CpuOwnerBorrowed`] for a reentrant owner query, or
38/// [`TaskError::NoRunnableThread`] before a current thread is installed.
39pub fn current_thread_handle() -> Result<ThreadHandle, TaskError> {
40    #[cfg(feature = "qperf-metrics")]
41    crate::diagnostics::counters::record_current_thread_handle_query();
42    let publication = current_thread_publication()?;
43    // SAFETY: the scheduler retains the executing task's owner-side Arc across
44    // preemption and migration until this synchronous operation returns.
45    unsafe { publication.acquire_handle() }
46}
47
48/// Returns the generation-bearing identity of the calling scheduler thread.
49#[inline(always)]
50pub fn current_thread_id() -> Result<ThreadId, TaskError> {
51    let identity = current_thread_identity()?;
52    Ok(ThreadId::from_parts(identity.slot, identity.generation))
53}
54
55/// Captures the scheduler thread executing this task context.
56#[inline(always)]
57pub fn current_thread_token() -> Result<CurrentThreadToken, TaskError> {
58    Ok(CurrentThreadToken::new(current_thread_id()?))
59}
60
61#[inline(always)]
62pub(crate) fn current_thread_identity()
63-> Result<crate::runtime::switch::ThreadIdentityV1, TaskError> {
64    let identity = task_runtime::current_thread_identity();
65    if identity.is_bound() {
66        return Ok(identity);
67    }
68
69    let publication = task_runtime::current_thread_publication();
70    if publication.identity() != identity || !publication.owner().is_none() {
71        return Err(TaskError::InvalidRuntimeHandle);
72    }
73    // Preserve the public distinction between a runtime that has not installed
74    // its task system and an initialized bootstrap context without a current
75    // scheduler thread. Bound task contexts never enter this cold path.
76    let _system = runtime_task_system()?;
77    Err(TaskError::NoRunnableThread)
78}
79
80pub(crate) fn current_thread_publication()
81-> Result<crate::runtime::switch::CurrentThreadPublication, TaskError> {
82    let publication = task_runtime::current_thread_publication();
83    let identity = publication.identity();
84    if !identity.is_bound() {
85        if !publication.owner().is_none() {
86            return Err(TaskError::InvalidRuntimeHandle);
87        }
88        // Preserve the public distinction between a runtime that has not
89        // installed its task system and an initialized bootstrap context that
90        // has not published a scheduler thread. This cold error path does not
91        // add a handle lookup to the bound-current fast path.
92        let _system = runtime_task_system()?;
93        return Err(TaskError::NoRunnableThread);
94    }
95    if publication.owner().is_none() {
96        return Err(TaskError::InvalidRuntimeHandle);
97    }
98    Ok(publication)
99}
100
101pub(crate) fn current_thread_core_arc() -> Result<Arc<ThreadCore>, TaskError> {
102    let publication = current_thread_publication()?;
103    // SAFETY: the runtime publication belongs to this architecture context.
104    // The returned Arc is scheduler-internal and remains in the synchronous
105    // current-thread operation; it does not acquire an external lease.
106    unsafe { publication.acquire_scheduler_core() }
107}
108
109/// Validates that the caller may publish a waiter or block its current thread.
110///
111/// Sleeping synchronization primitives should call this before changing any
112/// waiter, owner, donation, or thread-lifecycle state.
113pub fn validate_blocking_context() -> Result<(), TaskError> {
114    acquire_blocking_permit().map(|_| ())
115}
116
117/// One validated opportunity to publish a blocking handshake.
118pub(crate) struct BlockingPermit {
119    _not_send: PhantomData<*mut ()>,
120}
121
122pub(crate) fn acquire_blocking_permit() -> Result<BlockingPermit, TaskError> {
123    validate_schedule_context(RuntimeScheduleOrigin::Block)?;
124    Ok(BlockingPermit {
125        _not_send: PhantomData,
126    })
127}
128
129/// Returns the opaque extension of the calling scheduler thread.
130///
131/// Runtime entry trampolines use the callback-table address as a type identity
132/// before recovering an OS-owned closure or process object from `data`.
133pub fn current_thread_extension() -> Result<Option<ThreadExtensionLease>, TaskError> {
134    let handle = current_thread_handle()?;
135    Ok(handle
136        .extension_view()
137        .map(|view| ThreadExtensionLease::new(view, handle)))
138}
139
140/// Updates the calling thread's affinity and completes a required migration.
141///
142/// A successful return guarantees that the caller is executing on a CPU in
143/// the new mask. Generic remote-thread affinity updates remain asynchronous and
144/// are completed by the remote owner's next scheduler safe point.
145pub fn set_current_thread_affinity(affinity: CpuSet) -> Result<(), TaskError> {
146    let mut scheduler_frame = RuntimeSchedulerFrameGuard::enter(
147        RuntimeScheduleOrigin::Yield,
148        RuntimeSchedulerEntry::Task,
149    )?;
150    let current = scheduler_frame.current_thread_ref()?;
151    let system = scheduler_frame.task_system();
152    let mut outcome = {
153        let mut cpu = runtime_current_cpu_mut(&mut scheduler_frame)?;
154        let must_migrate = system.set_current_affinity(cpu.as_mut(), affinity)?;
155        if !must_migrate {
156            return Ok(());
157        }
158
159        // The new mask is now visible and excludes this CPU. Keep the scheduler
160        // baton and raw IRQ mask continuously owned until this context has moved;
161        // exposing an IRQ-enabled validation window here could let IRQ-return
162        // scheduling migrate the caller between publishing the mask and yielding.
163        // SAFETY: `scheduler_frame` owns the IRQ-off scheduler baton.
164        unsafe { system.yield_current_in_scheduler_frame(cpu.as_mut()) }.unwrap_or_else(|_| {
165            // Affinity publication cannot be rolled back safely after another CPU
166            // may have observed the migration target. Scheduler commit failures are
167            // therefore runtime invariants, like failures after exit publication.
168            task_runtime::fatal_invariant(0x4558_0021, current.id().as_u64() as usize);
169        })
170    };
171    let decision = outcome.decision_mut().unwrap_or_else(|| {
172        task_runtime::fatal_invariant(0x4558_0022, current.id().as_u64() as usize)
173    });
174    execute_switch_plan(&mut scheduler_frame, decision);
175    Ok(())
176}
177pub(crate) mod park;