1use crate::{
3 cpu_penalty::{self, Config, State},
4 execution_control::{CommandRunner, SystemCommandRunner},
5 execution_store::{ExecutionRecord, ExecutionStore, ExecutionStoreOptions},
6 store_lock::LockManager,
7};
8use chrono::Utc;
9use serde_json::{json, Value};
10use std::sync::atomic::{AtomicBool, Ordering};
11use std::{
12 fs,
13 path::PathBuf,
14 thread,
15 time::{Duration, Instant},
16};
17static STOP: AtomicBool = AtomicBool::new(false);
18const STOP_POLL: Duration = Duration::from_millis(50);
19#[cfg(unix)]
20extern "C" fn stop_monitor(_: libc::c_int) {
21 STOP.store(true, Ordering::Relaxed);
22}
23fn sleep_until_stopped(interval: Duration) {
26 let started = Instant::now();
27 while !STOP.load(Ordering::Relaxed) {
28 let Some(rest) = interval.checked_sub(started.elapsed()) else {
29 break;
30 };
31 thread::sleep(rest.min(STOP_POLL));
32 }
33}
34fn debug(record: &ExecutionRecord, error: &str) {
35 if std::env::var("START_DEBUG").as_deref() == Ok("1") {
36 let text = format!("[DEBUG] CPU monitor: {}\n", error);
37 eprint!("{}", text);
38 crate::isolation::isolation_log::append_log_file(&PathBuf::from(&record.log_path), &text);
39 }
40}
41pub fn state_path(record: &ExecutionRecord) -> PathBuf {
42 PathBuf::from(format!("{}.cpu-penalty.json", record.log_path))
43}
44pub fn read_state(record: &ExecutionRecord) -> Option<Value> {
45 if !record.options.contains_key("cpuPenaltyConfig") || record.log_path.is_empty() {
46 return None;
47 }
48 let mut value: Value =
49 serde_json::from_str(&fs::read_to_string(state_path(record)).ok()?).ok()?;
50 if value["containerName"]
51 != record
52 .options
53 .get("sessionName")
54 .cloned()
55 .unwrap_or(Value::Null)
56 || value["attemptNumber"] != json!(record.attempt.as_ref().map(|a| a.number).unwrap_or(1))
57 {
58 return None;
59 }
60 if value["phase"] == "penalized" {
61 if let Some(updated) = value["updatedAt"].as_i64() {
62 let until = record
63 .end_time
64 .as_deref()
65 .and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok())
66 .map(|s| s.timestamp_millis())
67 .unwrap_or_else(|| Utc::now().timestamp_millis());
68 value["penalizedMs"] =
69 json!(value["penalizedMs"].as_i64().unwrap_or(0) + (until - updated).max(0));
70 value["updatedAt"] = json!(until.max(updated));
71 }
72 }
73 Some(value)
74}
75fn publish(
76 record: &ExecutionRecord,
77 state: &State,
78 name: &str,
79 attempt: u64,
80) -> Result<(), String> {
81 let file = state_path(record);
82 let tmp = file.with_extension(format!("{}.tmp", std::process::id()));
83 let mut value = serde_json::to_value(state).map_err(|e| e.to_string())?;
84 value["containerName"] = json!(name);
85 value["attemptNumber"] = json!(attempt);
86 value["updatedAt"] = json!(Utc::now().timestamp_millis());
87 let result = fs::write(&tmp, value.to_string())
88 .and_then(|_| fs::rename(&tmp, &file))
89 .map_err(|e| e.to_string());
90 let _ = fs::remove_file(tmp);
91 result
92}
93fn monitor(uuid: &str, name: &str, attempt: u64, standalone: Option<&str>) -> Result<(), String> {
94 let store = (!uuid.is_empty()).then(|| {
95 ExecutionStore::with_options(ExecutionStoreOptions {
96 app_folder: std::env::var_os("START_APP_FOLDER").map(PathBuf::from),
97 ..Default::default()
98 })
99 });
100 let Some(record) = store
101 .as_ref()
102 .and_then(|s| s.get(uuid))
103 .or_else(|| standalone.and_then(|s| serde_json::from_str::<ExecutionRecord>(s).ok()))
104 else {
105 return Ok(());
106 };
107 if record.log_path.is_empty()
108 || record.options.get("sessionName").and_then(Value::as_str) != Some(name)
109 || record.attempt.as_ref().map(|a| a.number).unwrap_or(1) != attempt
110 {
111 return Ok(());
112 }
113 let Some(config) = record
114 .options
115 .get("cpuPenaltyConfig")
116 .and_then(|v| serde_json::from_value::<Config>(v.clone()).ok())
117 else {
118 return Ok(());
119 };
120 let mut lock = LockManager::new(PathBuf::from(format!(
121 "{}.lock",
122 state_path(&record).display()
123 )))
124 .for_cpu_monitor();
125 if !lock.acquire(100) {
126 return Ok(());
127 }
128 let result = (|| {
129 let mut base = crate::docker_resource_limits::normalize_resource_limits(
130 record
131 .options
132 .get("baseResourceLimits")
133 .or_else(|| record.options.get("resourceLimits")),
134 );
135 let base_cpus = crate::docker_resource_limits::cpu_count(&base);
136 let saved = read_state(&record).and_then(|s| serde_json::from_value::<State>(s).ok());
137 let mut state = cpu_penalty::initial(base_cpus, Utc::now().timestamp_millis(), saved);
138 let interval =
139 (config.trigger_window_ms.min(config.release_window_ms) / 8).clamp(50, 30000);
140 let docker = crate::docker_cleanup::docker_command()
141 .to_string_lossy()
142 .into_owned();
143 let runner = SystemCommandRunner;
144 while !STOP.load(Ordering::Relaxed) {
145 let tick = (|| -> Result<bool, String> {
146 let Some(current) = store
147 .as_ref()
148 .map(|s| s.get(uuid))
149 .unwrap_or_else(|| Some(record.clone()))
150 else {
151 return Ok(false);
152 };
153 if current.options.get("sessionName").and_then(Value::as_str) != Some(name)
154 || current.attempt.as_ref().map(|a| a.number).unwrap_or(1) != attempt
155 {
156 return Ok(false);
157 }
158 let live = runner.run(
159 &docker,
160 &[
161 "inspect".into(),
162 "-f".into(),
163 "{{.State.Running}}".into(),
164 name.into(),
165 ],
166 );
167 if (live.success && live.stdout.trim() == "false")
168 || (!live.success
169 && (live.stderr.contains("No such object")
170 || live.stderr.contains("No such container")))
171 {
172 return Ok(false);
173 }
174 let observed =
175 crate::docker_resource_limits::read_docker_resource_limits(name, &runner);
176 if let Some(limits) = observed.as_ref() {
177 base = limits.clone();
178 state = cpu_penalty::observe_base(
179 &state,
180 crate::docker_resource_limits::cpu_count(limits),
181 Utc::now().timestamp_millis(),
182 );
183 }
184 let info = runner.run(
185 &docker,
186 &["info".into(), "--format".into(), "{{json .}}".into()],
187 );
188 let stats = runner.run(
189 &docker,
190 &[
191 "stats".into(),
192 "--no-stream".into(),
193 "--format".into(),
194 "{{json .}}".into(),
195 ],
196 );
197 let cpus = if info.success {
198 serde_json::from_str::<Value>(&info.stdout)
199 .ok()
200 .and_then(|v| v["NCPU"].as_f64())
201 .unwrap_or(f64::NAN)
202 } else {
203 f64::NAN
204 };
205 let cores = if stats.success && live.success && observed.is_some() {
206 stats
207 .stdout
208 .lines()
209 .filter_map(|s| serde_json::from_str::<Value>(s).ok())
210 .find(|v| v["Name"].as_str() == Some(name))
211 .and_then(|v| {
212 v["CPUPerc"]
213 .as_str()
214 .and_then(|s| s.strip_suffix('%'))
215 .and_then(|s| s.parse::<f64>().ok())
216 })
217 .map(|p| p / 100.0)
218 .unwrap_or(f64::NAN)
219 } else {
220 f64::NAN
221 };
222 debug(
223 &record,
224 &format!(
225 "sample: phase={} base={:?} limit={:?} cores={} daemon={}",
226 state.phase, state.base_cpus, state.limit_cpus, cores, cpus
227 ),
228 );
229 let (mut next, action) = cpu_penalty::evaluate(
230 &state,
231 Utc::now().timestamp_millis(),
232 cores,
233 cpus,
234 (interval * 3).max(5000),
235 &config,
236 );
237 if let Some(action) = action {
238 if action.kind == "apply" {
239 publish(&record, &next, name, attempt)?;
240 }
241 let update = runner.run(
242 &docker,
243 &[
244 vec!["update".into()],
245 crate::docker_resource_limits::cpu_update_args(&base, action.cpus),
246 vec![name.into()],
247 ]
248 .concat(),
249 );
250 if update.success {
251 crate::isolation::isolation_log::append_log_file(
252 &PathBuf::from(&record.log_path),
253 &format!("{}\n", cpu_penalty::transition_log(&action, &config)),
254 );
255 } else {
256 debug(&record, &format!("docker update failed: {}", update.stderr));
257 next.phase = state.phase.clone();
258 next.since = state.since;
259 next.limit_cpus = state.limit_cpus;
260 next.penalty_count = state.penalty_count;
261 next.samples.clear();
262 }
263 }
264 state = next;
265 publish(&record, &state, name, attempt)?;
266 Ok(true)
267 })();
268 match tick {
269 Ok(false) => break,
270 Ok(true) => {}
271 Err(error) => {
272 state.samples.clear();
273 debug(&record, &format!("tick failed: {}", error));
274 }
275 }
276 sleep_until_stopped(Duration::from_millis(interval as u64));
277 }
278 let current = store
279 .as_ref()
280 .map(|s| s.get(uuid))
281 .unwrap_or_else(|| Some(record.clone()));
282 if state.phase == "penalized"
283 && current.is_some_and(|r| {
284 r.options.get("sessionName").and_then(Value::as_str) == Some(name)
285 && r.attempt.as_ref().map(|a| a.number).unwrap_or(1) == attempt
286 })
287 {
288 let live = runner.run(
289 &docker,
290 &[
291 "inspect".into(),
292 "-f".into(),
293 "{{.State.Running}}".into(),
294 name.into(),
295 ],
296 );
297 if live.success && live.stdout.trim() == "true" {
298 if let Some(limits) =
299 crate::docker_resource_limits::read_docker_resource_limits(name, &runner)
300 {
301 if crate::docker_resource_limits::cpu_count(&limits) == state.limit_cpus {
302 let base_limits = if let Some(base) = state.base_cpus {
303 crate::docker_resource_limits::with_cpu_limit(limits.clone(), base)
304 } else {
305 limits
306 .iter()
307 .filter(|s| {
308 !["--cpus=", "--cpu-quota=", "--cpu-period="]
309 .iter()
310 .any(|p| s.starts_with(p))
311 })
312 .cloned()
313 .collect()
314 };
315 let restore =
316 crate::resume_resources::clamp_cpu_to_daemon(base_limits, &runner)?;
317 let cpus = crate::docker_resource_limits::cpu_count(&restore)
318 .ok_or("Cannot restore CPU base")?;
319 let update = runner.run(
320 &docker,
321 &[
322 vec!["update".into()],
323 crate::docker_resource_limits::cpu_update_args(&limits, cpus),
324 vec![name.into()],
325 ]
326 .concat(),
327 );
328 if update.success {
329 let now = Utc::now().timestamp_millis();
330 cpu_penalty::accrue(&mut state, now);
331 state.phase = "observing".into();
332 state.limit_cpus = Some(cpus);
333 state.since = now;
334 crate::isolation::isolation_log::append_log_file(
335 &PathBuf::from(&record.log_path),
336 &format!(
337 "[start] CPU penalty lifted: {} CPUs (monitor stopped)\n",
338 cpus
339 ),
340 );
341 if let Err(error) = publish(&record, &state, name, attempt) {
342 debug(&record, &error);
343 }
344 }
345 }
346 }
347 }
348 }
349 Ok(())
350 })();
351 lock.release();
352 result
353}
354pub fn main(args: &[String]) {
355 if args.len() < 3 {
356 return;
357 }
358 STOP.store(false, Ordering::Relaxed);
359 #[cfg(unix)]
360 unsafe {
361 libc::signal(libc::SIGTERM, stop_monitor as *const () as usize);
362 libc::signal(libc::SIGINT, stop_monitor as *const () as usize);
363 }
364 if let Err(error) = monitor(
365 &args[0],
366 &args[1],
367 args[2].parse().unwrap_or(1),
368 args.get(3).map(String::as_str),
369 ) {
370 if std::env::var("START_DEBUG").as_deref() == Ok("1") {
371 eprintln!("[DEBUG] CPU monitor: {}", error);
372 }
373 std::process::exit(1);
374 }
375}
376pub fn start_snippet(uuid: Option<&str>, name: &str, attempt: Option<u64>) -> String {
377 use crate::isolation::isolation_log::shell_quote;
378 let Some(uuid) = uuid else {
379 return String::new();
380 };
381 let exe = std::env::current_exe()
382 .ok()
383 .map(|p| p.to_string_lossy().into_owned())
384 .unwrap_or_else(|| "$".into());
385 supervise_snippet(&format!(
386 "{} __start-cpu-monitor {} {} {}",
387 shell_quote(&exe),
388 shell_quote(uuid),
389 shell_quote(name),
390 attempt.unwrap_or(1)
391 ))
392}
393fn supervise_snippet(command: &str) -> String {
394 format!("( __start_cpu_child=; trap 'if [ -n \"$__start_cpu_child\" ]; then kill \"$__start_cpu_child\" 2>/dev/null; wait \"$__start_cpu_child\" 2>/dev/null; fi; exit 0' TERM INT HUP; while :; do {} >/dev/null 2>&1 & __start_cpu_child=$!; wait \"$__start_cpu_child\"; __start_cpu_code=$?; [ \"$__start_cpu_code\" -eq 0 ] && break; sleep 1 & __start_cpu_child=$!; wait \"$__start_cpu_child\"; done ) & __start_cpu_monitor=$!", command)
395}
396pub fn stop_snippet() -> String {
397 "if [ -n \"$__start_cpu_monitor\" ]; then kill \"$__start_cpu_monitor\" 2>/dev/null; wait \"$__start_cpu_monitor\" 2>/dev/null; fi".into()
398}
399
400fn standalone_record(name: &str, options: &crate::isolation::IsolationOptions) -> Option<String> {
401 let config = options.cpu_penalty_config.as_ref()?;
402 let mut record = ExecutionRecord::new("CPU monitor");
403 record.log_path = options.log_path.as_ref()?.to_string_lossy().into();
404 record.options.insert("sessionName".into(), json!(name));
405 record
406 .options
407 .insert("cpuPenaltyConfig".into(), json!(config));
408 record
409 .options
410 .insert("baseResourceLimits".into(), json!(options.resource_limits));
411 serde_json::to_string(&record).ok()
412}
413pub fn start_snippet_with_options(
414 name: &str,
415 options: &crate::isolation::IsolationOptions,
416) -> String {
417 if options.execution_id.is_some() {
418 return start_snippet(options.execution_id.as_deref(), name, Some(1));
419 }
420 let Some(record) = standalone_record(name, options) else {
421 return String::new();
422 };
423 let exe = std::env::current_exe().unwrap_or_default();
424 use crate::isolation::isolation_log::shell_quote;
425 supervise_snippet(&format!(
426 "{} __start-cpu-monitor '' {} 1 {}",
427 shell_quote(&exe.to_string_lossy()),
428 shell_quote(name),
429 shell_quote(&record)
430 ))
431}
432pub fn spawn_standalone(name: &str, options: &crate::isolation::IsolationOptions) {
433 let snippet = start_snippet_with_options(name, options);
434 if !snippet.is_empty() {
435 let _ = std::process::Command::new("sh")
436 .args(["-c", &snippet])
437 .stdout(std::process::Stdio::null())
438 .stderr(std::process::Stdio::null())
439 .spawn();
440 }
441}