Skip to main content

cpu_temp/cpu/
info.rs

1use std::num::NonZeroU16;
2
3use raw_cpuid::{CpuId as RawCpuId, TopologyType};
4
5#[cfg(target_os = "windows")]
6use windows::Win32::{Foundation::*, System::SystemInformation::*};
7
8#[derive(Debug, Clone, Copy, PartialEq)]
9pub enum Vendor {
10    Intel,
11    Amd,
12    Unknown,
13}
14
15pub struct CpuInfo {
16    vendor: Vendor,
17    family: u32,
18    model: u32,
19    stepping: u32,
20    core_count: Option<NonZeroU16>,
21}
22
23impl CpuInfo {
24    pub fn new() -> Option<Self> {
25        let cpuid = RawCpuId::new();
26
27        // 获取厂商信息
28        let vendor = match cpuid.get_vendor_info() {
29            Some(info) => match info.as_str() {
30                "GenuineIntel" => Vendor::Intel,
31                "AuthenticAMD" => Vendor::Amd,
32                _ => Vendor::Unknown,
33            },
34            None => Vendor::Unknown,
35        };
36
37        // 获取基础CPU信息
38        let (family, model, stepping) = match cpuid.get_feature_info() {
39            Some(info) => (
40                info.family_id() as u32,
41                info.model_id() as u32,
42                info.stepping_id() as u32,
43            ),
44            None => (0, 0, 0),
45        };
46
47        let core_count = cpuid
48            .get_extended_topology_info()
49            .and_then(|mut info| info.find(|level| level.level_type() == TopologyType::Core))
50            .and_then(|level| NonZeroU16::new(level.processors()));
51
52        Some(CpuInfo {
53            vendor,
54            family,
55            model,
56            stepping,
57            core_count,
58        })
59    }
60
61    pub fn get_vendor(&self) -> Vendor {
62        self.vendor
63    }
64
65    pub fn get_family(&self) -> u32 {
66        self.family
67    }
68
69    pub fn get_model(&self) -> u32 {
70        self.model
71    }
72
73    pub fn get_stepping(&self) -> u32 {
74        self.stepping
75    }
76
77    pub fn get_core_count(&self) -> Option<NonZeroU16> {
78        self.core_count
79    }
80
81    /// 检查CPU是否支持温度传感器
82    pub fn has_temperature_sensor(&self) -> bool {
83        let cpuid = RawCpuId::new();
84        if cpuid.get_feature_info().is_some() {
85            // 检查是否支持数字温度传感器
86            // 通过CPUID功能06H的EAX返回值的第0位检查
87            if let Some(thermal_info) = cpuid.get_thermal_power_info() {
88                return thermal_info.has_dts(); // 检查是否支持数字温度传感器
89            }
90        }
91        false
92    }
93
94    /// 获取CPU名称
95    pub fn get_name(&self) -> String {
96        let cpuid = RawCpuId::new();
97        if let Some(info) = cpuid.get_processor_brand_string() {
98            info.as_str().to_string()
99        } else {
100            "Unknown CPU".to_string()
101        }
102    }
103
104    /// 获取物理核心对应的逻辑CPU核心映射
105    /// 返回一个向量,每个元素代表一个物理核心,包含该核心对应的逻辑CPU核心列表
106    pub fn get_core_cpu_mapping() -> anyhow::Result<Vec<Vec<usize>>> {
107        #[cfg(target_os = "windows")]
108        {
109            Self::get_core_cpu_mapping_windows()
110        }
111        #[cfg(target_os = "linux")]
112        {
113            let cpu_info = CpuInfo::new().ok_or_else(|| anyhow::anyhow!("Failed to get CPU info"))?;
114            cpu_info.get_core_cpu_mapping_linux()
115        }
116        #[cfg(not(any(target_os = "windows", target_os = "linux")))]
117        {
118            // 默认实现,返回空向量
119            anyhow::bail!("not implementated");
120        }
121    }
122
123    /// Windows实现:获取物理核心对应的逻辑CPU核心映射
124    #[cfg(target_os = "windows")]
125    fn get_core_cpu_mapping_windows() -> anyhow::Result<Vec<Vec<usize>>> {
126        use std::mem;
127
128        let mut mappings = Vec::new();
129
130        // 获取逻辑处理器信息所需的缓冲区大小
131        let mut buffer_length: u32 = 0;
132        let result =
133            unsafe { GetLogicalProcessorInformationEx(RelationAll, None, &mut buffer_length as _) };
134
135        match result {
136            Ok(()) => return Ok(mappings),
137            Err(e) => {
138                if unsafe { GetLastError() } != ERROR_INSUFFICIENT_BUFFER {
139                    anyhow::bail!(
140                        "GetLogicalProcessorInformationEx failed (0x{:X}): {}",
141                        e.code().0,
142                        e.message()
143                    )
144                }
145            }
146        }
147
148        // 分配缓冲区
149        let mut buffer: Vec<u8> = vec![0; buffer_length as usize];
150
151        // 获取逻辑处理器信息
152        let result = unsafe {
153            GetLogicalProcessorInformationEx(
154                RelationAll,
155                Some(buffer.as_mut_ptr() as _),
156                &mut buffer_length as _,
157            )
158        };
159
160        if let Err(e) = result {
161            anyhow::bail!(
162                "GetLogicalProcessorInformationEx failed (0x{:X}): {}",
163                e.code().0,
164                e.message()
165            )
166        }
167
168        // 解析处理器关系信息
169        let mut offset: usize = 0;
170        while offset < buffer.len() {
171            use anyhow::Context;
172            let size = mem::size_of::<SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX>();
173            let cur_buffer = buffer
174                .get(offset..offset + size)
175                .context("slice info error")?;
176            let info: SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX =
177                unsafe { std::ptr::read_unaligned(cur_buffer.as_ptr() as _) };
178            if info.Relationship == RelationProcessorCore {
179                let mut logical_processors = Vec::new();
180
181                // 获取处理器关系信息
182                let processor_info = unsafe { &info.Anonymous.Processor };
183
184                let group_mask_offset = offset
185                    + mem::offset_of!(
186                        SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX,
187                        Anonymous.Processor.GroupMask
188                    );
189                let group_mask_size = processor_info.GroupCount as usize
190                    * mem::size_of_val(&processor_info.GroupMask[0]);
191                let mask_buffer = buffer
192                    .get(group_mask_offset..group_mask_offset + group_mask_size)
193                    .context("slice group mask error")?;
194
195                // 遍历组掩码
196                for i in 0..processor_info.GroupCount as usize {
197                    let group_mask = unsafe {
198                        mask_buffer
199                            .as_ptr()
200                            .cast::<GROUP_AFFINITY>()
201                            .add(i)
202                            .read_unaligned()
203                    };
204                    let mut mask = group_mask.Mask;
205
206                    // 解析掩码中的逻辑处理器
207                    while mask != 0 {
208                        // 获取最低位 1 的索引
209                        let bit = mask.trailing_zeros();
210                        logical_processors.push(i * 64 + bit as usize);
211
212                        // 清除最低位的 1
213                        mask &= mask - 1;
214                    }
215                }
216
217                mappings.push(logical_processors);
218            }
219
220            // 移动到下一个条目
221            offset += info.Size as usize;
222        }
223
224        Ok(mappings)
225    }
226
227    /// Linux实现:获取物理核心对应的逻辑CPU核心映射
228    #[cfg(target_os = "linux")]
229    fn get_core_cpu_mapping_linux(&self) -> anyhow::Result<Vec<Vec<usize>>> {
230        use std::collections::BTreeMap;
231        use std::fs;
232
233        use anyhow::Context as _;
234
235        const CORE_ID_PATTERN: &str = "/sys/devices/system/cpu/cpu[0-9]*/topology/core_id";
236
237        let mut groups: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
238        for entry in glob::glob(CORE_ID_PATTERN).context("Failed to read glob pattern")? {
239            if let Ok(path) = entry {
240                // 获取目录名中的 cpu 号,例如从 ".../cpu3/topology/core_id" 提取 3
241                let cpu_id = path
242                    .ancestors()
243                    .nth(2)
244                    .and_then(|p| p.file_name())
245                    .and_then(|s| s.to_str())
246                    .and_then(|s| s.trim_start_matches("cpu").parse::<usize>().ok());
247
248                // 读取 core_id 文件内容
249                let core_id = fs::read_to_string(&path)
250                    .ok()
251                    .and_then(|s| s.trim().parse::<usize>().ok());
252
253                if let (Some(cpu), Some(core)) = (cpu_id, core_id) {
254                    groups.entry(core).or_default().push(cpu);
255                }
256            }
257        }
258
259        // 将分组后的结果转为 Vec<Vec<usize>>
260        Ok(groups.into_values().collect())
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    #[test]
269    fn test_get_core_cpu_mapping() {
270        // use windows::Win32::System::Threading::*;
271        // let mut groupcount = 0u16;
272        // let handle = unsafe { GetCurrentProcess() };
273        // let _ =
274        //     unsafe { GetProcessGroupAffinity(handle, &mut groupcount as _, Default::default()) };
275        // let grouparray: Vec<u16> = vec![0; groupcount as usize];
276        // let res = unsafe {
277        //     GetProcessGroupAffinity(handle, &mut groupcount as _, grouparray.as_ptr() as _)
278        // };
279        // if !res.as_bool() {
280        //     panic!("failed: {}", unsafe {
281        //         &GetLastError().to_hresult().message()
282        //     });
283        // }
284        let res = CpuInfo::get_core_cpu_mapping().unwrap();
285        println!("{:?}", res);
286    }
287}