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