1use crate::args_parser::WrapperOptions;
3use serde::{Deserialize, Serialize};
4#[derive(Debug, Clone, Serialize, Deserialize)]
5#[serde(rename_all = "camelCase")]
6pub struct Config {
7 pub cpus: f64,
8 pub trigger: f64,
9 pub trigger_window_ms: i64,
10 pub release: f64,
11 pub release_window_ms: i64,
12}
13impl Default for Config {
14 fn default() -> Self {
15 Self {
16 cpus: 2.0,
17 trigger: 95.0,
18 trigger_window_ms: 900000,
19 release: 65.0,
20 release_window_ms: 900000,
21 }
22 }
23}
24pub fn duration(value: &str) -> Result<i64, String> {
25 let end = value
26 .find(|c: char| !c.is_ascii_digit() && c != '.')
27 .unwrap_or(value.len());
28 let amount = value[..end]
29 .parse::<f64>()
30 .map_err(|_| format!("Invalid CPU penalty duration: {}", value))?;
31 let unit = match &value[end..] {
32 "ms" => 1.0,
33 "s" => 1000.0,
34 "m" => 60000.0,
35 "h" => 3600000.0,
36 _ => return Err(format!("Invalid CPU penalty duration: {}", value)),
37 };
38 let ms = amount * unit;
39 if !ms.is_finite() || ms <= 0.0 || ms > 9007199254740991.0 || ms.fract() != 0.0 {
40 return Err(format!("Invalid CPU penalty duration: {}", value));
41 }
42 Ok(ms as i64)
43}
44pub fn parse(args: &[String], index: usize, options: &mut WrapperOptions) -> Result<usize, String> {
45 let (flag, inline) = args[index]
46 .split_once('=')
47 .map(|(f, v)| (f, Some(v)))
48 .unwrap_or((&args[index], None));
49 if flag == "--cpu-penalty" {
50 if inline.is_some() {
51 return Err("--cpu-penalty does not take a value".into());
52 }
53 options.cpu_penalty = true;
54 return Ok(1);
55 }
56 if ![
57 "--cpu-penalty-cpus",
58 "--cpu-penalty-trigger",
59 "--cpu-penalty-trigger-window",
60 "--cpu-penalty-release",
61 "--cpu-penalty-release-window",
62 ]
63 .contains(&flag)
64 {
65 return Ok(0);
66 }
67 let value = inline
68 .or_else(|| args.get(index + 1).map(String::as_str))
69 .ok_or_else(|| format!("{} requires a value", flag))?;
70 let config = options
71 .cpu_penalty_config
72 .get_or_insert_with(Config::default);
73 if flag.ends_with("-window") {
74 let ms = duration(value)?;
75 if flag.contains("trigger") {
76 config.trigger_window_ms = ms;
77 } else {
78 config.release_window_ms = ms;
79 }
80 } else {
81 let n = value
82 .trim_end_matches('%')
83 .parse::<f64>()
84 .map_err(|_| format!("Invalid {}: {}", flag, value))?;
85 if !n.is_finite()
86 || n <= 0.0
87 || (flag != "--cpu-penalty-cpus" && n > 100.0)
88 || (flag == "--cpu-penalty-cpus" && value.ends_with('%'))
89 {
90 return Err(format!("Invalid {}: {}", flag, value));
91 }
92 match flag {
93 "--cpu-penalty-cpus" => config.cpus = n,
94 "--cpu-penalty-trigger" => config.trigger = n,
95 _ => config.release = n,
96 }
97 }
98 Ok(if inline.is_some() { 1 } else { 2 })
99}
100pub fn validate(options: &mut WrapperOptions) -> Result<(), String> {
101 if options.cpu_penalty || options.cpu_penalty_config.is_some() {
102 if !options.cpu_penalty {
103 return Err("CPU penalty settings require --cpu-penalty".into());
104 }
105 if options.isolated.as_deref() != Some("docker") {
106 return Err("--cpu-penalty requires Docker isolation".into());
107 }
108 options
109 .cpu_penalty_config
110 .get_or_insert_with(Config::default);
111 }
112 Ok(())
113}
114#[derive(Debug, Clone)]
115pub struct Sample {
116 pub at: i64,
117 pub cores: f64,
118}
119#[derive(Debug, Clone, Serialize, Deserialize)]
120#[serde(rename_all = "camelCase")]
121pub struct State {
122 pub phase: String,
123 pub since: i64,
124 pub limit_cpus: Option<f64>,
125 pub base_cpus: Option<f64>,
126 pub penalty_count: u64,
127 pub penalized_ms: u64,
128 #[serde(skip)]
129 pub samples: Vec<Sample>,
130 #[serde(skip)]
131 pub last_at: i64,
132 #[serde(skip)]
133 pub capacity: Option<f64>,
134}
135pub fn initial(base: Option<f64>, now: i64, saved: Option<State>) -> State {
136 if let Some(mut s) = saved {
137 s.samples.clear();
138 s.capacity = None;
139 s.last_at = now;
140 if s.phase != "penalized" {
141 s.phase = "observing".into();
142 s.since = now;
143 }
144 return s;
145 }
146 State {
147 phase: "observing".into(),
148 since: now,
149 limit_cpus: base,
150 base_cpus: base,
151 penalty_count: 0,
152 penalized_ms: 0,
153 samples: vec![],
154 last_at: now,
155 capacity: None,
156 }
157}
158pub fn observe_base(state: &State, base: Option<f64>, now: i64) -> State {
159 if base == state.limit_cpus {
160 return state.clone();
161 }
162 let mut next = state.clone();
163 next.phase = "observing".into();
164 next.base_cpus = base;
165 next.limit_cpus = base;
166 next.since = now;
167 next.samples.clear();
168 next
169}
170
171pub fn transition_log(action: &Action, config: &Config) -> String {
172 let applying = action.kind == "apply";
173 let window = if applying {
174 config.trigger_window_ms
175 } else {
176 config.release_window_ms
177 };
178 let duration = if window % 60000 == 0 {
179 format!("{}m", window / 60000)
180 } else {
181 format!("{}ms", window)
182 };
183 let measured = action
184 .average
185 .map(|n| {
186 format!(
187 "{:.2} of {} CPUs for {}",
188 n,
189 if applying {
190 action.capacity
191 } else {
192 config.cpus
193 },
194 duration
195 )
196 })
197 .unwrap_or_else(|| "capacity changed".into());
198 format!(
199 "[start] CPU penalty {}: {} CPUs ({})",
200 if applying { "applied" } else { "lifted" },
201 action.cpus,
202 measured
203 )
204}
205pub fn average(samples: &[Sample], now: i64, window: i64) -> Option<f64> {
206 let start = now - window;
207 if samples.len() < 2 || samples[0].at > start || samples.last()?.at != now {
208 return None;
209 }
210 let sum = samples
211 .windows(2)
212 .map(|pair| {
213 let left = start.max(pair[0].at);
214 let right = pair[1].at;
215 (right - left).max(0) as f64 * pair[0].cores
216 })
217 .sum::<f64>();
218 Some(sum / window as f64)
219}
220#[derive(Debug, Clone)]
221pub struct Action {
222 pub kind: &'static str,
223 pub cpus: f64,
224 pub average: Option<f64>,
225 pub capacity: f64,
226}
227pub fn evaluate(
228 state: &State,
229 now: i64,
230 cores: f64,
231 daemon_cpus: f64,
232 max_gap: i64,
233 config: &Config,
234) -> (State, Option<Action>) {
235 let mut s = state.clone();
236 if s.phase == "penalized" {
237 s.penalized_ms += (now - s.last_at).max(0) as u64;
238 }
239 s.last_at = now;
240 let capacity = s.base_cpus.unwrap_or(f64::INFINITY).min(daemon_cpus);
241 if !cores.is_finite()
242 || cores < 0.0
243 || !daemon_cpus.is_finite()
244 || capacity <= 0.0
245 || s.samples
246 .last()
247 .is_some_and(|last| now <= last.at || now - last.at > max_gap)
248 || s.capacity != Some(capacity)
249 {
250 s.samples.clear();
251 }
252 s.capacity = Some(capacity);
253 if !cores.is_finite() || cores < 0.0 || !daemon_cpus.is_finite() || capacity <= 0.0 {
254 return (s, None);
255 }
256 s.samples.push(Sample { at: now, cores });
257 let window = if s.phase == "penalized" {
258 config.release_window_ms
259 } else {
260 config.trigger_window_ms
261 };
262 while s.samples.len() > 2 && s.samples[1].at <= now - window {
263 s.samples.remove(0);
264 }
265 let avg = average(&s.samples, now, window);
266 let mut action = None;
267 if s.phase == "observing"
268 && config.cpus < capacity
269 && avg.is_some_and(|a| a >= config.trigger / 100.0 * capacity)
270 {
271 action = Some(Action {
272 kind: "apply",
273 cpus: config.cpus,
274 average: avg,
275 capacity,
276 });
277 s.phase = "penalized".into();
278 s.since = now;
279 s.limit_cpus = Some(config.cpus);
280 s.penalty_count += 1;
281 } else if s.phase == "penalized"
282 && (config.cpus >= capacity
283 || (now - s.since >= config.release_window_ms
284 && avg.is_some_and(|a| a < config.release / 100.0 * config.cpus)))
285 {
286 action = Some(Action {
287 kind: "lift",
288 cpus: capacity,
289 average: avg,
290 capacity,
291 });
292 s.phase = "observing".into();
293 s.since = now;
294 s.limit_cpus = Some(capacity);
295 }
296 if action.is_some() {
297 s.samples.clear();
298 }
299 (s, action)
300}