1 use core::sync::atomic::AtomicU32; 2 3 use alloc::vec::Vec; 4 5 use crate::{ 6 libs::lazy_init::Lazy, 7 smp::{ 8 core::smp_get_processor_id, 9 cpu::{smp_cpu_manager, ProcessorId}, 10 }, 11 }; 12 13 /// 系统中的CPU数量 14 /// 15 /// todo: 待smp模块重构后，从smp模块获取CPU数量。 16 /// 目前由于smp模块初始化时机较晚，导致大部分内核模块无法在早期初始化PerCpu变量。 17 static CPU_NUM_ATOMIC: AtomicU32 = AtomicU32::new(PerCpu::MAX_CPU_NUM); 18 19 #[derive(Debug)] 20 pub struct PerCpu; 21 22 impl PerCpu { 23 #[cfg(target_arch = "x86_64")] 24 pub const MAX_CPU_NUM: u32 = 128; 25 #[cfg(target_arch = "riscv64")] 26 pub const MAX_CPU_NUM: u32 = 64; 27 28 /// # 初始化PerCpu 29 /// 30 /// 该函数应该在内核初始化时调用一次。 31 /// 32 /// 该函数会调用`smp_get_total_cpu()`获取CPU数量，然后将其存储在`CPU_NUM`中。 33 #[allow(dead_code)] 34 pub fn init() { 35 let cpu_num: &AtomicU32 = &CPU_NUM_ATOMIC; 36 if cpu_num.load(core::sync::atomic::Ordering::SeqCst) != 0 { 37 panic!("PerCpu::init() called twice"); 38 } 39 let cpus = smp_cpu_manager().present_cpus_count(); 40 assert!(cpus > 0, "PerCpu::init(): present_cpus_count() returned 0"); 41 42 CPU_NUM_ATOMIC.store(cpus, core::sync::atomic::Ordering::SeqCst); 43 } 44 } 45 46 /// PerCpu变量 47 /// 48 /// 该结构体的每个实例都是线程安全的，因为每个CPU都有自己的变量。 49 /// 50 /// 一种简单的使用方法是：使用该结构体提供的`define_lazy`方法定义一个全局变量， 51 /// 然后在内核初始化时调用`init`、`new`方法去初始化它。 52 /// 53 /// 当然，由于Lazy<T>有运行时开销，所以也可以直接全局声明一个Option， 54 /// 然后手动初始化然后赋值到Option中。（这样需要在初始化的时候，手动确保并发安全） 55 #[derive(Debug)] 56 #[allow(dead_code)] 57 pub struct PerCpuVar<T> { 58 inner: Vec<T>, 59 } 60 61 #[allow(dead_code)] 62 impl<T> PerCpuVar<T> { 63 /// # 初始化PerCpu变量 64 /// 65 /// ## 参数 66 /// 67 /// - `data` - 每个CPU的数据的初始值。 传入的Vec的长度必须等于CPU的数量，否则返回None。 68 pub fn new(data: Vec<T>) -> Option<Self> { 69 let cpu_num = CPU_NUM_ATOMIC.load(core::sync::atomic::Ordering::SeqCst); 70 if cpu_num == 0 { 71 panic!("PerCpu::init() not called"); 72 } 73 74 if data.len() != cpu_num.try_into().unwrap() { 75 return None; 76 } 77 78 return Some(Self { inner: data }); 79 } 80 81 /// 定义一个Lazy的PerCpu变量，稍后再初始化 82 pub const fn define_lazy() -> Lazy<Self> { 83 Lazy::<Self>::new() 84 } 85 86 pub fn get(&self) -> &T { 87 let cpu_id = smp_get_processor_id(); 88 &self.inner[cpu_id.data() as usize] 89 } 90 91 #[allow(clippy::mut_from_ref)] 92 pub fn get_mut(&self) -> &mut T { 93 let cpu_id = smp_get_processor_id(); 94 unsafe { 95 &mut (self as *const Self as *mut Self).as_mut().unwrap().inner[cpu_id.data() as usize] 96 } 97 } 98 99 pub unsafe fn force_get(&self, cpu_id: ProcessorId) -> &T { 100 &self.inner[cpu_id.data() as usize] 101 } 102 103 #[allow(clippy::mut_from_ref)] 104 pub unsafe fn force_get_mut(&self, cpu_id: ProcessorId) -> &mut T { 105 &mut (self as *const Self as *mut Self).as_mut().unwrap().inner[cpu_id.data() as usize] 106 } 107 } 108 109 /// PerCpu变量是线程安全的，因为每个CPU都有自己的变量。 110 unsafe impl<T> Sync for PerCpuVar<T> {} 111 unsafe impl<T> Send for PerCpuVar<T> {} 112