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