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