xref: /DragonOS/kernel/src/mm/percpu.rs (revision 2eab6dd743e94a86a685f1f3c01e599adf86610a)
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