Skip to main content

scx_utils/
misc.rs

1use anyhow::Result;
2use anyhow::{Context, anyhow};
3use log::{info, warn};
4use nix::sys::resource::{RLIM_INFINITY, Resource, getrlimit, setrlimit};
5use scx_stats::prelude::*;
6use serde::Deserialize;
7use std::path::Path;
8use std::thread::sleep;
9use std::time::Duration;
10
11pub fn monitor_stats<T>(
12    stats_args: &[(String, String)],
13    intv: Duration,
14    should_exit: impl FnMut() -> bool,
15    output: impl FnMut(T) -> Result<()>,
16) -> Result<()>
17where
18    T: for<'a> Deserialize<'a>,
19{
20    monitor_stats_sched(None, stats_args, intv, should_exit, output)
21}
22
23/// Like monitor_stats(), but for the stats server at `sched_path`, the path the
24/// server was given through StatsServer::set_sched_path(). None selects the
25/// default server.
26pub fn monitor_stats_sched<T>(
27    sched_path: Option<&Path>,
28    stats_args: &[(String, String)],
29    intv: Duration,
30    mut should_exit: impl FnMut() -> bool,
31    mut output: impl FnMut(T) -> Result<()>,
32) -> Result<()>
33where
34    T: for<'a> Deserialize<'a>,
35{
36    let mut retry_cnt: u32 = 0;
37
38    const RETRYABLE_ERRORS: [std::io::ErrorKind; 2] = [
39        std::io::ErrorKind::NotFound,
40        std::io::ErrorKind::ConnectionRefused,
41    ];
42
43    while !should_exit() {
44        let mut client = StatsClient::new();
45        if let Some(sched_path) = sched_path {
46            client = client.set_sched_path(sched_path);
47        }
48        let mut client = match client.connect(None) {
49            Ok(v) => v,
50            Err(e) => match e.downcast_ref::<std::io::Error>() {
51                Some(ioe) if RETRYABLE_ERRORS.contains(&ioe.kind()) => {
52                    if retry_cnt == 1 {
53                        info!("Stats server not available, retrying...");
54                    }
55                    retry_cnt += 1;
56                    sleep(Duration::from_secs(1));
57                    continue;
58                }
59                _ => Err(e)?,
60            },
61        };
62        retry_cnt = 0;
63
64        while !should_exit() {
65            let stats = match client.request::<T>("stats", stats_args.to_owned()) {
66                Ok(v) => v,
67                Err(e) => {
68                    if let Some(ioe) = e.downcast_ref::<std::io::Error>() {
69                        info!("Connection to stats_server failed ({ioe})");
70                    } else {
71                        warn!("Error handling stats_server result: {e}");
72                    }
73                    sleep(Duration::from_secs(1));
74                    break;
75                }
76            };
77            output(stats)?;
78            sleep(intv);
79        }
80    }
81
82    Ok(())
83}
84
85pub fn try_set_rlimit_infinity() {
86    // Increase MEMLOCK size since the BPF scheduler might use
87    // more than the current limit
88    if setrlimit(Resource::RLIMIT_MEMLOCK, RLIM_INFINITY, RLIM_INFINITY).is_err() {
89        // If there is an error in expanding rlimit to infinity,
90        // show the current rlimits then proceed.
91        if let Ok((soft, hard)) = getrlimit(Resource::RLIMIT_MEMLOCK) {
92            warn!("Current MEMLOCK limit: soft={soft}, hard={hard}");
93        } else {
94            warn!("Cannot change or query MEMLOCK limit");
95        }
96    }
97}
98
99/// Read a file and parse its content into the specified type.
100///
101/// Trims null and whitespace before parsing.
102///
103/// # Errors
104/// Returns an error if reading or parsing fails.
105pub fn read_from_file<T>(path: &Path) -> Result<T>
106where
107    T: std::str::FromStr,
108    T::Err: std::error::Error + Send + Sync + 'static,
109{
110    let val = std::fs::read_to_string(path)
111        .with_context(|| format!("Failed to open or read file: {}", path.display()))?;
112
113    let val = val.trim_end_matches('\0').trim();
114
115    val.parse::<T>()
116        .with_context(|| format!("Failed to parse content '{}' from {}", val, path.display()))
117}
118
119pub fn read_file_usize_vec(path: &Path, separator: char) -> Result<Vec<usize>> {
120    let val = std::fs::read_to_string(path)?;
121    let val = val.trim_end_matches('\0');
122
123    val.split(separator)
124        .map(|s| {
125            s.trim()
126                .parse::<usize>()
127                .map_err(|_| anyhow!("Failed to parse '{}' as usize", s))
128        })
129        .collect::<Result<Vec<usize>>>()
130}
131
132pub fn read_file_byte(path: &Path) -> Result<usize> {
133    let val = std::fs::read_to_string(path)?;
134    let val = val.trim_end_matches('\0');
135    let val = val.trim();
136
137    // E.g., 10K, 10M, 10G, 10
138    if let Some(sval) = val.strip_suffix("K") {
139        let byte = sval.parse::<usize>()?;
140        return Ok(byte * 1024);
141    }
142    if let Some(sval) = val.strip_suffix("M") {
143        let byte = sval.parse::<usize>()?;
144        return Ok(byte * 1024 * 1024);
145    }
146    if let Some(sval) = val.strip_suffix("G") {
147        let byte = sval.parse::<usize>()?;
148        return Ok(byte * 1024 * 1024 * 1024);
149    }
150
151    let byte = val.parse::<usize>()?;
152    Ok(byte)
153}
154
155/* Load is reported as weight * duty cycle
156 *
157 * In the Linux kernel, EEDVF uses default weight = 1 s.t.
158 * load for a nice-0 thread runnable for time slice = 1
159 *
160 * To conform with cgroup weights convention, sched-ext uses
161 * the convention of default weight = 100 with the formula
162 * 100 * nice ^ 1.5. This means load for a nice-0 thread
163 * runnable for time slice = 100.
164 *
165 * To ensure we report load metrics consistently with the Linux
166 * kernel, we divide load by 100.0 prior to reporting metrics.
167 * This is also more intuitive for users since 1 CPU roughly
168 * means 1 unit of load.
169 *
170 * We only do this prior to reporting as its easier to work with
171 * weight as integers in BPF / userspace than floating point.
172 */
173pub fn normalize_load_metric(metric: f64) -> f64 {
174    metric / 100.0
175}
176
177/// Find the best split size for dividing a total number of items
178/// given a range of sizes.
179///
180/// Searches from starting with min to max to ideally find an even
181/// split with no remainders. Otherwise, choose a size that minimizes
182/// the remainder favoring smaller sizes.
183///
184/// # Arguments
185/// * `total_items` - Total number of items to split
186/// * `min_size` - Minimum size per split
187/// * `max_size` - Maximum size per split
188///
189/// # Returns
190/// The optimal partition size
191pub fn find_best_split_size(total_items: usize, min_size: usize, max_size: usize) -> usize {
192    if total_items <= min_size {
193        return total_items;
194    }
195
196    let mut optimal_size = min_size;
197
198    for size in min_size..=max_size {
199        if total_items.is_multiple_of(size) {
200            optimal_size = size;
201            break;
202        }
203        // If no perfect division, use the one that minimizes remainder
204        if total_items % size < total_items % optimal_size {
205            optimal_size = size;
206        }
207    }
208
209    optimal_size
210}