cgroups_explorer/
explorer.rs

1use std::{
2    collections::{HashSet, hash_set},
3    path::{Path, PathBuf},
4};
5
6use cgroups_rs::{
7    Cgroup, Hierarchy,
8    hierarchies::{V1, V2, is_cgroup2_unified_mode},
9};
10use derive_builder::Builder;
11use walkdir::WalkDir;
12
13/// An interface to explore cgroups in the system.
14///
15/// # Example
16/// ```rust
17/// use cgroups_explorer::Explorer;
18/// let explorer = Explorer::detect_version()
19///     .include(vec!["user.slice/*".to_string()])
20///     .build()
21///     .expect("Failed to build explorer");
22/// let found = explorer
23///     .iter_cgroups()
24///     .for_each(|c| println!("Found cgroup: {}", c.path()));
25///
26/// ```
27#[derive(Builder)]
28#[builder(pattern = "owned")]
29pub struct Explorer {
30    /// The cgroup hierarchy to explore.
31    hierarchy: Box<dyn Hierarchy>,
32
33    /// The globs to include in the exploration.
34    #[builder(field(ty = "Vec<String>", build = "parse_include(self.include)?"))]
35    include: Vec<glob::Pattern>,
36}
37
38/// An iterator over cgroups in the system that match the globs.
39struct CgroupsV2Iterator {
40    walker: walkdir::IntoIter,
41    include: Vec<glob::Pattern>,
42    base_path: PathBuf,
43}
44
45struct CgroupsV1Iterator {
46    discovered: hash_set::IntoIter<PathBuf>,
47}
48
49impl Explorer {
50    /// Create a new `ExplorerBuilder` for cgroups v1.
51    #[must_use]
52    pub fn v1() -> ExplorerBuilder {
53        ExplorerBuilder::default().hierarchy(Box::new(V1::new()))
54    }
55
56    /// Create a new `ExplorerBuilder` for cgroups v2.
57    #[must_use]
58    pub fn v2() -> ExplorerBuilder {
59        ExplorerBuilder::default().hierarchy(Box::new(V2::new()))
60    }
61
62    /// Create a new `ExplorerBuilder` by detecting the cgroups version on the system.
63    #[must_use]
64    pub fn detect_version() -> ExplorerBuilder {
65        if is_cgroup2_unified_mode() {
66            ExplorerBuilder::default().hierarchy(Box::new(V2::new()))
67        } else {
68            ExplorerBuilder::default().hierarchy(Box::new(V1::new()))
69        }
70    }
71
72    /// Create an iterator over all cgroups in the system, based on the criteria.
73    #[must_use]
74    pub fn iter_cgroups(&self) -> Box<dyn Iterator<Item = Cgroup>> {
75        if self.hierarchy.v2() {
76            Box::new(self.iter_cgroups_v2())
77        } else {
78            Box::new(self.iter_cgroups_v1())
79        }
80    }
81
82    fn iter_cgroups_v2(&self) -> CgroupsV2Iterator {
83        let base_path = self.hierarchy.root();
84        let walker = WalkDir::new(base_path.clone())
85            .min_depth(1)
86            .sort_by_file_name()
87            .into_iter();
88        CgroupsV2Iterator {
89            walker,
90            include: self.include.clone(),
91            base_path,
92        }
93    }
94
95    fn iter_cgroups_v1(&self) -> CgroupsV1Iterator {
96        let hierarchy = V1::new();
97        let subystems = hierarchy.subsystems();
98        let base_path = hierarchy.root();
99
100        let mut matching_rel_paths = HashSet::new();
101        for subsystem in subystems {
102            let name = subsystem.controller_name();
103            let walker = WalkDir::new(base_path.join(&name))
104                .min_depth(1)
105                .sort_by_file_name()
106                .into_iter();
107            let base_controller_path = base_path.join(name);
108            for entry in walker {
109                let Ok(entry) = entry else { continue };
110                let path = entry.path();
111                if !entry.file_type().is_dir() {
112                    continue;
113                }
114                let Ok(relative_path) = path.strip_prefix(&base_controller_path) else {
115                    continue;
116                };
117                if relative_path.components().count() == 0 {
118                    continue;
119                }
120                if self.include.is_empty() || path_matches_include(&self.include, relative_path) {
121                    matching_rel_paths.insert(relative_path.to_path_buf());
122                }
123            }
124        }
125
126        CgroupsV1Iterator {
127            discovered: matching_rel_paths.into_iter(),
128        }
129    }
130}
131
132impl Iterator for CgroupsV2Iterator {
133    type Item = Cgroup;
134
135    fn next(&mut self) -> Option<Self::Item> {
136        loop {
137            let entry = self.walker.next();
138            match entry {
139                Some(Ok(entry)) => {
140                    let path = entry.path();
141                    if !entry.file_type().is_dir() {
142                        continue;
143                    }
144                    let Ok(relative_path) = path.strip_prefix(&self.base_path) else {
145                        continue;
146                    };
147                    if relative_path.components().count() == 0 {
148                        continue;
149                    }
150                    if !path_matches_include(&self.include, relative_path) {
151                        continue;
152                    }
153                    return Some(Cgroup::load(Box::new(V2::new()), relative_path));
154                }
155                Some(Err(_e)) => return None,
156                None => return None,
157            }
158        }
159    }
160}
161
162impl Iterator for CgroupsV1Iterator {
163    type Item = Cgroup;
164
165    fn next(&mut self) -> Option<Self::Item> {
166        self.discovered
167            .next()
168            .map(|path| Cgroup::load(Box::new(V1::new()), path))
169    }
170}
171
172fn path_matches_include(include: &[glob::Pattern], path: &Path) -> bool {
173    if include.is_empty() {
174        return true;
175    }
176    let path_str = path.to_string_lossy();
177    include.iter().any(|pattern| pattern.matches(&path_str))
178}
179
180fn parse_include(include: Vec<String>) -> Result<Vec<glob::Pattern>, ExplorerBuilderError> {
181    if include.is_empty() {
182        Ok(Vec::new())
183    } else {
184        include
185            .into_iter()
186            .map(|include| {
187                glob::Pattern::new(&include)
188                    .map_err(|e| ExplorerBuilderError::ValidationError(e.to_string()))
189            })
190            .collect()
191    }
192}
cgroups_explorer/explorer.rs

cgroups_explorer/
explorer.rs
