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.base_cpus = base;
138 s.samples.clear();
139 s.capacity = None;
140 s.last_at = now;
141 if s.phase != "penalized" {
142 s.phase = "observing".into();
143 s.since = now;
144 s.limit_cpus = base;
145 }
146 return s;
147 }
148 State {
149 phase: "observing".into(),
150 since: now,
151 limit_cpus: base,
152 base_cpus: base,
153 penalty_count: 0,
154 penalized_ms: 0,
155 samples: vec![],
156 last_at: now,
157 capacity: None,
158 }
159}
160pub fn average(samples: &[Sample], now: i64, window: i64) -> Option<f64> {
161 let start = now - window;
162 if samples.len() < 2 || samples[0].at > start || samples.last()?.at != now {
163 return None;
164 }
165 let sum = samples
166 .windows(2)
167 .map(|pair| {
168 let left = start.max(pair[0].at);
169 let right = pair[1].at;
170 (right - left).max(0) as f64 * pair[0].cores
171 })
172 .sum::<f64>();
173 Some(sum / window as f64)
174}
175#[derive(Debug, Clone)]
176pub struct Action {
177 pub kind: &'static str,
178 pub cpus: f64,
179 pub average: Option<f64>,
180 pub capacity: f64,
181}
182pub fn evaluate(
183 state: &State,
184 now: i64,
185 cores: f64,
186 daemon_cpus: f64,
187 max_gap: i64,
188 config: &Config,
189) -> (State, Option<Action>) {
190 let mut s = state.clone();
191 if s.phase == "penalized" {
192 s.penalized_ms += (now - s.last_at).max(0) as u64;
193 }
194 s.last_at = now;
195 let capacity = s.base_cpus.unwrap_or(f64::INFINITY).min(daemon_cpus);
196 if !cores.is_finite()
197 || cores < 0.0
198 || !daemon_cpus.is_finite()
199 || capacity <= 0.0
200 || s.samples
201 .last()
202 .is_some_and(|last| now <= last.at || now - last.at > max_gap)
203 || s.capacity != Some(capacity)
204 {
205 s.samples.clear();
206 }
207 s.capacity = Some(capacity);
208 if !cores.is_finite() || cores < 0.0 || !daemon_cpus.is_finite() || capacity <= 0.0 {
209 return (s, None);
210 }
211 s.samples.push(Sample { at: now, cores });
212 let window = if s.phase == "penalized" {
213 config.release_window_ms
214 } else {
215 config.trigger_window_ms
216 };
217 while s.samples.len() > 2 && s.samples[1].at <= now - window {
218 s.samples.remove(0);
219 }
220 let avg = average(&s.samples, now, window);
221 let mut action = None;
222 if s.phase == "observing"
223 && config.cpus < capacity
224 && avg.is_some_and(|a| a >= config.trigger / 100.0 * capacity)
225 {
226 action = Some(Action {
227 kind: "apply",
228 cpus: config.cpus,
229 average: avg,
230 capacity,
231 });
232 s.phase = "penalized".into();
233 s.since = now;
234 s.limit_cpus = Some(config.cpus);
235 s.penalty_count += 1;
236 } else if s.phase == "penalized"
237 && (config.cpus >= capacity
238 || (now - s.since >= config.release_window_ms
239 && avg.is_some_and(|a| a < config.release / 100.0 * config.cpus)))
240 {
241 action = Some(Action {
242 kind: "lift",
243 cpus: capacity,
244 average: avg,
245 capacity,
246 });
247 s.phase = "observing".into();
248 s.since = now;
249 s.limit_cpus = Some(capacity);
250 }
251 if action.is_some() {
252 s.samples.clear();
253 }
254 (s, action)
255}