cpu_utils/
cpu_topology.rs

1/*
2 * Copyright 2024 Fluence DAO
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17pub use nonempty::NonEmpty;
18
19use crate::errors::CPUTopologyError;
20use crate::CTResult;
21use crate::LogicalCoreId;
22use crate::PhysicalCoreId;
23
24#[cfg_attr(feature = "mockall", mockall::automock)]
25pub trait CPUTopology {
26    fn physical_cores(&self) -> CTResult<NonEmpty<PhysicalCoreId>>;
27
28    fn logical_cores_for_physical(
29        &self,
30        core_id: PhysicalCoreId,
31    ) -> CTResult<NonEmpty<LogicalCoreId>>;
32}
33
34#[derive(Debug)]
35pub struct HwlocCPUTopology {
36    topology: hwlocality::Topology,
37}
38
39impl HwlocCPUTopology {
40    pub fn new() -> CTResult<Self> {
41        let topology = hwlocality::Topology::new()?;
42        Ok(Self { topology })
43    }
44
45    pub fn physical_cores_count_len(&self) -> usize {
46        self.physical_cores().map(|r| r.len()).unwrap_or(0)
47    }
48}
49
50impl CPUTopology for HwlocCPUTopology {
51    fn physical_cores(&self) -> CTResult<NonEmpty<PhysicalCoreId>> {
52        use hwlocality::object::types::ObjectType;
53
54        let physical_core_ids = self
55            .topology
56            .objects_with_type(ObjectType::Core)
57            .map(|value| PhysicalCoreId::from(value.logical_index() as u32))
58            .collect::<Vec<_>>();
59
60        NonEmpty::from_vec(physical_core_ids).ok_or_else(|| CPUTopologyError::PhysicalCoresNotFound)
61    }
62
63    fn logical_cores_for_physical(
64        &self,
65        core_id: PhysicalCoreId,
66    ) -> CTResult<NonEmpty<LogicalCoreId>> {
67        use hwlocality::object::types::ObjectType;
68
69        let core_depth = self.topology.depth_or_below_for_type(ObjectType::Core)?;
70        let physical_cores = self
71            .topology
72            .objects_at_depth(core_depth)
73            .collect::<Vec<_>>();
74
75        let physical_core = physical_cores
76            .get(<PhysicalCoreId as Into<usize>>::into(core_id))
77            .ok_or(CPUTopologyError::physical_core_not_found(core_id))?;
78
79        let physical_core_cpuset = physical_core
80            .cpuset()
81            .ok_or(CPUTopologyError::cpuset_not_found(core_id))?;
82
83        let logical_core_ids = physical_core_cpuset
84            .into_iter()
85            .map(usize::from)
86            .map(|value| LogicalCoreId::from(value as u32))
87            .collect::<Vec<_>>();
88
89        NonEmpty::from_vec(logical_core_ids)
90            .ok_or_else(|| CPUTopologyError::logical_cores_not_found(core_id))
91    }
92}