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::{fs, path::PathBuf, thread, time::Duration};
11pub fn state_path(record: &ExecutionRecord) -> PathBuf {
12 PathBuf::from(format!("{}.cpu-penalty.json", record.log_path))
13}
14pub fn read_state(record: &ExecutionRecord) -> Option<Value> {
15 if !record.options.contains_key("cpuPenaltyConfig") || record.log_path.is_empty() {
16 return None;
17 }
18 let mut value: Value =
19 serde_json::from_str(&fs::read_to_string(state_path(record)).ok()?).ok()?;
20 if value["containerName"]
21 != record
22 .options
23 .get("sessionName")
24 .cloned()
25 .unwrap_or(Value::Null)
26 || value["attemptNumber"] != json!(record.attempt.as_ref().map(|a| a.number).unwrap_or(1))
27 {
28 return None;
29 }
30 if value["phase"] == "penalized" {
31 if let Some(updated) = value["updatedAt"].as_i64() {
32 let until = record
33 .end_time
34 .as_deref()
35 .and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok())
36 .map(|s| s.timestamp_millis())
37 .unwrap_or_else(|| Utc::now().timestamp_millis());
38 value["penalizedMs"] =
39 json!(value["penalizedMs"].as_i64().unwrap_or(0) + (until - updated).max(0));
40 value["updatedAt"] = json!(until.max(updated));
41 }
42 }
43 Some(value)
44}
45fn publish(
46 record: &ExecutionRecord,
47 state: &State,
48 name: &str,
49 attempt: u64,
50) -> Result<(), String> {
51 let file = state_path(record);
52 let tmp = file.with_extension(format!("{}.tmp", std::process::id()));
53 let mut value = serde_json::to_value(state).map_err(|e| e.to_string())?;
54 value["containerName"] = json!(name);
55 value["attemptNumber"] = json!(attempt);
56 value["updatedAt"] = json!(Utc::now().timestamp_millis());
57 let result = fs::write(&tmp, value.to_string())
58 .and_then(|_| fs::rename(&tmp, &file))
59 .map_err(|e| e.to_string());
60 let _ = fs::remove_file(tmp);
61 result
62}
63fn monitor(uuid: &str, name: &str, attempt: u64, standalone: Option<&str>) -> Result<(), String> {
64 let store = (!uuid.is_empty()).then(|| {
65 ExecutionStore::with_options(ExecutionStoreOptions {
66 app_folder: std::env::var_os("START_APP_FOLDER").map(PathBuf::from),
67 ..Default::default()
68 })
69 });
70 let Some(record) = store
71 .as_ref()
72 .and_then(|s| s.get(uuid))
73 .or_else(|| standalone.and_then(|s| serde_json::from_str::<ExecutionRecord>(s).ok()))
74 else {
75 return Ok(());
76 };
77 if record.log_path.is_empty()
78 || record.options.get("sessionName").and_then(Value::as_str) != Some(name)
79 {
80 return Ok(());
81 }
82 let Some(config) = record
83 .options
84 .get("cpuPenaltyConfig")
85 .and_then(|v| serde_json::from_value::<Config>(v.clone()).ok())
86 else {
87 return Ok(());
88 };
89 let mut lock = LockManager::new(PathBuf::from(format!(
90 "{}.lock",
91 state_path(&record).display()
92 )));
93 if !lock.acquire(100) {
94 return Ok(());
95 }
96 let result = (|| {
97 let base = crate::docker_resource_limits::normalize_resource_limits(
98 record
99 .options
100 .get("baseResourceLimits")
101 .or_else(|| record.options.get("resourceLimits")),
102 );
103 let base_cpus = crate::docker_resource_limits::cpu_count(&base);
104 let saved = read_state(&record).and_then(|s| serde_json::from_value::<State>(s).ok());
105 let mut state = cpu_penalty::initial(base_cpus, Utc::now().timestamp_millis(), saved);
106 let interval =
107 (config.trigger_window_ms.min(config.release_window_ms) / 8).clamp(50, 30000);
108 let docker = crate::docker_cleanup::docker_command()
109 .to_string_lossy()
110 .into_owned();
111 let runner = SystemCommandRunner;
112 let started = Utc::now().timestamp_millis();
113 publish(&record, &state, name, attempt)?;
114 while let Some(current) = store
115 .as_ref()
116 .map(|s| s.get(uuid))
117 .unwrap_or_else(|| Some(record.clone()))
118 {
119 if current.options.get("sessionName").and_then(Value::as_str) != Some(name)
120 || current.attempt.as_ref().map(|a| a.number).unwrap_or(1) != attempt
121 {
122 break;
123 }
124 let live = runner.run(
125 &docker,
126 &[
127 "inspect".into(),
128 "-f".into(),
129 "{{.State.Running}}".into(),
130 name.into(),
131 ],
132 );
133 if (live.success && live.stdout.trim() == "false")
134 || (!live.success && Utc::now().timestamp_millis() - started > 30000)
135 {
136 break;
137 }
138 let info = runner.run(
139 &docker,
140 &["info".into(), "--format".into(), "{{json .}}".into()],
141 );
142 let stats = runner.run(
143 &docker,
144 &[
145 "stats".into(),
146 "--no-stream".into(),
147 "--format".into(),
148 "{{json .}}".into(),
149 ],
150 );
151 let cpus = if info.success {
152 serde_json::from_str::<Value>(&info.stdout)
153 .ok()
154 .and_then(|v| v["NCPU"].as_f64())
155 .unwrap_or(f64::NAN)
156 } else {
157 f64::NAN
158 };
159 let cores = if stats.success {
160 stats
161 .stdout
162 .lines()
163 .filter_map(|s| serde_json::from_str::<Value>(s).ok())
164 .find(|v| v["Name"].as_str() == Some(name))
165 .and_then(|v| {
166 v["CPUPerc"]
167 .as_str()
168 .and_then(|s| s.strip_suffix('%'))
169 .and_then(|s| s.parse::<f64>().ok())
170 })
171 .map(|p| p / 100.0)
172 .unwrap_or(f64::NAN)
173 } else {
174 f64::NAN
175 };
176 let (mut next, action) = cpu_penalty::evaluate(
177 &state,
178 Utc::now().timestamp_millis(),
179 cores,
180 cpus,
181 (interval * 3).max(5000),
182 &config,
183 );
184 if let Some(action) = action {
185 let update = runner.run(
186 &docker,
187 &[
188 vec!["update".into()],
189 crate::docker_resource_limits::cpu_update_args(&base, action.cpus),
190 vec![name.into()],
191 ]
192 .concat(),
193 );
194 if update.success {
195 crate::isolation::isolation_log::append_log_file(
196 &PathBuf::from(&record.log_path),
197 &format!(
198 "[start] CPU penalty {}: {} CPUs ({} of {} CPUs)\n",
199 if action.kind == "apply" {
200 "applied"
201 } else {
202 "lifted"
203 },
204 action.cpus,
205 action
206 .average
207 .map(|n| format!("{:.2}", n))
208 .unwrap_or_else(|| "capacity changed".into()),
209 action.capacity
210 ),
211 );
212 } else {
213 next.phase = state.phase.clone();
214 next.since = state.since;
215 next.limit_cpus = state.limit_cpus;
216 next.penalty_count = state.penalty_count;
217 next.samples.clear();
218 }
219 }
220 state = next;
221 publish(&record, &state, name, attempt)?;
222 thread::sleep(Duration::from_millis(interval as u64));
223 }
224 Ok(())
225 })();
226 lock.release();
227 result
228}
229pub fn main(args: &[String]) {
230 if args.len() < 3 {
231 return;
232 }
233 if let Err(error) = monitor(
234 &args[0],
235 &args[1],
236 args[2].parse().unwrap_or(1),
237 args.get(3).map(String::as_str),
238 ) {
239 if std::env::var("START_DEBUG").as_deref() == Ok("1") {
240 eprintln!("[DEBUG] CPU monitor: {}", error);
241 }
242 }
243}
244pub fn start_snippet(uuid: Option<&str>, name: &str, attempt: Option<u64>) -> String {
245 use crate::isolation::isolation_log::shell_quote;
246 let Some(uuid) = uuid else {
247 return String::new();
248 };
249 let exe = std::env::current_exe()
250 .ok()
251 .map(|p| p.to_string_lossy().into_owned())
252 .unwrap_or_else(|| "$".into());
253 format!(
254 "{} __start-cpu-monitor {} {} {} >/dev/null 2>&1 & __start_cpu_monitor=$!",
255 shell_quote(&exe),
256 shell_quote(uuid),
257 shell_quote(name),
258 attempt.unwrap_or(1)
259 )
260}
261pub fn stop_snippet() -> String {
262 "if [ -n \"$__start_cpu_monitor\" ]; then kill \"$__start_cpu_monitor\" 2>/dev/null; wait \"$__start_cpu_monitor\" 2>/dev/null; fi".into()
263}
264
265fn standalone_record(name: &str, options: &crate::isolation::IsolationOptions) -> Option<String> {
266 let config = options.cpu_penalty_config.as_ref()?;
267 let mut record = ExecutionRecord::new("CPU monitor");
268 record.log_path = options.log_path.as_ref()?.to_string_lossy().into();
269 record.options.insert("sessionName".into(), json!(name));
270 record
271 .options
272 .insert("cpuPenaltyConfig".into(), json!(config));
273 record
274 .options
275 .insert("baseResourceLimits".into(), json!(options.resource_limits));
276 serde_json::to_string(&record).ok()
277}
278pub fn start_snippet_with_options(
279 name: &str,
280 options: &crate::isolation::IsolationOptions,
281) -> String {
282 if options.execution_id.is_some() {
283 return start_snippet(options.execution_id.as_deref(), name, Some(1));
284 }
285 let Some(record) = standalone_record(name, options) else {
286 return String::new();
287 };
288 let exe = std::env::current_exe().unwrap_or_default();
289 use crate::isolation::isolation_log::shell_quote;
290 format!(
291 "{} __start-cpu-monitor '' {} 1 {} >/dev/null 2>&1 & __start_cpu_monitor=$!",
292 shell_quote(&exe.to_string_lossy()),
293 shell_quote(name),
294 shell_quote(&record)
295 )
296}
297pub fn spawn_standalone(name: &str, options: &crate::isolation::IsolationOptions) {
298 if let (Ok(exe), Some(record)) = (std::env::current_exe(), standalone_record(name, options)) {
299 let _ = std::process::Command::new(exe)
300 .args(["__start-cpu-monitor", "", name, "1", &record])
301 .stdout(std::process::Stdio::null())
302 .stderr(std::process::Stdio::null())
303 .spawn();
304 }
305}