From f7a9510be0a21ada005e59d8c525ed318376e146 Mon Sep 17 00:00:00 2001 From: Zhang Junyang Date: Mon, 19 Aug 2024 20:32:14 +0800 Subject: [PATCH] Refactor the `this_cpu` API with `PinCurrentCpu` --- kernel/src/sched/priority_scheduler.rs | 9 ++- kernel/src/time/clocks/system_wide.rs | 18 ++++-- osdk/src/base_crate/x86_64.ld.template | 4 -- ostd/src/arch/x86/mod.rs | 7 ++- ostd/src/boot/smp.rs | 4 +- ostd/src/cpu/local/cpu_local.rs | 10 +-- ostd/src/cpu/local/mod.rs | 14 +---- ostd/src/cpu/mod.rs | 76 ++++++++++++++++++----- ostd/src/mm/page_table/node.rs | 8 +-- ostd/src/sync/rwlock.rs | 4 +- ostd/src/sync/spin.rs | 4 +- ostd/src/task/mod.rs | 2 +- ostd/src/task/preempt/cpu_local.rs | 2 +- ostd/src/task/preempt/guard.rs | 12 ++-- ostd/src/task/preempt/mod.rs | 2 +- ostd/src/task/scheduler/fifo_scheduler.rs | 14 +++-- ostd/src/task/scheduler/mod.rs | 8 ++- 17 files changed, 123 insertions(+), 75 deletions(-) diff --git a/kernel/src/sched/priority_scheduler.rs b/kernel/src/sched/priority_scheduler.rs index 8185124df..415576430 100644 --- a/kernel/src/sched/priority_scheduler.rs +++ b/kernel/src/sched/priority_scheduler.rs @@ -1,11 +1,12 @@ // SPDX-License-Identifier: MPL-2.0 use ostd::{ - cpu::{num_cpus, this_cpu}, + cpu::{num_cpus, PinCurrentCpu}, task::{ scheduler::{inject_scheduler, EnqueueFlags, LocalRunQueue, Scheduler, UpdateFlags}, AtomicCpuId, Priority, Task, }, + trap::disable_local, }; use crate::prelude::*; @@ -71,13 +72,15 @@ impl Scheduler for PreemptScheduler { } fn local_rq_with(&self, f: &mut dyn FnMut(&dyn LocalRunQueue)) { - let local_rq: &PreemptRunQueue = &self.rq[this_cpu() as usize].disable_irq().lock(); + let irq_guard = disable_local(); + let local_rq: &PreemptRunQueue = &self.rq[irq_guard.current_cpu() as usize].lock(); f(local_rq); } fn local_mut_rq_with(&self, f: &mut dyn FnMut(&mut dyn LocalRunQueue)) { + let irq_guard = disable_local(); let local_rq: &mut PreemptRunQueue = - &mut self.rq[this_cpu() as usize].disable_irq().lock(); + &mut self.rq[irq_guard.current_cpu() as usize].lock(); f(local_rq); } } diff --git a/kernel/src/time/clocks/system_wide.rs b/kernel/src/time/clocks/system_wide.rs index fdd82e1be..0969d53e5 100644 --- a/kernel/src/time/clocks/system_wide.rs +++ b/kernel/src/time/clocks/system_wide.rs @@ -6,9 +6,10 @@ use core::time::Duration; use aster_time::read_monotonic_time; use ostd::{ arch::timer::Jiffies, - cpu::{num_cpus, this_cpu}, + cpu::{num_cpus, PinCurrentCpu}, cpu_local, sync::SpinLock, + task::disable_preempt, }; use paste::paste; use spin::Once; @@ -35,7 +36,11 @@ impl RealTimeClock { /// Get the cpu-local system-wide `TimerManager` singleton of this clock. pub fn timer_manager() -> &'static Arc { - CLOCK_REALTIME_MANAGER.get_on_cpu(this_cpu()).get().unwrap() + let preempt_guard = disable_preempt(); + CLOCK_REALTIME_MANAGER + .get_on_cpu(preempt_guard.current_cpu()) + .get() + .unwrap() } } @@ -53,8 +58,9 @@ impl MonotonicClock { /// Get the cpu-local system-wide `TimerManager` singleton of this clock. pub fn timer_manager() -> &'static Arc { + let preempt_guard = disable_preempt(); CLOCK_MONOTONIC_MANAGER - .get_on_cpu(this_cpu()) + .get_on_cpu(preempt_guard.current_cpu()) .get() .unwrap() } @@ -135,7 +141,11 @@ impl BootTimeClock { /// Get the cpu-local system-wide `TimerManager` singleton of this clock. pub fn timer_manager() -> &'static Arc { - CLOCK_BOOTTIME_MANAGER.get_on_cpu(this_cpu()).get().unwrap() + let preempt_guard = disable_preempt(); + CLOCK_BOOTTIME_MANAGER + .get_on_cpu(preempt_guard.current_cpu()) + .get() + .unwrap() } } diff --git a/osdk/src/base_crate/x86_64.ld.template b/osdk/src/base_crate/x86_64.ld.template index 2803d1953..59775408b 100644 --- a/osdk/src/base_crate/x86_64.ld.template +++ b/osdk/src/base_crate/x86_64.ld.template @@ -119,10 +119,6 @@ SECTIONS . = ALIGN(4096); .cpu_local : AT(ADDR(.cpu_local) - KERNEL_VMA) { __cpu_local_start = .; - - # These 4 bytes are used to store the CPU ID. - . += 4; - KEEP(*(SORT(.cpu_local))) __cpu_local_end = .; } diff --git a/ostd/src/arch/x86/mod.rs b/ostd/src/arch/x86/mod.rs index 51c181d1f..11475f0c2 100644 --- a/ostd/src/arch/x86/mod.rs +++ b/ostd/src/arch/x86/mod.rs @@ -63,8 +63,11 @@ pub(crate) fn init_on_bsp() { irq::init(); kernel::acpi::init(); - // SAFETY: it is only called once and ACPI has been initialized. - unsafe { crate::cpu::init() }; + // SAFETY: they are only called once on BSP and ACPI has been initialized. + unsafe { + crate::cpu::init_num_cpus(); + crate::cpu::set_this_cpu_id(0); + } match kernel::apic::init() { Ok(_) => { diff --git a/ostd/src/boot/smp.rs b/ostd/src/boot/smp.rs index 7748e3f63..b442cd974 100644 --- a/ostd/src/boot/smp.rs +++ b/ostd/src/boot/smp.rs @@ -113,9 +113,11 @@ pub fn register_ap_entry(entry: fn() -> !) { fn ap_early_entry(local_apic_id: u32) -> ! { crate::arch::enable_cpu_features(); - // SAFETY: we are on the AP. + // SAFETY: we are on the AP and they are only called once with the correct + // CPU ID. unsafe { cpu::local::init_on_ap(local_apic_id); + cpu::set_this_cpu_id(local_apic_id); } trap::init(); diff --git a/ostd/src/cpu/local/cpu_local.rs b/ostd/src/cpu/local/cpu_local.rs index 402990262..960761aa8 100644 --- a/ostd/src/cpu/local/cpu_local.rs +++ b/ostd/src/cpu/local/cpu_local.rs @@ -23,7 +23,7 @@ use crate::{ /// # Example /// /// ```rust -/// use ostd::{cpu_local, cpu::this_cpu}; +/// use ostd::{cpu_local, cpu::PinCurrentCpu, task::disable_preempt}; /// use core::{sync::atomic::{AtomicU32, Ordering}, cell::Cell}; /// /// cpu_local! { @@ -32,16 +32,12 @@ use crate::{ /// } /// /// fn not_an_atomic_function() { -/// let ref_of_foo = FOO.get_on_cpu(this_cpu()); -/// // Note that the value of `FOO` here doesn't necessarily equal to the value -/// // of `FOO` of exactly the __current__ CPU. Since that task may be preempted -/// // and moved to another CPU since `ref_of_foo` is created. +/// let preempt_guard = disable_preempt(); +/// let ref_of_foo = FOO.get_on_cpu(preempt_guard.current_cpu()); /// let val_of_foo = ref_of_foo.load(Ordering::Relaxed); /// println!("FOO VAL: {}", val_of_foo); /// /// let bar_guard = BAR.borrow_irq_disabled(); -/// // Here the value of `BAR` is always the one in the __current__ CPU since -/// // interrupts are disabled and we do not explicitly yield execution here. /// let val_of_bar = bar_guard.get(); /// println!("BAR VAL: {}", val_of_bar); /// } diff --git a/ostd/src/cpu/local/mod.rs b/ostd/src/cpu/local/mod.rs index 92e836423..0f64479a9 100644 --- a/ostd/src/cpu/local/mod.rs +++ b/ostd/src/cpu/local/mod.rs @@ -98,7 +98,7 @@ pub unsafe fn init_on_bsp() { let num_cpus = super::num_cpus(); let mut cpu_local_storages = Vec::with_capacity(num_cpus as usize - 1); - for cpu_i in 1..num_cpus { + for _ in 1..num_cpus { let ap_pages = { let nbytes = (bsp_end_va - bsp_base_va).align_up(PAGE_SIZE); page::allocator::alloc_contiguous(nbytes, |_| KernelMeta::default()).unwrap() @@ -116,23 +116,11 @@ pub unsafe fn init_on_bsp() { ); } - // SAFETY: bytes `0:4` are reserved for storing CPU ID. - unsafe { - (ap_pages_ptr as *mut u32).write(cpu_i); - } - cpu_local_storages.push(ap_pages); } CPU_LOCAL_STORAGES.call_once(|| cpu_local_storages); - // Write the CPU ID of BSP to the first 4 bytes of the CPU-local area. - let bsp_cpu_id_ptr = bsp_base_va as *mut u32; - // SAFETY: the first 4 bytes is reserved for storing CPU ID. - unsafe { - bsp_cpu_id_ptr.write(0); - } - arch::cpu::local::set_base(bsp_base_va as u64); has_init::set_true(); diff --git a/ostd/src/cpu/mod.rs b/ostd/src/cpu/mod.rs index 17289230c..40913c763 100644 --- a/ostd/src/cpu/mod.rs +++ b/ostd/src/cpu/mod.rs @@ -11,43 +11,85 @@ cfg_if::cfg_if! { } use alloc::vec::Vec; -use core::sync::atomic::{AtomicU32, Ordering}; use bitvec::{ prelude::{BitVec, Lsb0}, slice::IterOnes, }; +use local::cpu_local_cell; +use spin::Once; -use crate::arch::{self, boot::smp::get_num_processors}; +use crate::{ + arch::boot::smp::get_num_processors, task::DisabledPreemptGuard, trap::DisabledLocalIrqGuard, +}; -/// The number of CPUs. Zero means uninitialized. -static NUM_CPUS: AtomicU32 = AtomicU32::new(0); +/// The number of CPUs. +static NUM_CPUS: Once = Once::new(); /// Initializes the number of CPUs. /// /// # Safety /// -/// The caller must ensure that this function is called only once at the -/// correct time when the number of CPUs is available from the platform. -pub unsafe fn init() { +/// The caller must ensure that this function is called only once on the BSP +/// at the correct time when the number of CPUs is available from the platform. +pub(crate) unsafe fn init_num_cpus() { let num_processors = get_num_processors().unwrap_or(1); - NUM_CPUS.store(num_processors, Ordering::Release) + NUM_CPUS.call_once(|| num_processors); +} + +/// Initializes the number of the current CPU. +/// +/// # Safety +/// +/// The caller must ensure that this function is called only once on the +/// correct CPU with the correct CPU ID. +pub(crate) unsafe fn set_this_cpu_id(id: u32) { + CURRENT_CPU.store(id); } /// Returns the number of CPUs. pub fn num_cpus() -> u32 { - let num = NUM_CPUS.load(Ordering::Acquire); - debug_assert_ne!(num, 0, "The number of CPUs is not initialized"); - num + debug_assert!( + NUM_CPUS.get().is_some(), + "The number of CPUs is not initialized" + ); + // SAFETY: The number of CPUs is initialized. The unsafe version is used + // to avoid the overhead of the check. + unsafe { *NUM_CPUS.get_unchecked() } } -/// Returns the ID of this CPU. +/// A marker trait for guard types that can "pin" the current task to the +/// current CPU. /// -/// The CPU ID is strategically placed at the beginning of the CPU local storage area. -pub fn this_cpu() -> u32 { - // SAFETY: the cpu ID is stored at the beginning of the cpu local area, provided - // by the linker script. - unsafe { (arch::cpu::local::get_base() as usize as *mut u32).read() } +/// Such guard types include [`DisabledLocalIrqGuard`] and +/// [`DisabledPreemptGuard`]. When such guards exist, the CPU executing the +/// current task is pinned. So getting the current CPU ID or CPU-local +/// variables are safe. +/// +/// # Safety +/// +/// The implementor must ensure that the current task is pinned to the current +/// CPU while any one of the instances of the implemented structure exists. +pub unsafe trait PinCurrentCpu { + /// Returns the number of the current CPU. + fn current_cpu(&self) -> u32 { + let id = CURRENT_CPU.load(); + debug_assert_ne!(id, u32::MAX, "This CPU is not initialized"); + id + } +} + +// SAFETY: When IRQs are disabled, the task cannot be passively preempted and +// migrates to another CPU. If the task actively calls `yield`, it will not be +// successful either. +unsafe impl PinCurrentCpu for DisabledLocalIrqGuard {} +// SAFETY: When preemption is disabled, the task cannot be preempted and migrates +// to another CPU. +unsafe impl PinCurrentCpu for DisabledPreemptGuard {} + +cpu_local_cell! { + /// The number of the current CPU. + static CURRENT_CPU: u32 = u32::MAX; } /// A subset of all CPUs in the system. diff --git a/ostd/src/mm/page_table/node.rs b/ostd/src/mm/page_table/node.rs index 583359a76..b6da9b8df 100644 --- a/ostd/src/mm/page_table/node.rs +++ b/ostd/src/mm/page_table/node.rs @@ -42,7 +42,7 @@ use crate::{ page_prop::PageProperty, Paddr, PagingConstsTrait, PagingLevel, PAGE_SIZE, }, - task::{disable_preempt, DisablePreemptGuard}, + task::{disable_preempt, DisabledPreemptGuard}, }; /// The raw handle to a page table node. @@ -190,11 +190,11 @@ pub(super) struct PageTableNode< [(); C::NR_LEVELS as usize]:, { pub(super) page: Page>, - preempt_guard: DisablePreemptGuard, + preempt_guard: DisabledPreemptGuard, } -// FIXME: We cannot `#[derive(Debug)]` here due to `DisablePreemptGuard`. Should we skip -// this field or implement the `Debug` trait also for `DisablePreemptGuard`? +// FIXME: We cannot `#[derive(Debug)]` here due to `DisabledPreemptGuard`. Should we skip +// this field or implement the `Debug` trait also for `DisabledPreemptGuard`? impl fmt::Debug for PageTableNode where E: PageTableEntryTrait, diff --git a/ostd/src/sync/rwlock.rs b/ostd/src/sync/rwlock.rs index ba1999156..2f02bd955 100644 --- a/ostd/src/sync/rwlock.rs +++ b/ostd/src/sync/rwlock.rs @@ -14,7 +14,7 @@ use core::{ }; use crate::{ - task::{disable_preempt, DisablePreemptGuard}, + task::{disable_preempt, DisabledPreemptGuard}, trap::{disable_local, DisabledLocalIrqGuard}, }; @@ -544,7 +544,7 @@ unsafe impl> + Clone + Sync> Sync enum InnerGuard { IrqGuard(DisabledLocalIrqGuard), - PreemptGuard(DisablePreemptGuard), + PreemptGuard(DisabledPreemptGuard), } impl InnerGuard { diff --git a/ostd/src/sync/spin.rs b/ostd/src/sync/spin.rs index dfa2bb8ee..87a6e7463 100644 --- a/ostd/src/sync/spin.rs +++ b/ostd/src/sync/spin.rs @@ -12,7 +12,7 @@ use core::{ }; use crate::{ - task::{disable_preempt, DisablePreemptGuard}, + task::{disable_preempt, DisabledPreemptGuard}, trap::{disable_local, DisabledLocalIrqGuard}, }; @@ -54,7 +54,7 @@ pub trait Guardian { pub struct PreemptDisabled; impl Guardian for PreemptDisabled { - type Guard = DisablePreemptGuard; + type Guard = DisabledPreemptGuard; fn guard() -> Self::Guard { disable_preempt() diff --git a/ostd/src/task/mod.rs b/ostd/src/task/mod.rs index 2dff187f6..fcb9a7fd6 100644 --- a/ostd/src/task/mod.rs +++ b/ostd/src/task/mod.rs @@ -9,6 +9,6 @@ pub mod scheduler; mod task; pub use self::{ - preempt::{disable_preempt, DisablePreemptGuard}, + preempt::{disable_preempt, DisabledPreemptGuard}, task::{AtomicCpuId, Priority, Task, TaskAdapter, TaskContextApi, TaskOptions}, }; diff --git a/ostd/src/task/preempt/cpu_local.rs b/ostd/src/task/preempt/cpu_local.rs index 23de46cf9..d3fb9d227 100644 --- a/ostd/src/task/preempt/cpu_local.rs +++ b/ostd/src/task/preempt/cpu_local.rs @@ -4,7 +4,7 @@ //! on a CPU with a single 32-bit, CPU-local integer value. //! //! * Bits from 0 to 30 represents an unsigned counter called `guard_count`, -//! which is the number of `DisablePreemptGuard` instances held by the +//! which is the number of `DisabledPreemptGuard` instances held by the //! current CPU; //! * Bit 31 is set to `!need_preempt`, where `need_preempt` is a boolean value //! that will be set by the scheduler when it decides that the current task diff --git a/ostd/src/task/preempt/guard.rs b/ostd/src/task/preempt/guard.rs index 6ce86b577..a8e9bd8d9 100644 --- a/ostd/src/task/preempt/guard.rs +++ b/ostd/src/task/preempt/guard.rs @@ -3,14 +3,14 @@ /// A guard for disable preempt. #[clippy::has_significant_drop] #[must_use] -pub struct DisablePreemptGuard { +pub struct DisabledPreemptGuard { // This private field prevents user from constructing values of this type directly. _private: (), } -impl !Send for DisablePreemptGuard {} +impl !Send for DisabledPreemptGuard {} -impl DisablePreemptGuard { +impl DisabledPreemptGuard { fn new() -> Self { super::cpu_local::inc_guard_count(); Self { _private: () } @@ -23,13 +23,13 @@ impl DisablePreemptGuard { } } -impl Drop for DisablePreemptGuard { +impl Drop for DisabledPreemptGuard { fn drop(&mut self) { super::cpu_local::dec_guard_count(); } } /// Disables preemption. -pub fn disable_preempt() -> DisablePreemptGuard { - DisablePreemptGuard::new() +pub fn disable_preempt() -> DisabledPreemptGuard { + DisabledPreemptGuard::new() } diff --git a/ostd/src/task/preempt/mod.rs b/ostd/src/task/preempt/mod.rs index 3e6d7b6e3..4ba707896 100644 --- a/ostd/src/task/preempt/mod.rs +++ b/ostd/src/task/preempt/mod.rs @@ -3,4 +3,4 @@ pub(super) mod cpu_local; mod guard; -pub use self::guard::{disable_preempt, DisablePreemptGuard}; +pub use self::guard::{disable_preempt, DisabledPreemptGuard}; diff --git a/ostd/src/task/scheduler/fifo_scheduler.rs b/ostd/src/task/scheduler/fifo_scheduler.rs index 9e1243885..c338c6d11 100644 --- a/ostd/src/task/scheduler/fifo_scheduler.rs +++ b/ostd/src/task/scheduler/fifo_scheduler.rs @@ -4,9 +4,9 @@ use alloc::{boxed::Box, collections::VecDeque, sync::Arc, vec::Vec}; use super::{inject_scheduler, EnqueueFlags, LocalRunQueue, Scheduler, UpdateFlags}; use crate::{ - cpu::{num_cpus, this_cpu}, + cpu::{num_cpus, PinCurrentCpu}, sync::SpinLock, - task::{AtomicCpuId, Task}, + task::{disable_preempt, AtomicCpuId, Task}, }; pub fn init() { @@ -61,12 +61,18 @@ impl Scheduler for FifoScheduler { } fn local_rq_with(&self, f: &mut dyn FnMut(&dyn LocalRunQueue)) { - let local_rq: &FifoRunQueue = &self.rq[this_cpu() as usize].disable_irq().lock(); + let preempt_guard = disable_preempt(); + let local_rq: &FifoRunQueue = &self.rq[preempt_guard.current_cpu() as usize] + .disable_irq() + .lock(); f(local_rq); } fn local_mut_rq_with(&self, f: &mut dyn FnMut(&mut dyn LocalRunQueue)) { - let local_rq: &mut FifoRunQueue = &mut self.rq[this_cpu() as usize].disable_irq().lock(); + let preempt_guard = disable_preempt(); + let local_rq: &mut FifoRunQueue = &mut self.rq[preempt_guard.current_cpu() as usize] + .disable_irq() + .lock(); f(local_rq); } } diff --git a/ostd/src/task/scheduler/mod.rs b/ostd/src/task/scheduler/mod.rs index 59ba0791e..c3da3d7e1 100644 --- a/ostd/src/task/scheduler/mod.rs +++ b/ostd/src/task/scheduler/mod.rs @@ -12,7 +12,7 @@ use core::sync::atomic::{AtomicBool, Ordering}; use spin::Once; use super::{preempt::cpu_local, processor, task::Task}; -use crate::{arch::timer, cpu::this_cpu, prelude::*}; +use crate::{arch::timer, cpu::PinCurrentCpu, prelude::*, task::disable_preempt}; /// Injects a scheduler implementation into framework. /// @@ -140,8 +140,9 @@ pub(crate) fn unpark_target(runnable: Arc) { .enqueue(runnable, EnqueueFlags::Wake); if need_preempt_info.is_some() { let cpu_id = need_preempt_info.unwrap(); + let preempt_guard = disable_preempt(); // FIXME: send IPI to set remote CPU's need_preempt if needed. - if cpu_id == this_cpu() { + if cpu_id == preempt_guard.current_cpu() { cpu_local::set_need_preempt(); } } @@ -163,8 +164,9 @@ pub(super) fn run_new_task(runnable: Arc) { .enqueue(runnable, EnqueueFlags::Spawn); if need_preempt_info.is_some() { let cpu_id = need_preempt_info.unwrap(); + let preempt_guard = disable_preempt(); // FIXME: send IPI to set remote CPU's need_preempt if needed. - if cpu_id == this_cpu() { + if cpu_id == preempt_guard.current_cpu() { cpu_local::set_need_preempt(); } }