cgroups_explorer/
explorer.rs

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

cgroups_explorer/
explorer.rs
