Skip to main content

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