Skip to main content

cpu_local/
switch.rs

1use core::{marker::PhantomData, mem::ManuallyDrop, pin::Pin, ptr::NonNull};
2
3use crate::{
4    ContextSwitchError, CpuBindingEpoch, CpuLocalError, CpuPin, ExecutionContextHeader,
5    current_context,
6};
7
8/// Prepared current-context publication owned by the final switch tail.
9#[must_use = "dropping an uncommitted switch rolls back the next CPU binding"]
10pub struct PreparedContextSwitch<'switch> {
11    next: NonNull<ExecutionContextHeader>,
12    next_epoch: CpuBindingEpoch,
13    current_context: usize,
14    area: crate::CpuAreaRef,
15    _scope: PhantomData<&'switch mut &'switch ()>,
16    _not_send_or_sync: PhantomData<*mut ()>,
17}
18
19impl PreparedContextSwitch<'_> {
20    /// Returns the exact next header bound by this transaction.
21    #[doc(hidden)]
22    pub const fn next_header(&self) -> NonNull<ExecutionContextHeader> {
23        self.next
24    }
25
26    /// Publishes the prepared CPU anchor slot immediately before naked switch.
27    ///
28    /// # Safety
29    ///
30    /// The caller serialization and IRQ exclusion used during preparation
31    /// must still be active. The caller must enter the architecture switch
32    /// without performing fallible or ownership-sensitive Rust work.
33    #[doc(hidden)]
34    #[inline(always)]
35    pub unsafe fn commit(self) {
36        // Disarm rollback before publication. After the final call below there
37        // is no destructor state update or ownership-sensitive Rust work; the
38        // architecture wrapper enters its naked switch tail immediately.
39        let prepared = ManuallyDrop::new(self);
40        unsafe { crate::register::commit_current_context(prepared.area, prepared.current_context) };
41    }
42}
43
44impl Drop for PreparedContextSwitch<'_> {
45    fn drop(&mut self) {
46        // SAFETY: an uncommitted token still owns the next binding, and its
47        // invariant lifetime keeps the caller's critical section live. The
48        // preparation contract keeps the pinned header alive until this drop.
49        let next = unsafe { Pin::new_unchecked(self.next.as_ref()) };
50        if unsafe { next.unbind_cpu(self.next_epoch) }.is_err() {
51            panic!("prepared context-switch rollback lost the next CPU binding");
52        }
53    }
54}
55
56/// Opaque previous-context binding consumed by the incoming switch tail.
57#[must_use = "the incoming context must withdraw the previous CPU binding"]
58#[derive(Debug)]
59pub struct PreviousContextBinding {
60    previous: NonNull<ExecutionContextHeader>,
61    epoch: CpuBindingEpoch,
62}
63
64impl PreviousContextBinding {
65    /// Withdraws the exact previous binding after architecture registers have
66    /// switched to the incoming context.
67    ///
68    /// # Errors
69    ///
70    /// Returns [`ContextSwitchError::PreviousContextMismatch`] if `previous`
71    /// differs from the prepared context, or
72    /// [`ContextSwitchError::StalePreviousBinding`] for an obsolete tail.
73    ///
74    /// # Safety
75    ///
76    /// The incoming switch tail must be the sole owner of this token and the
77    /// previous context allocation must remain pinned and alive.
78    pub unsafe fn finish(
79        self,
80        previous: Pin<&ExecutionContextHeader>,
81    ) -> Result<(), ContextSwitchError> {
82        if previous.as_non_null() != self.previous {
83            return Err(ContextSwitchError::PreviousContextMismatch);
84        }
85        unsafe { previous.unbind_cpu(self.epoch) }
86    }
87}
88
89/// Validates and binds a complete execution-context switch transaction.
90///
91/// All fallible validation occurs before the returned prepared token can be
92/// committed. Dropping the prepared token automatically rolls back `next`.
93///
94/// # Safety
95///
96/// The caller must own the IRQ-disabled context-switch path. Both headers
97/// must remain pinned and alive through the raw switch and incoming tail.
98pub unsafe fn prepare_context_switch<'switch>(
99    pin: &'switch CpuPin<'_>,
100    previous: Pin<&ExecutionContextHeader>,
101    next: Pin<&ExecutionContextHeader>,
102) -> Result<(PreparedContextSwitch<'switch>, PreviousContextBinding), ContextSwitchError> {
103    let published = current_context(pin).map_err(|error| match error {
104        CpuLocalError::CurrentContextMismatch => ContextSwitchError::CurrentContextMismatch,
105        other => ContextSwitchError::CpuLocal(other),
106    })?;
107    if published != previous.as_non_null() {
108        return Err(ContextSwitchError::CurrentContextMismatch);
109    }
110    let previous_binding = previous
111        .cpu_binding()
112        .filter(|binding| binding.area == pin.area())
113        .ok_or(ContextSwitchError::CurrentContextMismatch)?;
114    let next_epoch = unsafe { next.bind_cpu(pin.area()) }?;
115    Ok((
116        PreparedContextSwitch {
117            next: next.as_non_null(),
118            next_epoch,
119            current_context: next.as_non_null().as_ptr() as usize,
120            area: pin.area(),
121            _scope: PhantomData,
122            _not_send_or_sync: PhantomData,
123        },
124        PreviousContextBinding {
125            previous: previous.as_non_null(),
126            epoch: previous_binding.epoch,
127        },
128    ))
129}
130
131#[cfg(all(test, feature = "host-test"))]
132mod tests {
133    use core::mem::MaybeUninit;
134
135    use super::*;
136    use crate::{
137        CpuAreaPrefix, CpuAreaRef, CpuIndex, install_bootstrap_context, install_cpu_area,
138        with_cpu_pin,
139    };
140
141    fn on_fresh_modeled_cpu(operation: impl FnOnce(CpuAreaRef) + Send + 'static) {
142        std::thread::spawn(move || {
143            let storage = Box::leak(Box::new(MaybeUninit::<CpuAreaPrefix>::uninit()));
144            let base = storage.as_mut_ptr() as usize;
145            storage.write(CpuAreaPrefix::initialize(CpuIndex::try_from(0).unwrap(), base).unwrap());
146            // SAFETY: the leaked prefix is initialized and remains mapped for
147            // the complete process lifetime.
148            let area = unsafe { CpuAreaRef::from_initialized_base(base) }.unwrap();
149            // SAFETY: this fresh host thread models one offline CPU and owns
150            // the completed area exclusively during register installation.
151            unsafe { install_cpu_area(area) }.unwrap();
152            operation(area);
153        })
154        .join()
155        .unwrap();
156    }
157
158    fn context_header() -> Pin<Box<ExecutionContextHeader>> {
159        Box::pin(ExecutionContextHeader::new())
160    }
161
162    #[test]
163    fn abandoned_prepare_rolls_back_next_binding() {
164        on_fresh_modeled_cpu(|area| {
165            let previous = context_header();
166            let next = context_header();
167
168            // SAFETY: the modeled CPU cannot migrate or receive interrupts.
169            unsafe {
170                with_cpu_pin(|pin| {
171                    install_bootstrap_context(pin, previous.as_ref()).unwrap();
172                    let (prepared, _previous_binding) =
173                        prepare_context_switch(pin, previous.as_ref(), next.as_ref()).unwrap();
174                    assert_eq!(current_context(pin), Ok(previous.as_ref().as_non_null()));
175                    assert_eq!(next.cpu_area(), Some(area));
176
177                    drop(prepared);
178
179                    assert_eq!(current_context(pin), Ok(previous.as_ref().as_non_null()));
180                    assert_eq!(next.cpu_area(), None);
181                })
182            }
183            .unwrap();
184        });
185    }
186
187    #[test]
188    fn prepare_reports_the_domain_mismatch_before_binding_next() {
189        on_fresh_modeled_cpu(|_| {
190            let published = context_header();
191            let wrong_previous = context_header();
192            let next = context_header();
193
194            // SAFETY: the modeled CPU cannot migrate or receive interrupts.
195            unsafe {
196                with_cpu_pin(|pin| {
197                    install_bootstrap_context(pin, published.as_ref()).unwrap();
198                    let result =
199                        prepare_context_switch(pin, wrong_previous.as_ref(), next.as_ref());
200                    assert!(matches!(
201                        result,
202                        Err(ContextSwitchError::CurrentContextMismatch)
203                    ));
204                    assert_eq!(next.cpu_area(), None);
205                })
206            }
207            .unwrap();
208        });
209    }
210
211    #[test]
212    fn publication_precedes_incoming_unbind() {
213        on_fresh_modeled_cpu(|area| {
214            let previous = context_header();
215            let next = context_header();
216
217            // SAFETY: this host model serializes the entire switch. Returning
218            // from `commit` represents resuming after the naked switch tail.
219            unsafe {
220                with_cpu_pin(|pin| {
221                    install_bootstrap_context(pin, previous.as_ref()).unwrap();
222                    let (prepared, previous_binding) =
223                        prepare_context_switch(pin, previous.as_ref(), next.as_ref()).unwrap();
224
225                    assert_eq!(current_context(pin), Ok(previous.as_ref().as_non_null()));
226                    assert_eq!(previous.cpu_area(), Some(area));
227                    assert_eq!(next.cpu_area(), Some(area));
228
229                    prepared.commit();
230
231                    assert_eq!(current_context(pin), Ok(next.as_ref().as_non_null()));
232                    assert_eq!(previous.cpu_area(), Some(area));
233                    previous_binding.finish(previous.as_ref()).unwrap();
234                    assert_eq!(previous.cpu_area(), None);
235                    assert_eq!(next.cpu_area(), Some(area));
236                })
237            }
238            .unwrap();
239        });
240    }
241
242    #[test]
243    fn stale_epoch_cannot_unbind_a_new_binding() {
244        on_fresh_modeled_cpu(|area| {
245            let header = context_header();
246            // SAFETY: this fresh modeled CPU exclusively owns the header.
247            let stale = unsafe { header.as_ref().bind_cpu(area) }.unwrap();
248            unsafe { header.as_ref().unbind_cpu(stale) }.unwrap();
249            let current = unsafe { header.as_ref().bind_cpu(area) }.unwrap();
250
251            assert_eq!(
252                unsafe { header.as_ref().unbind_cpu(stale) },
253                Err(ContextSwitchError::StalePreviousBinding)
254            );
255            assert_eq!(header.cpu_area(), Some(area));
256            unsafe { header.as_ref().unbind_cpu(current) }.unwrap();
257        });
258    }
259}