cpu_local/register/
mod.rs

1//! Architecture register primitives and shared validation.
2
3use core::{pin::Pin, ptr::NonNull, sync::atomic::Ordering};
4
5use crate::{CpuAreaRef, CpuLocalError, CpuPin, CurrentThreadHeader, ThreadSwitchError};
6
7#[cfg(all(not(feature = "host-test"), target_arch = "aarch64"))]
8mod aarch64;
9#[cfg(feature = "host-test")]
10mod host;
11#[cfg(all(not(feature = "host-test"), target_arch = "loongarch64"))]
12mod loongarch64;
13#[cfg(all(
14    not(feature = "host-test"),
15    any(target_arch = "riscv32", target_arch = "riscv64")
16))]
17mod riscv;
18#[cfg(all(not(feature = "host-test"), target_arch = "x86_64"))]
19mod x86_64;
20
21#[cfg(all(not(feature = "host-test"), target_arch = "aarch64"))]
22use aarch64 as imp;
23#[cfg(feature = "host-test")]
24use host as imp;
25#[cfg(all(not(feature = "host-test"), target_arch = "loongarch64"))]
26use loongarch64 as imp;
27#[cfg(all(
28    not(feature = "host-test"),
29    any(target_arch = "riscv32", target_arch = "riscv64")
30))]
31use riscv as imp;
32#[cfg(all(not(feature = "host-test"), target_arch = "x86_64"))]
33use x86_64 as imp;
34
35#[cfg(all(
36    not(feature = "host-test"),
37    not(any(
38        target_arch = "x86_64",
39        target_arch = "aarch64",
40        target_arch = "riscv32",
41        target_arch = "riscv64",
42        target_arch = "loongarch64"
43    ))
44))]
45compile_error!("cpu-local supports x86_64, AArch64, RISC-V, and LoongArch64 only");
46
47/// Installs the final area of an offline CPU.
48///
49/// # Safety
50///
51/// The area must remain mapped until shutdown. The CPU must be offline with
52/// traps disabled, and no previous area may be installed on this physical CPU.
53#[doc(hidden)]
54pub unsafe fn install_cpu_area(area: CpuAreaRef) -> Result<(), CpuLocalError> {
55    imp::validate_environment()?;
56    let boot_thread = area.prefix().boot_thread().header();
57    let boot_pointer = boot_thread as *const CurrentThreadHeader as usize;
58    // SAFETY: the caller owns the offline register installation boundary.
59    unsafe { imp::install_cpu_base(area.base(), boot_pointer) };
60    if unsafe { imp::read_cpu_base()? } != area.base() {
61        fatal_register_invariant();
62    }
63    Ok(())
64}
65
66pub(crate) fn current_area() -> Result<CpuAreaRef, CpuLocalError> {
67    let area_base = unsafe { imp::read_cpu_base()? };
68    if area_base == 0 {
69        return Err(CpuLocalError::AreaNotInstalled);
70    }
71    // SAFETY: only install_cpu_area writes the architecture-owned base, and
72    // its contract requires a shutdown-lifetime initialized area.
73    unsafe { CpuAreaRef::from_initialized_base(area_base) }
74}
75
76/// Publishes the scheduler anchor before the architecture switch tail.
77///
78/// # Safety
79///
80/// The caller must own the final IRQ-disabled context-switch boundary. `value`
81/// must identify the prepared pinned header and remain alive while current.
82pub(crate) unsafe fn commit_current_thread(area: CpuAreaRef, value: usize) {
83    area.runtime_anchor()
84        .current_thread_slot()
85        .store(value, Ordering::Release);
86}
87
88/// Returns the pinned current-thread header after checking both sources.
89pub fn current_thread(pin: &CpuPin<'_>) -> Result<NonNull<CurrentThreadHeader>, CpuLocalError> {
90    let area = pin.area();
91    let slot = area.runtime_anchor().current_thread_raw();
92    let register = unsafe { imp::read_current_thread(area.base()) };
93    if slot == 0
94        || slot != register
95        || !slot.is_multiple_of(core::mem::align_of::<CurrentThreadHeader>())
96    {
97        return Err(CpuLocalError::CurrentThreadMismatch);
98    }
99    let pointer = NonNull::new(slot as *mut CurrentThreadHeader)
100        .ok_or(CpuLocalError::CurrentThreadMismatch)?;
101    // SAFETY: scheduler publication only accepts pinned headers that remain
102    // alive while current, and the caller holds the required CPU pin.
103    let thread_area = unsafe { pointer.as_ref() }
104        .cpu_area()
105        .ok_or(CpuLocalError::CurrentThreadMismatch)?;
106    if thread_area != area {
107        return Err(CpuLocalError::CurrentThreadMismatch);
108    }
109    Ok(pointer)
110}
111
112/// Reads the current header before the scheduler can construct its guard.
113///
114/// # Safety
115///
116/// The caller must keep the scheduler-owned current task alive and must not
117/// dereference the result after a context switch.
118#[doc(hidden)]
119pub unsafe fn scheduler_current_thread() -> Result<NonNull<CurrentThreadHeader>, CpuLocalError> {
120    let area = current_area()?;
121    let slot = area.runtime_anchor().current_thread_raw();
122    let register = unsafe { imp::read_current_thread(area.base()) };
123    if slot == 0 || slot != register {
124        return Err(CpuLocalError::CurrentThreadMismatch);
125    }
126    NonNull::new(slot as *mut CurrentThreadHeader).ok_or(CpuLocalError::CurrentThreadMismatch)
127}
128
129/// Binds and publishes the first scheduler task on an offline CPU.
130///
131/// # Safety
132///
133/// The CPU must remain offline and trap-free. `header` must stay pinned and
134/// alive until the scheduler replaces it through a prepared switch.
135#[doc(hidden)]
136pub unsafe fn install_bootstrap_thread(
137    pin: &CpuPin<'_>,
138    header: Pin<&CurrentThreadHeader>,
139) -> Result<(), ThreadSwitchError> {
140    let epoch = unsafe { header.bind_cpu(pin.area()) }?;
141    let pointer = header.as_non_null().as_ptr() as usize;
142    unsafe { commit_current_thread(pin.area(), pointer) };
143    // Bootstrap has no raw switch tail. Install the architecture-owned current
144    // register directly while this CPU remains offline and trap-free.
145    unsafe { imp::write_current_thread(pointer) };
146    if current_thread(pin) != Ok(header.as_non_null()) {
147        // The register is already committed, so continuing would make all
148        // later Rust execution unsound. Rollback is intentionally impossible.
149        let _ = epoch;
150        fatal_register_invariant();
151    }
152    Ok(())
153}
154
155/// Reads task-owned kernel TLS under an explicit CPU pin.
156#[cfg(feature = "tls")]
157pub fn kernel_tls(_pin: &CpuPin<'_>) -> usize {
158    unsafe { imp::read_kernel_tls() }
159}
160
161/// Installs task-owned kernel TLS at an offline bootstrap boundary.
162///
163/// # Safety
164///
165/// The caller must own the offline CPU or IRQ-disabled final task switch, and
166/// `value` must remain a valid TLS base for the installed execution context.
167#[cfg(feature = "tls")]
168#[doc(hidden)]
169pub unsafe fn install_kernel_tls(_pin: &CpuPin<'_>, value: usize) {
170    unsafe { imp::write_kernel_tls(value) };
171}
172
173#[cold]
174#[inline(never)]
175fn fatal_register_invariant() -> ! {
176    panic!("CPU-local register commit did not retain the validated state")
177}
cpu_local/register/mod.rs

cpu_local/register/
mod.rs
