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_epoch = previous
111        .binding_epoch_for_area(pin.area())
112        .ok_or(ContextSwitchError::CurrentContextMismatch)?;
113    let next_epoch = unsafe { next.bind_cpu(pin.area()) }?;
114    Ok((
115        PreparedContextSwitch {
116            next: next.as_non_null(),
117            next_epoch,
118            current_context: next.as_non_null().as_ptr() as usize,
119            area: pin.area(),
120            _scope: PhantomData,
121            _not_send_or_sync: PhantomData,
122        },
123        PreviousContextBinding {
124            previous: previous.as_non_null(),
125            epoch: previous_epoch,
126        },
127    ))
128}
129
130#[cfg(all(test, feature = "host-test"))]
131mod tests {
132    use core::mem::MaybeUninit;
133
134    use super::*;
135    use crate::{
136        CpuAreaPrefix, CpuAreaRef, CpuIndex, install_bootstrap_context, install_cpu_area,
137        with_cpu_pin,
138    };
139
140    fn on_fresh_modeled_cpu(operation: impl FnOnce(CpuAreaRef) + Send + 'static) {
141        std::thread::spawn(move || {
142            let storage = Box::leak(Box::new(MaybeUninit::<CpuAreaPrefix>::uninit()));
143            let base = storage.as_mut_ptr() as usize;
144            storage.write(CpuAreaPrefix::initialize(CpuIndex::try_from(0).unwrap(), base).unwrap());
145            // SAFETY: the leaked prefix is initialized and remains mapped for
146            // the complete process lifetime.
147            let area = unsafe { CpuAreaRef::from_initialized_base(base) }.unwrap();
148            // SAFETY: this fresh host thread models one offline CPU and owns
149            // the completed area exclusively during register installation.
150            unsafe { install_cpu_area(area) }.unwrap();
151            operation(area);
152        })
153        .join()
154        .unwrap();
155    }
156
157    fn context_header() -> Pin<Box<ExecutionContextHeader>> {
158        Box::pin(ExecutionContextHeader::new())
159    }
160
161    #[test]
162    fn abandoned_prepare_rolls_back_next_binding() {
163        on_fresh_modeled_cpu(|area| {
164            let previous = context_header();
165            let next = context_header();
166
167            // SAFETY: the modeled CPU cannot migrate or receive interrupts.
168            unsafe {
169                with_cpu_pin(|pin| {
170                    install_bootstrap_context(pin, previous.as_ref()).unwrap();
171                    let (prepared, _previous_binding) =
172                        prepare_context_switch(pin, previous.as_ref(), next.as_ref()).unwrap();
173                    assert_eq!(current_context(pin), Ok(previous.as_ref().as_non_null()));
174                    assert_eq!(next.cpu_area(), Some(area));
175
176                    drop(prepared);
177
178                    assert_eq!(current_context(pin), Ok(previous.as_ref().as_non_null()));
179                    assert_eq!(next.cpu_area(), None);
180                })
181            }
182            .unwrap();
183        });
184    }
185
186    #[test]
187    fn prepare_reports_the_domain_mismatch_before_binding_next() {
188        on_fresh_modeled_cpu(|_| {
189            let published = context_header();
190            let wrong_previous = context_header();
191            let next = context_header();
192
193            // SAFETY: the modeled CPU cannot migrate or receive interrupts.
194            unsafe {
195                with_cpu_pin(|pin| {
196                    install_bootstrap_context(pin, published.as_ref()).unwrap();
197                    let result =
198                        prepare_context_switch(pin, wrong_previous.as_ref(), next.as_ref());
199                    assert!(matches!(
200                        result,
201                        Err(ContextSwitchError::CurrentContextMismatch)
202                    ));
203                    assert_eq!(next.cpu_area(), None);
204                })
205            }
206            .unwrap();
207        });
208    }
209
210    #[test]
211    fn publication_precedes_incoming_unbind() {
212        on_fresh_modeled_cpu(|area| {
213            let previous = context_header();
214            let next = context_header();
215
216            // SAFETY: this host model serializes the entire switch. Returning
217            // from `commit` represents resuming after the naked switch tail.
218            unsafe {
219                with_cpu_pin(|pin| {
220                    install_bootstrap_context(pin, previous.as_ref()).unwrap();
221                    let (prepared, previous_binding) =
222                        prepare_context_switch(pin, previous.as_ref(), next.as_ref()).unwrap();
223
224                    assert_eq!(current_context(pin), Ok(previous.as_ref().as_non_null()));
225                    assert_eq!(previous.cpu_area(), Some(area));
226                    assert_eq!(next.cpu_area(), Some(area));
227
228                    prepared.commit();
229
230                    assert_eq!(current_context(pin), Ok(next.as_ref().as_non_null()));
231                    assert_eq!(previous.cpu_area(), Some(area));
232                    previous_binding.finish(previous.as_ref()).unwrap();
233                    assert_eq!(previous.cpu_area(), None);
234                    assert_eq!(next.cpu_area(), Some(area));
235                })
236            }
237            .unwrap();
238        });
239    }
240
241    #[test]
242    fn stale_epoch_cannot_unbind_a_new_binding() {
243        on_fresh_modeled_cpu(|area| {
244            let header = context_header();
245            // SAFETY: this fresh modeled CPU exclusively owns the header.
246            let stale = unsafe { header.as_ref().bind_cpu(area) }.unwrap();
247            unsafe { header.as_ref().unbind_cpu(stale) }.unwrap();
248            let current = unsafe { header.as_ref().bind_cpu(area) }.unwrap();
249
250            assert_eq!(
251                unsafe { header.as_ref().unbind_cpu(stale) },
252                Err(ContextSwitchError::StalePreviousBinding)
253            );
254            assert_eq!(header.cpu_area(), Some(area));
255            unsafe { header.as_ref().unbind_cpu(current) }.unwrap();
256        });
257    }
258}