start_command/
cpu_penalty.rs

1//! Opt-in timestamp-weighted CPU hysteresis. Missing coverage cannot trigger.
2use 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}
158/// Penalized time counts until the phase leaves `penalized`; `--status` only
159/// extrapolates while penalized, so every transition settles it first.
160pub fn accrue(state: &mut State, now: i64) {
161    if state.phase == "penalized" {
162        state.penalized_ms += (now - state.last_at).max(0) as u64;
163    }
164    state.last_at = now;
165}
166pub fn observe_base(state: &State, base: Option<f64>, now: i64) -> State {
167    if base == state.limit_cpus {
168        return state.clone();
169    }
170    let mut next = state.clone();
171    accrue(&mut next, now);
172    next.phase = "observing".into();
173    next.base_cpus = base;
174    next.limit_cpus = base;
175    next.since = now;
176    next.samples.clear();
177    next
178}
179
180pub fn transition_log(action: &Action, config: &Config) -> String {
181    let applying = action.kind == "apply";
182    let window = if applying {
183        config.trigger_window_ms
184    } else {
185        config.release_window_ms
186    };
187    let duration = if window % 60000 == 0 {
188        format!("{}m", window / 60000)
189    } else {
190        format!("{}ms", window)
191    };
192    let measured = action
193        .average
194        .map(|n| {
195            format!(
196                "{:.2} of {} CPUs for {}",
197                n,
198                if applying {
199                    action.capacity
200                } else {
201                    config.cpus
202                },
203                duration
204            )
205        })
206        .unwrap_or_else(|| "capacity changed".into());
207    format!(
208        "[start] CPU penalty {}: {} CPUs ({})",
209        if applying { "applied" } else { "lifted" },
210        action.cpus,
211        measured
212    )
213}
214pub fn average(samples: &[Sample], now: i64, window: i64) -> Option<f64> {
215    let start = now - window;
216    if samples.len() < 2 || samples[0].at > start || samples.last()?.at != now {
217        return None;
218    }
219    let sum = samples
220        .windows(2)
221        .map(|pair| {
222            let left = start.max(pair[0].at);
223            let right = pair[1].at;
224            (right - left).max(0) as f64 * pair[0].cores
225        })
226        .sum::<f64>();
227    Some(sum / window as f64)
228}
229#[derive(Debug, Clone)]
230pub struct Action {
231    pub kind: &'static str,
232    pub cpus: f64,
233    pub average: Option<f64>,
234    pub capacity: f64,
235}
236pub fn evaluate(
237    state: &State,
238    now: i64,
239    cores: f64,
240    daemon_cpus: f64,
241    max_gap: i64,
242    config: &Config,
243) -> (State, Option<Action>) {
244    let mut s = state.clone();
245    accrue(&mut s, now);
246    let capacity = s.base_cpus.unwrap_or(f64::INFINITY).min(daemon_cpus);
247    if !cores.is_finite()
248        || cores < 0.0
249        || !daemon_cpus.is_finite()
250        || capacity <= 0.0
251        || s.samples
252            .last()
253            .is_some_and(|last| now <= last.at || now - last.at > max_gap)
254        || s.capacity != Some(capacity)
255    {
256        s.samples.clear();
257    }
258    s.capacity = Some(capacity);
259    if !cores.is_finite() || cores < 0.0 || !daemon_cpus.is_finite() || capacity <= 0.0 {
260        return (s, None);
261    }
262    s.samples.push(Sample { at: now, cores });
263    let window = if s.phase == "penalized" {
264        config.release_window_ms
265    } else {
266        config.trigger_window_ms
267    };
268    while s.samples.len() > 2 && s.samples[1].at <= now - window {
269        s.samples.remove(0);
270    }
271    let avg = average(&s.samples, now, window);
272    let mut action = None;
273    if s.phase == "observing"
274        && config.cpus < capacity
275        && avg.is_some_and(|a| a >= config.trigger / 100.0 * capacity)
276    {
277        action = Some(Action {
278            kind: "apply",
279            cpus: config.cpus,
280            average: avg,
281            capacity,
282        });
283        s.phase = "penalized".into();
284        s.since = now;
285        s.limit_cpus = Some(config.cpus);
286        s.penalty_count += 1;
287    } else if s.phase == "penalized"
288        && (config.cpus >= capacity
289            || (now - s.since >= config.release_window_ms
290                && avg.is_some_and(|a| a < config.release / 100.0 * config.cpus)))
291    {
292        action = Some(Action {
293            kind: "lift",
294            cpus: capacity,
295            average: avg,
296            capacity,
297        });
298        s.phase = "observing".into();
299        s.since = now;
300        s.limit_cpus = Some(capacity);
301    }
302    if action.is_some() {
303        s.samples.clear();
304    }
305    (s, action)
306}
start_command/cpu_penalty.rs

start_command/
cpu_penalty.rs
