Skip to main content

cpu_local/
thread.rs

1use core::{
2    mem::{offset_of, size_of},
3    pin::Pin,
4    ptr::NonNull,
5    sync::atomic::{AtomicUsize, Ordering},
6};
7
8use crate::{ContextSwitchError, CpuAreaRef, preempt::PreemptionState};
9
10#[derive(Clone, Copy, Debug, Eq, PartialEq)]
11pub(crate) struct CpuBindingEpoch(usize);
12
13#[derive(Clone, Copy, Debug)]
14pub(crate) struct CurrentCpuBinding {
15    pub(crate) area: CpuAreaRef,
16}
17
18const CPU_PHASE_MASK: usize = 0b11;
19const CPU_UNBOUND: usize = 0b00;
20const CPU_BINDING: usize = 0b01;
21const CPU_BOUND: usize = 0b10;
22const CPU_UNBINDING: usize = 0b11;
23
24#[repr(u8)]
25#[derive(Clone, Copy, Debug, Eq, PartialEq)]
26enum ExecutionContextKind {
27    Owned,
28    PermanentBoot,
29}
30
31const fn execution_context_reserved_size() -> usize {
32    64 - 4 * size_of::<usize>() - size_of::<PreemptionState>() - size_of::<ExecutionContextKind>()
33}
34
35/// Pinned architecture header for one execution context.
36///
37/// CPU binding uses a four-phase publication word:
38/// `Unbound -> Binding -> Bound -> Unbinding -> next Unbound`. The epoch is
39/// retained solely to reject a stale incoming switch tail.
40#[repr(C, align(64))]
41pub struct ExecutionContextHeader {
42    cpu_area: AtomicUsize,
43    binding_epoch: AtomicUsize,
44    architecture_state: [AtomicUsize; 2],
45    preemption_state: PreemptionState,
46    kind: ExecutionContextKind,
47    reserved: [u8; execution_context_reserved_size()],
48}
49
50impl ExecutionContextHeader {
51    /// Creates an unbound header before placing it in stable pinned storage.
52    pub const fn new() -> Self {
53        Self {
54            cpu_area: AtomicUsize::new(0),
55            binding_epoch: AtomicUsize::new(CPU_UNBOUND),
56            architecture_state: [const { AtomicUsize::new(0) }; 2],
57            preemption_state: PreemptionState::new(),
58            kind: ExecutionContextKind::Owned,
59            reserved: [0; execution_context_reserved_size()],
60        }
61    }
62
63    pub(crate) const fn boot(area_base: usize) -> Self {
64        Self {
65            cpu_area: AtomicUsize::new(area_base),
66            binding_epoch: AtomicUsize::new(CPU_BOUND),
67            architecture_state: [const { AtomicUsize::new(0) }; 2],
68            preemption_state: PreemptionState::bootstrap_disabled(),
69            kind: ExecutionContextKind::PermanentBoot,
70            reserved: [0; execution_context_reserved_size()],
71        }
72    }
73
74    /// Creates the execution owner's continuation of the boot context.
75    ///
76    /// Its initial depth retains bootstrap exclusion until the owner has
77    /// published all state required by its external safe point.
78    #[doc(hidden)]
79    pub const fn new_bootstrap() -> Self {
80        Self {
81            cpu_area: AtomicUsize::new(0),
82            binding_epoch: AtomicUsize::new(CPU_UNBOUND),
83            architecture_state: [const { AtomicUsize::new(0) }; 2],
84            preemption_state: PreemptionState::bootstrap_disabled(),
85            kind: ExecutionContextKind::Owned,
86            reserved: [0; execution_context_reserved_size()],
87        }
88    }
89
90    /// Reports whether this is a CPU area's permanent pre-runtime placeholder.
91    ///
92    /// The kind is immutable after construction, so a current-context reader
93    /// may classify its already-live header without sampling CPU-local state.
94    #[doc(hidden)]
95    #[inline(always)]
96    pub const fn is_permanent_boot_context(&self) -> bool {
97        matches!(self.kind, ExecutionContextKind::PermanentBoot)
98    }
99
100    /// Returns the stable CPU area while this header is fully bound.
101    pub fn cpu_area(&self) -> Option<CpuAreaRef> {
102        self.cpu_binding().map(|binding| binding.area)
103    }
104
105    /// Returns the raw bound area base used by architecture trap entry.
106    pub fn cpu_area_base(&self) -> Option<usize> {
107        self.cpu_binding().map(|binding| binding.area.base())
108    }
109
110    pub(crate) unsafe fn bind_cpu(
111        self: Pin<&Self>,
112        area: CpuAreaRef,
113    ) -> Result<CpuBindingEpoch, ContextSwitchError> {
114        let this = self.get_ref();
115        let unbound = this.binding_epoch.load(Ordering::Acquire);
116        if unbound & CPU_PHASE_MASK != CPU_UNBOUND {
117            return Err(ContextSwitchError::NextContextAlreadyBound);
118        }
119        this.binding_epoch
120            .compare_exchange(
121                unbound,
122                unbound | CPU_BINDING,
123                Ordering::AcqRel,
124                Ordering::Acquire,
125            )
126            .map_err(|_| ContextSwitchError::NextContextAlreadyBound)?;
127        this.cpu_area.store(area.base(), Ordering::Relaxed);
128        let bound = (unbound & !CPU_PHASE_MASK) | CPU_BOUND;
129        this.binding_epoch.store(bound, Ordering::Release);
130        Ok(CpuBindingEpoch(bound))
131    }
132
133    pub(crate) unsafe fn unbind_cpu(
134        self: Pin<&Self>,
135        expected: CpuBindingEpoch,
136    ) -> Result<(), ContextSwitchError> {
137        if expected.0 & CPU_PHASE_MASK != CPU_BOUND {
138            return Err(ContextSwitchError::StalePreviousBinding);
139        }
140        let this = self.get_ref();
141        let unbinding = (expected.0 & !CPU_PHASE_MASK) | CPU_UNBINDING;
142        this.binding_epoch
143            .compare_exchange(expected.0, unbinding, Ordering::AcqRel, Ordering::Acquire)
144            .map_err(|_| ContextSwitchError::StalePreviousBinding)?;
145        this.cpu_area.store(0, Ordering::Relaxed);
146        let next_unbound = (expected.0 & !CPU_PHASE_MASK).wrapping_add(4);
147        this.binding_epoch.store(next_unbound, Ordering::Release);
148        Ok(())
149    }
150
151    pub(crate) fn cpu_binding(&self) -> Option<CurrentCpuBinding> {
152        let (area_base, _) = self.raw_cpu_binding()?;
153        // SAFETY: only bind_cpu can publish this field, and it accepts an
154        // already validated shutdown-lifetime CpuAreaRef.
155        let area = unsafe { CpuAreaRef::from_initialized_base(area_base) }.ok()?;
156        Some(CurrentCpuBinding { area })
157    }
158
159    pub(crate) fn is_bound_to(&self, area: CpuAreaRef) -> bool {
160        self.binding_epoch_for_area(area).is_some()
161    }
162
163    pub(crate) fn binding_epoch_for_area(&self, area: CpuAreaRef) -> Option<CpuBindingEpoch> {
164        // The caller already owns a validated CpuAreaRef. Preserve the stable
165        // epoch observation without reconstructing the same area identity.
166        self.raw_cpu_binding()
167            .and_then(|(area_base, epoch)| (area_base == area.base()).then_some(epoch))
168    }
169
170    pub(crate) fn raw_cpu_binding(&self) -> Option<(usize, CpuBindingEpoch)> {
171        #[cfg(feature = "host-test")]
172        crate::register::host_test::record_binding_observation();
173        loop {
174            let before = self.binding_epoch.load(Ordering::Acquire);
175            if before & CPU_PHASE_MASK != CPU_BOUND {
176                return None;
177            }
178            let area_base = self.cpu_area.load(Ordering::Relaxed);
179            let after = self.binding_epoch.load(Ordering::Acquire);
180            if before == after {
181                return Some((area_base, CpuBindingEpoch(after)));
182            }
183            core::hint::spin_loop();
184        }
185    }
186
187    /// Returns the stable pointer installed in the current-context source.
188    pub fn as_non_null(self: Pin<&Self>) -> NonNull<Self> {
189        NonNull::from(self.get_ref())
190    }
191
192    pub(crate) const fn preemption_state(&self) -> &PreemptionState {
193        &self.preemption_state
194    }
195}
196
197impl Default for ExecutionContextHeader {
198    fn default() -> Self {
199        Self::new()
200    }
201}
202
203/// Byte offset of the execution context's bound CPU-area base.
204pub const EXECUTION_CONTEXT_CPU_BASE_OFFSET: usize = offset_of!(ExecutionContextHeader, cpu_area);
205/// Byte offset of architecture-owned execution-context trap state.
206pub const EXECUTION_CONTEXT_ARCH_STATE_OFFSET: usize =
207    offset_of!(ExecutionContextHeader, architecture_state);
208/// Reserved bytes available to architecture-owned execution-context trap state.
209pub const EXECUTION_CONTEXT_ARCH_STATE_SIZE: usize = 2 * size_of::<usize>();
210
211const _: () = {
212    assert!(EXECUTION_CONTEXT_CPU_BASE_OFFSET == 0);
213    assert!(size_of::<ExecutionContextHeader>() == 64);
214    assert!(core::mem::align_of::<ExecutionContextHeader>() == 64);
215};
216
217#[cfg(test)]
218mod tests {
219    use core::mem::MaybeUninit;
220
221    use super::*;
222    use crate::{CpuAreaPrefix, CpuIndex};
223
224    fn modeled_area(cpu_index: usize) -> CpuAreaRef {
225        let storage = Box::leak(Box::new(MaybeUninit::<CpuAreaPrefix>::uninit()));
226        let base = storage.as_mut_ptr() as usize;
227        storage.write(
228            CpuAreaPrefix::initialize(CpuIndex::try_from(cpu_index).unwrap(), base).unwrap(),
229        );
230        // SAFETY: the initialized fixture is leaked for the process lifetime.
231        unsafe { CpuAreaRef::from_initialized_base(base) }.unwrap()
232    }
233
234    #[test]
235    fn execution_context_header_starts_with_cpu_binding() {
236        assert_eq!(EXECUTION_CONTEXT_CPU_BASE_OFFSET, 0);
237    }
238
239    #[test]
240    fn stable_binding_matches_only_the_published_area() {
241        let first = modeled_area(0);
242        let second = modeled_area(1);
243        let header = Box::pin(ExecutionContextHeader::new());
244
245        // SAFETY: the pinned fixture is unbound and this test owns it.
246        let epoch = unsafe { header.as_ref().bind_cpu(first) }.unwrap();
247        assert!(header.is_bound_to(first));
248        assert!(!header.is_bound_to(second));
249
250        // SAFETY: the test owns the same live binding epoch.
251        unsafe { header.as_ref().unbind_cpu(epoch) }.unwrap();
252        assert!(!header.is_bound_to(first));
253    }
254}