ax_task/runtime/
cpu.rs

1//! Pinned CPU capabilities and scheduler observations.
2
3pub use crate::{
4    runtime::clock::{RqClockSample, SchedulerDeadlineUpdate, SchedulerRuntimeDeadline},
5    sched::system::{
6        CpuLifecycleState, CpuLoadSummary, CpuLocal, CpuLocalOwnerBorrow, CpuRemote, CpuSnapshot,
7    },
8};
9use crate::{
10    runtime::{
11        context::{current_cpu_remote, runtime_current_cpu, validate_schedule_context},
12        lock::PreemptScope,
13        switch::RuntimeScheduleOrigin,
14        task_runtime,
15    },
16    thread::TaskError,
17};
18
19/// Tests the current CPU's sticky reschedule request while migration is pinned.
20///
21/// # Safety
22///
23/// The caller must prevent migration until it has finished the decision that
24/// uses this snapshot. Sleeping-lock owner spinning normally satisfies this
25/// with a preemption guard.
26pub unsafe fn current_needs_reschedule_pinned() -> Result<bool, TaskError> {
27    Ok(current_cpu_remote()
28        .ok_or(TaskError::NotInitialized)?
29        .needs_reschedule())
30}
31
32/// Tests only scheduler work consumed by kernel preempt-enable/IRQ return.
33///
34/// # Safety
35///
36/// The caller must prevent migration until it has finished the decision that
37/// uses this snapshot.
38pub unsafe fn current_needs_immediate_scheduler_work_pinned() -> Result<bool, TaskError> {
39    Ok(current_cpu_remote()
40        .ok_or(TaskError::NotInitialized)?
41        .needs_immediate_scheduler_work())
42}
43
44/// Tests the sticky reschedule state of the calling CPU.
45pub fn current_cpu_needs_resched() -> Result<bool, TaskError> {
46    let _pin = PreemptScope::enter();
47    // SAFETY: `_pin` prevents migration through the remote reschedule-state
48    // observation. Stronger IRQ/scheduler owner scopes are inherited.
49    unsafe { current_needs_reschedule_pinned() }
50}
51
52/// Observes only the immediate preemption bit in real-runtime regression tests.
53/// Owner maintenance and lazy preemption remain separate scheduler requests.
54#[cfg(feature = "fault-injection")]
55pub fn current_immediate_preemption_requested() -> Result<bool, TaskError> {
56    let _pin = PreemptScope::enter();
57    Ok(current_cpu_remote()
58        .ok_or(TaskError::NotInitialized)?
59        .immediate_preemption_requested())
60}
61
62/// Clears the current CPU's idle-polling state at the runtime sleep boundary.
63///
64/// # Safety
65///
66/// The runtime must have disabled local interrupts and must prevent migration
67/// through the immediately following sticky-work and clockevent recheck. This
68/// is Linux's `current_clr_polling_and_test()` boundary: work published before
69/// the clear is found by that recheck, while work published afterwards must
70/// own a physical interrupt edge.
71#[doc(hidden)]
72pub unsafe fn finish_current_cpu_idle_polling() -> Result<(), TaskError> {
73    let remote = current_cpu_remote().ok_or(TaskError::NotInitialized)?;
74    remote.finish_idle_wait();
75    Ok(())
76}
77
78/// Executes one lossless idle publication/recheck/WFI iteration.
79pub fn idle_current_cpu_once() -> Result<(), TaskError> {
80    validate_schedule_context(RuntimeScheduleOrigin::Preempt)?;
81    let may_wait = {
82        let cpu = runtime_current_cpu()?;
83        cpu.prepare_idle_wait()
84    };
85    if may_wait {
86        task_runtime::wait_for_interrupt();
87    }
88    Ok(())
89}
90use crate::runtime::handle::opaque_handle;
91
92opaque_handle!(
93    /// Opaque address of the current CPU's pinned owner-only scheduler object.
94    ///
95    /// Consumers must claim the corresponding [`crate::runtime::cpu::CpuRemote`] owner gate
96    /// before reconstructing any reference from this address.
97    CurrentCpuLocalHandle,
98    "runtime::cpu"
99);
100opaque_handle!(
101    /// Opaque pointer-sized handle to one Arc-backed remote CPU endpoint.
102    ///
103    /// Remote and owner-only CPU handles are intentionally not interchangeable:
104    ///
105    /// ```compile_fail
106    /// use ax_task::runtime::cpu::{CpuRemoteHandle, CurrentCpuLocalHandle};
107    ///
108    /// fn borrow_owner(_handle: CurrentCpuLocalHandle) {}
109    /// borrow_owner(CpuRemoteHandle::NONE);
110    /// ```
111    CpuRemoteHandle,
112    "runtime::cpu"
113);
114opaque_handle!(
115    /// Token returned by the nested IRQ guard service.
116    IrqGuardToken,
117    "runtime::cpu"
118);
119opaque_handle!(
120    /// Token returned by the nested task-preemption guard service.
121    PreemptGuardToken,
122    "runtime::cpu"
123);
124
125/// Runtime-defined raw local-IRQ state saved by a synchronization guard.
126///
127/// Unlike [`IrqGuardToken`], this value does not own a scheduler publication
128/// scope. It only transports the architecture interrupt state back to the
129/// runtime that produced it.
130#[derive(Clone, Copy, Debug, Eq, PartialEq)]
131#[repr(transparent)]
132pub struct LocalIrqState(usize);
133
134impl LocalIrqState {
135    /// Creates a saved local-IRQ state at the runtime provider boundary.
136    ///
137    /// # Safety
138    ///
139    /// `raw` must be a state value accepted by the linked runtime's matching
140    /// local-IRQ restore operation.
141    pub const unsafe fn from_raw(raw: usize) -> Self {
142        Self(raw)
143    }
144
145    /// Returns the runtime-owned representation of this saved state.
146    pub const fn into_raw(self) -> usize {
147        self.0
148    }
149}
150
151/// Logical CPU identifier exchanged with the operating-system runtime.
152#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
153#[repr(transparent)]
154pub struct RuntimeCpuId(u32);
155
156impl RuntimeCpuId {
157    /// Creates a logical CPU identifier.
158    pub const fn new(value: u32) -> Self {
159        Self(value)
160    }
161
162    /// Returns the numeric logical CPU identifier.
163    pub const fn as_u32(self) -> u32 {
164        self.0
165    }
166}
167
168/// Runtime-owned capability snapshot for one pinned scheduler CPU.
169///
170/// The paired fields are captured in one runtime operation, mirroring Linux's
171/// direct `this_rq()` lookup. The remote endpoint is the sole owner identity;
172/// its embedded CPU ID prevents a second architecture or registry lookup.
173#[derive(Clone, Copy, Debug, Eq, PartialEq)]
174#[repr(C)]
175pub struct CurrentCpuOwnerHandles {
176    local: CurrentCpuLocalHandle,
177    remote: CpuRemoteHandle,
178}
179
180impl CurrentCpuOwnerHandles {
181    /// Empty capability used when a scheduler-frame entry is rejected.
182    pub const NONE: Self = Self {
183        local: CurrentCpuLocalHandle::NONE,
184        remote: CpuRemoteHandle::NONE,
185    };
186
187    /// Creates one pinned current-CPU capability snapshot.
188    ///
189    /// # Safety
190    ///
191    /// `local` and `remote` must identify the paired owner-only and Arc-backed
192    /// scheduler endpoints for the pinned CPU. Every non-empty handle must
193    /// remain live until shutdown, and the caller must keep migration excluded
194    /// while the snapshot is used.
195    pub const unsafe fn new(local: CurrentCpuLocalHandle, remote: CpuRemoteHandle) -> Self {
196        Self { local, remote }
197    }
198
199    /// Returns the current CPU's owner-only scheduler handle.
200    pub const fn local(self) -> CurrentCpuLocalHandle {
201        self.local
202    }
203
204    /// Returns the current CPU's Arc-backed remote endpoint.
205    pub const fn remote(self) -> CpuRemoteHandle {
206        self.remote
207    }
208}
209
210pub use crate::sched::system::OwnerControlDrain;
211
212/// Failed prerequisite observed by the serial real-idle test probe.
213#[cfg(feature = "fault-injection")]
214#[derive(Clone, Copy, Debug)]
215#[repr(u8)]
216pub enum IdleOfflineRejection {
217    /// No instrumented prerequisite failed.
218    Unclassified         = 0,
219    /// A publisher still owns the placement endpoint.
220    PlacementPublication = 1,
221    /// A thread cannot leave this CPU's placement domain.
222    ThreadTarget         = 2,
223    /// Owner-directed delivery has not relinquished publication.
224    OwnerPublication     = 3,
225    /// Runqueue, timer, handoff or remote work remains.
226    CpuState             = 4,
227    /// A thread retains CPU ownership or a migration pin.
228    ThreadOwnership      = 5,
229    /// Scheduler work arrived before owner publication was closed.
230    SchedulerWork        = 6,
231}
232
233#[cfg(feature = "fault-injection")]
234static IDLE_OFFLINE_REJECTION: core::sync::atomic::AtomicU8 = core::sync::atomic::AtomicU8::new(0);
235
236#[cfg(feature = "fault-injection")]
237pub(crate) fn record_idle_offline_rejection(reason: IdleOfflineRejection) {
238    IDLE_OFFLINE_REJECTION.store(reason as u8, core::sync::atomic::Ordering::Release);
239}
240
241/// Reads the last serial probe's rejection after its locks have been released.
242#[cfg(feature = "fault-injection")]
243pub fn idle_offline_rejection() -> IdleOfflineRejection {
244    match IDLE_OFFLINE_REJECTION.load(core::sync::atomic::Ordering::Acquire) {
245        1 => IdleOfflineRejection::PlacementPublication,
246        2 => IdleOfflineRejection::ThreadTarget,
247        3 => IdleOfflineRejection::OwnerPublication,
248        4 => IdleOfflineRejection::CpuState,
249        5 => IdleOfflineRejection::ThreadOwnership,
250        6 => IdleOfflineRejection::SchedulerWork,
251        _ => IdleOfflineRejection::Unclassified,
252    }
253}
254
255/// Exercises scheduler CPU offline/online on the real idle owner.
256///
257/// This test-only transaction retains IRQ exclusion and the exclusive owner
258/// borrow across both transitions. It never returns to scheduling while offline
259/// and does not implement platform power-off or an externally parked CPU.
260/// Returns `NotReady` while ordinary scheduler work must be drained first.
261/// `publish_work_after_drain` injects scheduler work, then a timer notification
262/// after placement closes, to verify that the per-CPU worker can still drain.
263#[cfg(feature = "fault-injection")]
264pub fn probe_idle_cpu_round_trip(publish_work_after_drain: bool) -> Result<(), TaskError> {
265    use crate::runtime::context::{RuntimeIrqGuard, runtime_current_cpu_mut, runtime_task_system};
266    validate_schedule_context(RuntimeScheduleOrigin::Preempt)?;
267    let system = runtime_task_system()?;
268    let mut irq = RuntimeIrqGuard::enter();
269    let mut cpu = runtime_current_cpu_mut(&mut irq)?;
270    if cpu.remote().current_thread() != cpu.remote().idle_thread() {
271        return Err(TaskError::NotReady);
272    }
273    if publish_work_after_drain {
274        // Force work published after idle's normal scheduler drain.
275        // The lifecycle owner must close placement before draining this work.
276        cpu.request_scheduler_work();
277    }
278    record_idle_offline_rejection(IdleOfflineRejection::Unclassified);
279    let offline = system.take_cpu_offline(cpu.as_mut());
280    if publish_work_after_drain {
281        assert!(
282            matches!(offline, Err(TaskError::NotReady)),
283            "owner work must defer CPU offline: {offline:?}"
284        );
285        assert_eq!(cpu.remote().lifecycle_state(), CpuLifecycleState::Inactive);
286        // A timer IRQ may publish soft work after placement closes. The fixed
287        // worker must still wake, drain the event, and park before final offline.
288        cpu.remote().publish_ktimer_work();
289    }
290    offline?;
291    assert_eq!(cpu.remote().lifecycle_state(), CpuLifecycleState::Offline);
292    assert!(system.cpu_remote(cpu.owner()).is_none());
293    // Returning an error here would strand the executing idle owner offline.
294    system
295        .bring_cpu_online(cpu.as_mut())
296        .expect("idle CPU re-online failed");
297    assert_eq!(cpu.remote().lifecycle_state(), CpuLifecycleState::Online);
298    Ok(())
299}
300
301/// Publishes ordinary owner work so an idle probe leaves NOHZ sleep.
302#[cfg(feature = "fault-injection")]
303pub fn notify_idle_cpu_probe(cpu: RuntimeCpuId) -> Result<(), TaskError> {
304    let system = crate::runtime::context::runtime_task_system()?;
305    let remote = system
306        .cpu_remote(crate::sched::CpuId::new(cpu.as_u32()))
307        .ok_or(TaskError::CpuOffline(cpu.as_u32()))?;
308    if remote.kick_scheduler_work() {
309        Ok(())
310    } else {
311        Err(TaskError::CpuOffline(cpu.as_u32()))
312    }
313}
314
315/// Actual scheduler readers retained by the cross-CPU offline regression.
316#[cfg(feature = "fault-injection")]
317#[derive(Clone, Copy, Debug)]
318pub enum IdleOfflineReader {
319    /// A remote control publisher between admission and completion.
320    OwnerDelivery,
321    /// A source CPU still finishing a committed idle-balance claim.
322    IdleBalance,
323}
324
325/// Retains a real scheduler reader across a controlled cross-CPU test.
326/// No IRQ or rq guard is held across the callback, so the target can run its
327/// ordinary idle lifecycle protocol while the callback observes the transition.
328#[cfg(feature = "fault-injection")]
329pub fn with_idle_offline_reader<T>(
330    cpu: RuntimeCpuId,
331    reader: IdleOfflineReader,
332    action: impl FnOnce(&CpuRemote) -> T,
333) -> Result<T, TaskError> {
334    crate::thread::current::validate_blocking_context()?;
335    let system = crate::runtime::context::runtime_task_system()?;
336    let remote = system
337        .cpu_remote(crate::sched::CpuId::new(cpu.as_u32()))
338        .ok_or(TaskError::CpuOffline(cpu.as_u32()))?;
339    record_idle_offline_rejection(IdleOfflineRejection::Unclassified);
340    match reader {
341        IdleOfflineReader::OwnerDelivery => {
342            let _publication = remote
343                .begin_owner_delivery()
344                .ok_or(TaskError::CpuOffline(cpu.as_u32()))?;
345            Ok(action(remote))
346        }
347        IdleOfflineReader::IdleBalance => {
348            let crate::sched::system::IdlePullReservation::Started(reservation) =
349                remote.begin_idle_pull()
350            else {
351                return Err(TaskError::NotReady);
352            };
353            let Some(mut claim) = remote.claim_idle_pull(reservation) else {
354                remote.cancel_idle_pull(reservation);
355                return Err(TaskError::NotReady);
356            };
357            if !claim.commit() {
358                return Err(TaskError::NotReady);
359            }
360            Ok(action(remote))
361        }
362    }
363}
ax_task/runtime/cpu.rs

ax_task/runtime/
cpu.rs
