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 itertools::Itertools as _;
6use regex::Error;
7use regex::bytes::RegexSet;
8use starnix_task_command::TaskCommand;
9
10/// Per-container overrides for thread roles.
11#[derive(Debug)]
12pub struct RoleOverrides {
13    process_filter: RegexSet,
14    thread_filter: RegexSet,
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            role_names: vec![],
25        }
26    }
27
28    /// Get the overridden role name (if any) for provided process and thread names.
29    pub fn get_role_name<'a>(
30        &self,
31        process_name: &TaskCommand,
32        thread_name: &TaskCommand,
33    ) -> Option<&str> {
34        debug_assert_eq!(self.process_filter.len(), self.role_names.len());
35        debug_assert_eq!(self.thread_filter.len(), self.role_names.len());
36
37        // RegexSet iterators are sorted, we can assume they'll be in ascending order.
38        let process_match_indices =
39            self.process_filter.matches(process_name.as_bytes()).into_iter();
40        let thread_match_indices = self.thread_filter.matches(thread_name.as_bytes()).into_iter();
41        let mut indices_with_both_matching =
42            process_match_indices.merge(thread_match_indices).duplicates();
43        indices_with_both_matching.next().map(|i| self.role_names[i].as_str())
44    }
45}
46
47/// Builder for `RoleOverrides`.
48pub struct RoleOverridesBuilder {
49    process_patterns: Vec<String>,
50    thread_patterns: Vec<String>,
51    role_names: Vec<String>,
52}
53
54impl RoleOverridesBuilder {
55    /// Add a new override to the configuration.
56    pub fn add(
57        &mut self,
58        process: impl Into<String>,
59        thread: impl Into<String>,
60        role_name: impl Into<String>,
61    ) {
62        self.process_patterns.push(process.into());
63        self.thread_patterns.push(thread.into());
64        self.role_names.push(role_name.into());
65    }
66
67    /// Compile all of the provided regular expressions and return a `RoleOverrides`.
68    pub fn build(self) -> Result<RoleOverrides, Error> {
69        Ok(RoleOverrides {
70            process_filter: RegexSet::new(self.process_patterns)?,
71            thread_filter: RegexSet::new(self.thread_patterns)?,
72            role_names: self.role_names,
73        })
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    fn str_role_name<'a>(
82        mappings: &'a RoleOverrides,
83        process_name: &str,
84        thread_name: &str,
85    ) -> Option<&'a str> {
86        mappings.get_role_name(
87            &TaskCommand::new(process_name.as_bytes()),
88            &TaskCommand::new(thread_name.as_bytes()),
89        )
90    }
91
92    #[fuchsia::test]
93    fn single_pattern() {
94        let mut builder = RoleOverrides::new();
95        builder.add("process_prefix_.+", "thread_prefix_.+", "replacement_role");
96        let mappings = builder.build().unwrap();
97
98        assert_eq!(
99            str_role_name(&mappings, "process_prefix_foo", "thread_prefix_bar"),
100            Some("replacement_role")
101        );
102        assert_eq!(str_role_name(&mappings, "process_prefix_foo", "non_matching"), None);
103        assert_eq!(str_role_name(&mappings, "non_matching", "process_prefix_bar"), None);
104        assert_eq!(str_role_name(&mappings, "non_matching", "non_matching"), None);
105    }
106
107    #[fuchsia::test]
108    fn multiple_patterns() {
109        let mut builder = RoleOverrides::new();
110        builder.add("pre_one.+", "pre_one.+", "replace_one");
111        builder.add("pre_two.+", "pre_two.+", "replace_two");
112        builder.add("pre_three.+", "pre_three.+", "replace_three");
113        builder.add("pre_four.+", "pre_four.+", "replace_four");
114        let mappings = builder.build().unwrap();
115
116        assert_eq!(str_role_name(&mappings, "pre_one_foo", "pre_one_bar"), Some("replace_one"));
117        assert_eq!(str_role_name(&mappings, "pre_one_foo", "non_matching"), None);
118        assert_eq!(str_role_name(&mappings, "non_matching", "pre_one_bar"), None);
119        assert_eq!(str_role_name(&mappings, "non_matching", "non_matching"), None);
120
121        assert_eq!(str_role_name(&mappings, "pre_two_foo", "pre_two_bar"), Some("replace_two"));
122        assert_eq!(str_role_name(&mappings, "pre_two_foo", "non_matching"), None);
123        assert_eq!(str_role_name(&mappings, "non_matching", "pre_two_bar"), None);
124
125        assert_eq!(
126            str_role_name(&mappings, "pre_three_foo", "pre_three_bar"),
127            Some("replace_three")
128        );
129        assert_eq!(str_role_name(&mappings, "pre_three_foo", "non_matching"), None);
130        assert_eq!(str_role_name(&mappings, "non_matching", "pre_three_bar"), None);
131
132        assert_eq!(str_role_name(&mappings, "pre_four_foo", "pre_four_bar"), Some("replace_four"));
133        assert_eq!(str_role_name(&mappings, "pre_four_foo", "non_matching"), None);
134        assert_eq!(str_role_name(&mappings, "non_matching", "pre_four_bar"), None);
135    }
136}