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
23pub 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 if setrlimit(Resource::RLIMIT_MEMLOCK, RLIM_INFINITY, RLIM_INFINITY).is_err() {
89 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
99pub 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 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
155pub fn normalize_load_metric(metric: f64) -> f64 {
174 metric / 100.0
175}
176
177pub 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 total_items % size < total_items % optimal_size {
205 optimal_size = size;
206 }
207 }
208
209 optimal_size
210}