ostd/cpu/
id.rs

1// SPDX-License-Identifier: MPL-2.0
2
3//! CPU identification numbers.
4
5pub use current::PinCurrentCpu;
6pub use set::{AtomicCpuSet, CpuSet};
7
8use crate::util::id_set::Id;
9
10/// The ID of a CPU in the system.
11///
12/// If converting from/to an integer, the integer must start from 0 and be less
13/// than the number of CPUs.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub struct CpuId(u32);
16
17impl CpuId {
18    /// Creates a new instance.
19    ///
20    /// # Panics
21    ///
22    /// The given number must be smaller than the total number of CPUs
23    /// (`ostd::cpu::num_cpus()`).
24    pub fn new(raw_id: u32) -> Self {
25        assert!(raw_id < num_cpus() as u32);
26        // SAFETY: The raw ID is smaller than `num_cpus()`.
27        unsafe { Self::new_unchecked(raw_id) }
28    }
29
30    /// Returns the CPU ID of the bootstrap processor (BSP).
31    ///
32    /// The number for the BSP is always zero.
33    pub const fn bsp() -> Self {
34        // BSP's `CURRENT_CPU` is assigned a value of 0.
35        let bsp_raw_cpu_id = 0;
36        // SAFETY: There is at least one CPU.
37        Self(bsp_raw_cpu_id)
38    }
39}
40
41impl From<CpuId> for u32 {
42    fn from(cpu_id: CpuId) -> Self {
43        cpu_id.0
44    }
45}
46
47/// The error type returned when converting an out-of-range integer to [`CpuId`].
48#[derive(Debug, Clone, Copy)]
49pub struct CpuIdFromIntError;
50
51impl TryFrom<usize> for CpuId {
52    type Error = CpuIdFromIntError;
53
54    fn try_from(raw_id: usize) -> Result<Self, Self::Error> {
55        if raw_id < num_cpus() {
56            // SAFETY: The raw ID is smaller than `num_cpus()`.
57            let new_self = unsafe { CpuId::new_unchecked(raw_id as u32) };
58            Ok(new_self)
59        } else {
60            Err(CpuIdFromIntError)
61        }
62    }
63}
64
65/// Returns the number of CPUs.
66pub fn num_cpus() -> usize {
67    // SAFETY: As far as the safe APIs are concerned, `NUM_CPUS` is
68    // read-only, so it is always valid to read.
69    (unsafe { NUM_CPUS }) as usize
70}
71
72static mut NUM_CPUS: u32 = 1;
73
74/// Returns an iterator over all CPUs.
75pub fn all_cpus() -> impl Iterator<Item = CpuId> {
76    (0..num_cpus()).map(|raw_id| {
77        // SAFETY: The raw ID is smaller than `num_cpus()`.
78        unsafe { CpuId::new_unchecked(raw_id as u32) }
79    })
80}
81
82mod set {
83    use super::{CpuId, num_cpus};
84    use crate::util::id_set::{AtomicIdSet, Id, IdSet};
85
86    /// A set of CPU IDs.
87    pub type CpuSet = IdSet<CpuId>;
88
89    /// A set of CPU IDs, with support for concurrent access.
90    pub type AtomicCpuSet = AtomicIdSet<CpuId>;
91
92    // SAFETY: `CpuId`s and the integers within 0 to `num_cpus` (exclusive)
93    // have 1:1 mapping.
94    unsafe impl Id for CpuId {
95        unsafe fn new_unchecked(raw_id: u32) -> Self {
96            Self(raw_id)
97        }
98
99        fn cardinality() -> u32 {
100            num_cpus() as u32
101        }
102    }
103}
104
105mod current {
106    //! The current CPU ID.
107
108    use super::CpuId;
109    use crate::{cpu_local_cell, task::atomic_mode::InAtomicMode, util::id_set::Id};
110
111    /// A marker trait for guard types that can "pin" the current task to the
112    /// current CPU.
113    ///
114    /// Such guard types include [`DisabledLocalIrqGuard`] and
115    /// [`DisabledPreemptGuard`]. When such guards exist, the CPU executing the
116    /// current task is pinned. So getting the current CPU ID or CPU-local
117    /// variables are safe.
118    ///
119    /// # Safety
120    ///
121    /// The implementor must ensure that the current task is pinned to the current
122    /// CPU while any one of the instances of the implemented structure exists.
123    ///
124    /// [`DisabledLocalIrqGuard`]: crate::irq::DisabledLocalIrqGuard
125    /// [`DisabledPreemptGuard`]: crate::task::DisabledPreemptGuard
126    pub unsafe trait PinCurrentCpu {
127        /// Returns the ID of the current CPU.
128        fn current_cpu(&self) -> CpuId {
129            CpuId::current_racy()
130        }
131    }
132
133    // SAFETY: A guard that enforces the atomic mode requires disabling any
134    // context switching. So naturally, the current task is pinned on the CPU.
135    unsafe impl<T: InAtomicMode> PinCurrentCpu for T {}
136    unsafe impl PinCurrentCpu for dyn InAtomicMode + '_ {}
137
138    impl CpuId {
139        /// Returns the ID of the current CPU.
140        ///
141        /// This function is safe to call, but is vulnerable to races. The returned CPU
142        /// ID may be outdated if the task migrates to another CPU.
143        ///
144        /// To ensure that the CPU ID is up-to-date, do it under any guards that
145        /// implement the [`PinCurrentCpu`] trait.
146        pub fn current_racy() -> Self {
147            #[cfg(debug_assertions)]
148            assert!(IS_CURRENT_CPU_INITED.load());
149
150            let current_raw_id = CURRENT_CPU.load();
151            // SAFETY: The CPU-local value is initialized to a correct one.
152            unsafe { Self::new_unchecked(current_raw_id) }
153        }
154    }
155
156    /// Initializes the module on the current CPU.
157    ///
158    /// Note that this method will use the current CPU's CPU-local storage.
159    ///
160    /// # Safety
161    ///
162    /// The caller must ensure:
163    /// 1. This method is called on each CPU in the boot context.
164    /// 2. The CPU ID for the current CPU is correct.
165    pub(super) unsafe fn init_on_cpu(current_cpu_id: u32) {
166        // FIXME: If there are safe APIs that rely on the correctness of
167        // the CPU ID for soundness, we'd better make the CPU ID a global
168        // invariant and initialize it before entering `ap_early_entry`.
169        CURRENT_CPU.store(current_cpu_id);
170
171        #[cfg(debug_assertions)]
172        IS_CURRENT_CPU_INITED.store(true);
173    }
174
175    cpu_local_cell! {
176        /// The current CPU ID.
177        pub(super) static CURRENT_CPU: u32 = 0;
178        /// The initialization state of the current CPU ID.
179        #[cfg(debug_assertions)]
180        pub(super) static IS_CURRENT_CPU_INITED: bool = false;
181    }
182}
183
184/// Initializes the CPU ID module (the BSP part).
185///
186/// Note that this method will use the BSP's CPU-local storage.
187///
188/// # Safety
189///
190/// The caller must ensure that
191/// 1. We're in the boot context of the BSP and APs have not yet booted.
192/// 2. The number of CPUs is correct.
193pub(super) unsafe fn init_on_bsp(num_cpus: u32) {
194    // SAFETY:
195    // 1. We're in the boot context of the BSP.
196    // 2. The CPU ID of BSP has a value of zero.
197    unsafe { current::init_on_cpu(0) };
198
199    // SAFETY:
200    // 1. We're in the boot context of the BSP and APs have not yet booted.
201    // 2. The number of CPUs is correct.
202    unsafe { init_num_cpus(num_cpus) };
203}
204
205/// Initializes the number of CPUs.
206///
207/// The number of CPUs is a fixed value,
208/// since we don't support CPU hot-plugging.
209///
210/// # Safety
211///
212/// The caller must ensure that
213/// 1. We're in the boot context of the BSP and APs have not yet booted.
214/// 2. The number of CPUs is correct.
215unsafe fn init_num_cpus(num_cpus: u32) {
216    // Thanks to this assertion,
217    // it's only legal to call this function to
218    // increase the number of CPUs from one (the initial value)
219    // to the actual number of CPUs.
220    assert!(num_cpus >= 1);
221
222    // SAFETY: It is safe to mutate this global variable because we
223    // are in the boot context.
224    unsafe { NUM_CPUS = num_cpus };
225}
226
227/// Initializes the CPU ID module (the AP part).
228///
229/// Note that this method will use the BSP's CPU-local storage.
230/// This should be fine
231/// because `crate::cpu::init_on_bsp` must have been invoked before APs boot.
232///
233/// # Safety
234///
235/// The caller must ensure that:
236/// 1. We're in the boot context of an AP.
237/// 2. The CPU ID of the AP is correct.
238pub(super) unsafe fn init_on_ap(cpu_id: u32) {
239    // SAFETY:
240    // 1. We're in the boot context of the AP.
241    // 2. The CPU ID for the AP is correct.
242    unsafe { current::init_on_cpu(cpu_id) };
243}
ostd/cpu/id.rs

ostd/cpu/
id.rs
