cpu_local/
switch.rs

1use core::{marker::PhantomData, mem::ManuallyDrop, pin::Pin, ptr::NonNull};
2
3use crate::{
4    CpuBindingEpoch, CpuLocalError, CpuPin, CurrentThreadHeader, ThreadSwitchError, current_thread,
5};
6
7/// Prepared current-thread publication owned by the final context-switch tail.
8#[must_use = "dropping an uncommitted switch rolls back the next CPU binding"]
9pub struct PreparedThreadSwitch<'switch> {
10    next: NonNull<CurrentThreadHeader>,
11    next_epoch: CpuBindingEpoch,
12    current_thread: usize,
13    area: crate::CpuAreaRef,
14    _scope: PhantomData<&'switch mut &'switch ()>,
15    _not_send_or_sync: PhantomData<*mut ()>,
16}
17
18impl PreparedThreadSwitch<'_> {
19    /// Returns the exact next header bound by this transaction.
20    #[doc(hidden)]
21    pub const fn next_header(&self) -> NonNull<CurrentThreadHeader> {
22        self.next
23    }
24
25    /// Publishes the prepared CPU runtime slot immediately before naked switch.
26    ///
27    /// # Safety
28    ///
29    /// The scheduler serialization and IRQ exclusion used during preparation
30    /// must still be active. The caller must enter the architecture switch
31    /// without performing fallible or ownership-sensitive Rust work.
32    #[doc(hidden)]
33    #[inline(always)]
34    pub unsafe fn commit(self) {
35        // Disarm rollback before publication. After the final call below there
36        // is no destructor state update or ownership-sensitive Rust work; the
37        // architecture wrapper enters its naked switch tail immediately.
38        let prepared = ManuallyDrop::new(self);
39        unsafe { crate::register::commit_current_thread(prepared.area, prepared.current_thread) };
40    }
41}
42
43impl Drop for PreparedThreadSwitch<'_> {
44    fn drop(&mut self) {
45        // SAFETY: an uncommitted token still owns the next binding, and its
46        // invariant lifetime keeps the scheduler critical section live. The
47        // preparation contract keeps the pinned header alive until this drop.
48        let next = unsafe { Pin::new_unchecked(self.next.as_ref()) };
49        if unsafe { next.unbind_cpu(self.next_epoch) }.is_err() {
50            panic!("prepared thread-switch rollback lost the next CPU binding");
51        }
52    }
53}
54
55/// Opaque previous-task binding consumed by the incoming switch tail.
56#[must_use = "the incoming task must withdraw the previous CPU binding"]
57#[derive(Debug)]
58pub struct PreviousThreadBinding {
59    previous: NonNull<CurrentThreadHeader>,
60    epoch: CpuBindingEpoch,
61}
62
63impl PreviousThreadBinding {
64    /// Withdraws the exact previous binding after architecture registers have
65    /// switched to the incoming task.
66    ///
67    /// # Errors
68    ///
69    /// Returns [`ThreadSwitchError::PreviousThreadMismatch`] if `previous`
70    /// differs from the prepared task, or
71    /// [`ThreadSwitchError::StalePreviousBinding`] for an obsolete tail.
72    ///
73    /// # Safety
74    ///
75    /// The incoming switch tail must be the sole owner of this token and the
76    /// previous task allocation must remain pinned and alive.
77    pub unsafe fn finish(
78        self,
79        previous: Pin<&CurrentThreadHeader>,
80    ) -> Result<(), ThreadSwitchError> {
81        if previous.as_non_null() != self.previous {
82            return Err(ThreadSwitchError::PreviousThreadMismatch);
83        }
84        unsafe { previous.unbind_cpu(self.epoch) }
85    }
86}
87
88/// Validates and binds a complete scheduler thread switch transaction.
89///
90/// All fallible validation occurs before the returned prepared token can be
91/// committed. Dropping the prepared token automatically rolls back `next`.
92///
93/// # Safety
94///
95/// The caller must own the IRQ-disabled scheduler switch path. Both headers
96/// must remain pinned and alive through the raw switch and incoming tail.
97pub unsafe fn prepare_thread_switch<'switch>(
98    pin: &'switch CpuPin<'_>,
99    previous: Pin<&CurrentThreadHeader>,
100    next: Pin<&CurrentThreadHeader>,
101) -> Result<(PreparedThreadSwitch<'switch>, PreviousThreadBinding), ThreadSwitchError> {
102    let published = current_thread(pin).map_err(|error| match error {
103        CpuLocalError::CurrentThreadMismatch => ThreadSwitchError::CurrentThreadMismatch,
104        other => ThreadSwitchError::CpuLocal(other),
105    })?;
106    if published != previous.as_non_null() {
107        return Err(ThreadSwitchError::CurrentThreadMismatch);
108    }
109    let previous_binding = previous
110        .cpu_binding()
111        .filter(|binding| binding.area == pin.area())
112        .ok_or(ThreadSwitchError::CurrentThreadMismatch)?;
113    let next_epoch = unsafe { next.bind_cpu(pin.area()) }?;
114    Ok((
115        PreparedThreadSwitch {
116            next: next.as_non_null(),
117            next_epoch,
118            current_thread: next.as_non_null().as_ptr() as usize,
119            area: pin.area(),
120            _scope: PhantomData,
121            _not_send_or_sync: PhantomData,
122        },
123        PreviousThreadBinding {
124            previous: previous.as_non_null(),
125            epoch: previous_binding.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, CurrentContext, install_bootstrap_thread,
137        install_cpu_area, 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 task_header(identity: usize) -> Pin<Box<CurrentThreadHeader>> {
158        Box::pin(CurrentThreadHeader::new(
159            CurrentContext::from_raw(identity).unwrap(),
160        ))
161    }
162
163    #[test]
164    fn abandoned_prepare_rolls_back_next_binding() {
165        on_fresh_modeled_cpu(|area| {
166            let previous = task_header(1);
167            let next = task_header(2);
168
169            // SAFETY: the modeled CPU cannot migrate or receive interrupts.
170            unsafe {
171                with_cpu_pin(|pin| {
172                    install_bootstrap_thread(pin, previous.as_ref()).unwrap();
173                    let (prepared, _previous_binding) =
174                        prepare_thread_switch(pin, previous.as_ref(), next.as_ref()).unwrap();
175                    assert_eq!(current_thread(pin), Ok(previous.as_ref().as_non_null()));
176                    assert_eq!(next.cpu_area(), Some(area));
177
178                    drop(prepared);
179
180                    assert_eq!(current_thread(pin), Ok(previous.as_ref().as_non_null()));
181                    assert_eq!(next.cpu_area(), None);
182                })
183            }
184            .unwrap();
185        });
186    }
187
188    #[test]
189    fn prepare_reports_the_domain_mismatch_before_binding_next() {
190        on_fresh_modeled_cpu(|_| {
191            let published = task_header(1);
192            let wrong_previous = task_header(2);
193            let next = task_header(3);
194
195            // SAFETY: the modeled CPU cannot migrate or receive interrupts.
196            unsafe {
197                with_cpu_pin(|pin| {
198                    install_bootstrap_thread(pin, published.as_ref()).unwrap();
199                    let result = prepare_thread_switch(pin, wrong_previous.as_ref(), next.as_ref());
200                    assert!(matches!(
201                        result,
202                        Err(ThreadSwitchError::CurrentThreadMismatch)
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 = task_header(1);
215            let next = task_header(2);
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_thread(pin, previous.as_ref()).unwrap();
222                    let (prepared, previous_binding) =
223                        prepare_thread_switch(pin, previous.as_ref(), next.as_ref()).unwrap();
224
225                    assert_eq!(current_thread(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_thread(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 = task_header(1);
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(ThreadSwitchError::StalePreviousBinding)
254            );
255            assert_eq!(header.cpu_area(), Some(area));
256            unsafe { header.as_ref().unbind_cpu(current) }.unwrap();
257        });
258    }
259}
cpu_local/switch.rs

cpu_local/
switch.rs
