starnix_core/task/scheduler/
role_overrides.rs

1// Copyright 2025 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use bstr::ByteSlice;
6use regex_lite::{Error, Regex};
7use starnix_task_command::TaskCommand;
8
9/// Per-container overrides for thread roles.
10#[derive(Debug)]
11pub struct RoleOverrides {
12    process_filter: Vec<Regex>,
13    thread_filter: Vec<Regex>,
14    cgroup_filter: Vec<Option<Regex>>,
15    is_leader_filter: Vec<Option<bool>>,
16    min_nice_filter: Vec<Option<i8>>,
17    max_nice_filter: Vec<Option<i8>>,
18    role_names: Vec<String>,
19}
20
21impl RoleOverrides {
22    /// Create a new builder for role overrides.
23    pub fn new() -> RoleOverridesBuilder {
24        RoleOverridesBuilder {
25            process_patterns: vec![],
26            thread_patterns: vec![],
27            cgroup_patterns: vec![],
28            is_leader_filters: vec![],
29            min_nice_filters: vec![],
30            max_nice_filters: vec![],
31            role_names: vec![],
32        }
33    }
34
35    /// Get the overridden role name (if any) for provided process and thread names.
36    pub fn get_role_name<'a>(
37        &self,
38        process_name: &TaskCommand,
39        thread_name: &TaskCommand,
40        cgroup_path: &str,
41    ) -> Option<&str> {
42        self.get_role_name_with_filters(process_name, thread_name, cgroup_path, false, 0)
43    }
44
45    /// Get the overridden role name (if any) for provided process, thread, cgroup, leadership, and
46    /// nice attributes.
47    pub fn get_role_name_with_filters<'a>(
48        &self,
49        process_name: &TaskCommand,
50        thread_name: &TaskCommand,
51        cgroup_path: &str,
52        is_leader: bool,
53        nice: i8,
54    ) -> Option<&str> {
55        debug_assert_eq!(self.process_filter.len(), self.role_names.len());
56        debug_assert_eq!(self.thread_filter.len(), self.role_names.len());
57        debug_assert_eq!(self.cgroup_filter.len(), self.role_names.len());
58        debug_assert_eq!(self.is_leader_filter.len(), self.role_names.len());
59        debug_assert_eq!(self.min_nice_filter.len(), self.role_names.len());
60        debug_assert_eq!(self.max_nice_filter.len(), self.role_names.len());
61
62        // NOTE(https://fxbug.dev/483609435): This used to be more elegantly expressed
63        // via use of regex::bytes::RegexSet, but regex_lite doesn't (yet?) offer RegexSet.
64        let process_name = process_name.as_bytes().to_str().ok()?;
65        let thread_name = thread_name.as_bytes().to_str().ok()?;
66        for index in 0..self.process_filter.len() {
67            if self.is_leader_filter[index].is_none_or(|expected| expected == is_leader)
68                && self.min_nice_filter[index].is_none_or(|min| nice >= min)
69                && self.max_nice_filter[index].is_none_or(|max| nice <= max)
70                && self.process_filter[index].is_match(process_name)
71                && self.thread_filter[index].is_match(thread_name)
72                && self.cgroup_filter[index].as_ref().is_none_or(|r| r.is_match(cgroup_path))
73            {
74                return Some(self.role_names[index].as_str());
75            }
76        }
77        None
78    }
79}
80
81/// Builder for `RoleOverrides`.
82pub struct RoleOverridesBuilder {
83    process_patterns: Vec<String>,
84    thread_patterns: Vec<String>,
85    cgroup_patterns: Vec<Option<String>>,
86    is_leader_filters: Vec<Option<bool>>,
87    min_nice_filters: Vec<Option<i8>>,
88    max_nice_filters: Vec<Option<i8>>,
89    role_names: Vec<String>,
90}
91
92impl RoleOverridesBuilder {
93    /// Add a new override to the configuration.
94    pub fn add(
95        &mut self,
96        process: impl Into<String>,
97        thread: impl Into<String>,
98        cgroup: Option<String>,
99        role_name: impl Into<String>,
100    ) {
101        self.add_with_filters(process, thread, cgroup, None, None, None, role_name);
102    }
103
104    /// Add a new override with optional thread-group leader and nice range filters.
105    pub fn add_with_filters(
106        &mut self,
107        process: impl Into<String>,
108        thread: impl Into<String>,
109        cgroup: Option<String>,
110        is_leader: Option<bool>,
111        min_nice: Option<i8>,
112        max_nice: Option<i8>,
113        role_name: impl Into<String>,
114    ) {
115        self.process_patterns.push(process.into());
116        self.thread_patterns.push(thread.into());
117        self.cgroup_patterns.push(cgroup);
118        self.is_leader_filters.push(is_leader);
119        self.min_nice_filters.push(min_nice);
120        self.max_nice_filters.push(max_nice);
121        self.role_names.push(role_name.into());
122    }
123
124    /// Compile all of the provided regular expressions and return a `RoleOverrides`.
125    pub fn build(self) -> Result<RoleOverrides, Error> {
126        let cgroup_filter = self
127            .cgroup_patterns
128            .into_iter()
129            .map(|opt| opt.map(|p| Regex::new(p.as_str())).transpose())
130            .collect::<Result<Vec<Option<Regex>>, Error>>()?;
131
132        Ok(RoleOverrides {
133            process_filter: self
134                .process_patterns
135                .into_iter()
136                .map(|pattern| Regex::new(pattern.as_str()))
137                .collect::<Result<Vec<Regex>, Error>>()?,
138            thread_filter: self
139                .thread_patterns
140                .into_iter()
141                .map(|pattern| Regex::new(pattern.as_str()))
142                .collect::<Result<Vec<Regex>, Error>>()?,
143            cgroup_filter,
144            is_leader_filter: self.is_leader_filters,
145            min_nice_filter: self.min_nice_filters,
146            max_nice_filter: self.max_nice_filters,
147            role_names: self.role_names,
148        })
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155
156    fn str_role_name<'a>(
157        mappings: &'a RoleOverrides,
158        process_name: &str,
159        thread_name: &str,
160        cpuset_path: &str,
161    ) -> Option<&'a str> {
162        mappings.get_role_name(
163            &TaskCommand::new(process_name.as_bytes()),
164            &TaskCommand::new(thread_name.as_bytes()),
165            cpuset_path,
166        )
167    }
168
169    #[fuchsia::test]
170    fn single_pattern() {
171        let mut builder = RoleOverrides::new();
172        builder.add("process_prefix_.+", "thread_prefix_.+", None, "replacement_role");
173        let mappings = builder.build().unwrap();
174
175        assert_eq!(
176            str_role_name(&mappings, "process_prefix_foo", "thread_prefix_bar", "/"),
177            Some("replacement_role")
178        );
179        assert_eq!(str_role_name(&mappings, "process_prefix_foo", "non_matching", "/"), None);
180        assert_eq!(str_role_name(&mappings, "non_matching", "process_prefix_bar", "/"), None);
181        assert_eq!(str_role_name(&mappings, "non_matching", "non_matching", "/"), None);
182    }
183
184    #[fuchsia::test]
185    fn multiple_patterns() {
186        let mut builder = RoleOverrides::new();
187        builder.add("pre_one.+", "pre_one.+", None, "replace_one");
188        builder.add("pre_two.+", "pre_two.+", None, "replace_two");
189        builder.add("pre_three.+", "pre_three.+", None, "replace_three");
190        builder.add("pre_four.+", "pre_four.+", None, "replace_four");
191        let mappings = builder.build().unwrap();
192
193        assert_eq!(
194            str_role_name(&mappings, "pre_one_foo", "pre_one_bar", "/"),
195            Some("replace_one")
196        );
197        assert_eq!(str_role_name(&mappings, "pre_one_foo", "non_matching", "/"), None);
198        assert_eq!(str_role_name(&mappings, "non_matching", "pre_one_bar", "/"), None);
199        assert_eq!(str_role_name(&mappings, "non_matching", "non_matching", "/"), None);
200
201        assert_eq!(
202            str_role_name(&mappings, "pre_two_foo", "pre_two_bar", "/"),
203            Some("replace_two")
204        );
205        assert_eq!(str_role_name(&mappings, "pre_two_foo", "non_matching", "/"), None);
206        assert_eq!(str_role_name(&mappings, "non_matching", "pre_two_bar", "/"), None);
207
208        assert_eq!(
209            str_role_name(&mappings, "pre_three_foo", "pre_three_bar", "/"),
210            Some("replace_three")
211        );
212        assert_eq!(str_role_name(&mappings, "pre_three_foo", "non_matching", "/"), None);
213        assert_eq!(str_role_name(&mappings, "non_matching", "pre_three_bar", "/"), None);
214
215        assert_eq!(
216            str_role_name(&mappings, "pre_four_foo", "pre_four_bar", "/"),
217            Some("replace_four")
218        );
219        assert_eq!(str_role_name(&mappings, "pre_four_foo", "non_matching", "/"), None);
220        assert_eq!(str_role_name(&mappings, "non_matching", "pre_four_bar", "/"), None);
221    }
222
223    #[fuchsia::test]
224    fn cgroup_patterns() {
225        let mut builder = RoleOverrides::new();
226        builder.add("proc", "thread", Some("/background".to_string()), "bg_role");
227        builder.add("proc", "thread", Some("/foreground".to_string()), "fg_role");
228        builder.add("proc", "thread", None, "default_role");
229        let mappings = builder.build().unwrap();
230
231        assert_eq!(str_role_name(&mappings, "proc", "thread", "/background"), Some("bg_role"));
232        assert_eq!(str_role_name(&mappings, "proc", "thread", "/foreground"), Some("fg_role"));
233        assert_eq!(str_role_name(&mappings, "proc", "thread", "/other"), Some("default_role"));
234        assert_eq!(str_role_name(&mappings, "proc", "thread", "/"), Some("default_role"));
235    }
236
237    #[fuchsia::test]
238    fn leader_and_nice_filters() {
239        let mut builder = RoleOverrides::new();
240        builder.add_with_filters(
241            ".+",
242            ".+",
243            Some("^/active$".to_string()),
244            Some(true),
245            None,
246            None,
247            "leader_role",
248        );
249        builder.add_with_filters(
250            ".+",
251            ".+",
252            Some("^/active$".to_string()),
253            None,
254            Some(1),
255            None,
256            "low_prio_worker_role",
257        );
258        builder.add_with_filters(
259            ".+",
260            ".+",
261            Some("^/active$".to_string()),
262            None,
263            None,
264            Some(-1),
265            "boosted_worker_role",
266        );
267        builder.add(".+", ".+", Some("^/active$".to_string()), "default_worker_role");
268        let mappings = builder.build().unwrap();
269
270        let proc_cmd = TaskCommand::new(b"app");
271        let worker_cmd = TaskCommand::new(b"worker");
272
273        assert_eq!(
274            mappings.get_role_name_with_filters(&proc_cmd, &proc_cmd, "/active", true, 0),
275            Some("leader_role")
276        );
277        assert_eq!(
278            mappings.get_role_name_with_filters(&proc_cmd, &worker_cmd, "/active", false, 0),
279            Some("default_worker_role")
280        );
281        assert_eq!(
282            mappings.get_role_name_with_filters(&proc_cmd, &worker_cmd, "/active", false, 10),
283            Some("low_prio_worker_role")
284        );
285        assert_eq!(
286            mappings.get_role_name_with_filters(&proc_cmd, &worker_cmd, "/active", false, -2),
287            Some("boosted_worker_role")
288        );
289    }
290}
starnix_core/task/scheduler/role_overrides.rs

starnix_core/task/scheduler/
role_overrides.rs
