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    role_names: Vec<String>,
16}
17
18impl RoleOverrides {
19    /// Create a new builder for role overrides.
20    pub fn new() -> RoleOverridesBuilder {
21        RoleOverridesBuilder {
22            process_patterns: vec![],
23            thread_patterns: vec![],
24            cgroup_patterns: vec![],
25            role_names: vec![],
26        }
27    }
28
29    /// Get the overridden role name (if any) for provided process and thread names.
30    pub fn get_role_name<'a>(
31        &self,
32        process_name: &TaskCommand,
33        thread_name: &TaskCommand,
34        cgroup_path: &str,
35    ) -> Option<&str> {
36        debug_assert_eq!(self.process_filter.len(), self.role_names.len());
37        debug_assert_eq!(self.thread_filter.len(), self.role_names.len());
38        debug_assert_eq!(self.cgroup_filter.len(), self.role_names.len());
39
40        // NOTE(https://fxbug.dev/483609435): This used to be more elegantly expressed
41        // via use of regex::bytes::RegexSet, but regex_lite doesn't (yet?) offer RegexSet.
42        let process_name = process_name.as_bytes().to_str().ok()?;
43        let thread_name = thread_name.as_bytes().to_str().ok()?;
44        for index in 0..self.process_filter.len() {
45            if self.process_filter[index].is_match(process_name)
46                && self.thread_filter[index].is_match(thread_name)
47                && self.cgroup_filter[index].as_ref().map_or(true, |r| r.is_match(cgroup_path))
48            {
49                return Some(self.role_names[index].as_str());
50            }
51        }
52        None
53    }
54}
55
56/// Builder for `RoleOverrides`.
57pub struct RoleOverridesBuilder {
58    process_patterns: Vec<String>,
59    thread_patterns: Vec<String>,
60    cgroup_patterns: Vec<Option<String>>,
61    role_names: Vec<String>,
62}
63
64impl RoleOverridesBuilder {
65    /// Add a new override to the configuration.
66    pub fn add(
67        &mut self,
68        process: impl Into<String>,
69        thread: impl Into<String>,
70        cgroup: Option<String>,
71        role_name: impl Into<String>,
72    ) {
73        self.process_patterns.push(process.into());
74        self.thread_patterns.push(thread.into());
75        self.cgroup_patterns.push(cgroup);
76        self.role_names.push(role_name.into());
77    }
78
79    /// Compile all of the provided regular expressions and return a `RoleOverrides`.
80    pub fn build(self) -> Result<RoleOverrides, Error> {
81        let cgroup_filter = self
82            .cgroup_patterns
83            .into_iter()
84            .map(|opt| opt.map(|p| Regex::new(p.as_str())).transpose())
85            .collect::<Result<Vec<Option<Regex>>, Error>>()?;
86
87        Ok(RoleOverrides {
88            process_filter: self
89                .process_patterns
90                .into_iter()
91                .map(|pattern| Regex::new(pattern.as_str()))
92                .collect::<Result<Vec<Regex>, Error>>()?,
93            thread_filter: self
94                .thread_patterns
95                .into_iter()
96                .map(|pattern| Regex::new(pattern.as_str()))
97                .collect::<Result<Vec<Regex>, Error>>()?,
98            cgroup_filter,
99            role_names: self.role_names,
100        })
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    fn str_role_name<'a>(
109        mappings: &'a RoleOverrides,
110        process_name: &str,
111        thread_name: &str,
112        cpuset_path: &str,
113    ) -> Option<&'a str> {
114        mappings.get_role_name(
115            &TaskCommand::new(process_name.as_bytes()),
116            &TaskCommand::new(thread_name.as_bytes()),
117            cpuset_path,
118        )
119    }
120
121    #[fuchsia::test]
122    fn single_pattern() {
123        let mut builder = RoleOverrides::new();
124        builder.add("process_prefix_.+", "thread_prefix_.+", None, "replacement_role");
125        let mappings = builder.build().unwrap();
126
127        assert_eq!(
128            str_role_name(&mappings, "process_prefix_foo", "thread_prefix_bar", "/"),
129            Some("replacement_role")
130        );
131        assert_eq!(str_role_name(&mappings, "process_prefix_foo", "non_matching", "/"), None);
132        assert_eq!(str_role_name(&mappings, "non_matching", "process_prefix_bar", "/"), None);
133        assert_eq!(str_role_name(&mappings, "non_matching", "non_matching", "/"), None);
134    }
135
136    #[fuchsia::test]
137    fn multiple_patterns() {
138        let mut builder = RoleOverrides::new();
139        builder.add("pre_one.+", "pre_one.+", None, "replace_one");
140        builder.add("pre_two.+", "pre_two.+", None, "replace_two");
141        builder.add("pre_three.+", "pre_three.+", None, "replace_three");
142        builder.add("pre_four.+", "pre_four.+", None, "replace_four");
143        let mappings = builder.build().unwrap();
144
145        assert_eq!(
146            str_role_name(&mappings, "pre_one_foo", "pre_one_bar", "/"),
147            Some("replace_one")
148        );
149        assert_eq!(str_role_name(&mappings, "pre_one_foo", "non_matching", "/"), None);
150        assert_eq!(str_role_name(&mappings, "non_matching", "pre_one_bar", "/"), None);
151        assert_eq!(str_role_name(&mappings, "non_matching", "non_matching", "/"), None);
152
153        assert_eq!(
154            str_role_name(&mappings, "pre_two_foo", "pre_two_bar", "/"),
155            Some("replace_two")
156        );
157        assert_eq!(str_role_name(&mappings, "pre_two_foo", "non_matching", "/"), None);
158        assert_eq!(str_role_name(&mappings, "non_matching", "pre_two_bar", "/"), None);
159
160        assert_eq!(
161            str_role_name(&mappings, "pre_three_foo", "pre_three_bar", "/"),
162            Some("replace_three")
163        );
164        assert_eq!(str_role_name(&mappings, "pre_three_foo", "non_matching", "/"), None);
165        assert_eq!(str_role_name(&mappings, "non_matching", "pre_three_bar", "/"), None);
166
167        assert_eq!(
168            str_role_name(&mappings, "pre_four_foo", "pre_four_bar", "/"),
169            Some("replace_four")
170        );
171        assert_eq!(str_role_name(&mappings, "pre_four_foo", "non_matching", "/"), None);
172        assert_eq!(str_role_name(&mappings, "non_matching", "pre_four_bar", "/"), None);
173    }
174
175    #[fuchsia::test]
176    fn cgroup_patterns() {
177        let mut builder = RoleOverrides::new();
178        builder.add("proc", "thread", Some("/background".to_string()), "bg_role");
179        builder.add("proc", "thread", Some("/foreground".to_string()), "fg_role");
180        builder.add("proc", "thread", None, "default_role");
181        let mappings = builder.build().unwrap();
182
183        assert_eq!(str_role_name(&mappings, "proc", "thread", "/background"), Some("bg_role"));
184        assert_eq!(str_role_name(&mappings, "proc", "thread", "/foreground"), Some("fg_role"));
185        assert_eq!(str_role_name(&mappings, "proc", "thread", "/other"), Some("default_role"));
186        assert_eq!(str_role_name(&mappings, "proc", "thread", "/"), Some("default_role"));
187    }
188}
starnix_core/task/scheduler/role_overrides.rs

starnix_core/task/scheduler/
role_overrides.rs
