diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index 08139f1a7a..350205082b 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -5,6 +5,8 @@ use serde_json::value::RawValue; use std::{ cmp::Reverse, collections::{HashMap, HashSet}, + path::Path, + str::FromStr, sync::{atomic::AtomicBool, Arc}, }; use tokio::sync::RwLock; @@ -15,7 +17,10 @@ lazy_static::lazy_static! { pub static ref WORKER_GROUP: String = std::env::var("WORKER_GROUP").unwrap_or_else(|_| "default".to_string()); pub static ref NO_LOGS: bool = std::env::var("NO_LOGS").ok().is_some_and(|x| x == "1" || x == "true"); - + pub static ref CGROUP_V2_PATH_RE: Regex = Regex::new(r#"(?m)^0::(/.*)$"#).unwrap(); + pub static ref CGROUP_V2_CPU_RE: Regex = Regex::new(r#"(?m)^(\d+) \S+$"#).unwrap(); + pub static ref CGROUP_V1_INACTIVE_FILE_RE: Regex = Regex::new(r#"(?m)^total_inactive_file (\d+)$"#).unwrap(); + pub static ref CGROUP_V2_INACTIVE_FILE_RE: Regex = Regex::new(r#"(?m)^inactive_file (\d+)$"#).unwrap(); pub static ref DEFAULT_TAGS: Vec = vec![ "deno".to_string(), @@ -132,156 +137,101 @@ fn process_custom_tags(tags: Vec) -> (Vec, HashMap Option { - let mut vcpus = std::process::Command::new("cat") - .args(["/sys/fs/cgroup/cpu.max"]) +fn parse_file(path: &str) -> Option { + std::process::Command::new("cat") + .args([path]) .output() .ok() .map(|o| { String::from_utf8_lossy(&o.stdout) .to_string() - .split(" ") - .map(|s| s.to_string()) - .collect::>() - .get(0) - .map(|s| s.to_string().trim().parse::().ok()) - .flatten() + .trim() + .parse::() + .ok() }) .flatten() - .map(|x| if x > 0 { Some(x) } else { None }) - .flatten(); +} - if vcpus.is_none() { - vcpus = std::process::Command::new("cat") - .args(["/sys/fs/cgroup/cpu/cpu.cfs_quota_us"]) - .output() - .ok() - .map(|o| { - String::from_utf8_lossy(&o.stdout) - .to_string() - .trim() - .parse::() - .ok() +fn get_cgroupv2_path() -> Option { + let cgroup_path: String = parse_file("/proc/self/cgroup")?; + + CGROUP_V2_PATH_RE + .captures(&cgroup_path) + .map(|x| format!("/sys/fs/cgroup{}", x.get(1).unwrap().as_str())) +} + +pub fn get_vcpus() -> Option { + if Path::new("/sys/fs/cgroup/cpu/cpu.cfs_quota_us").exists() { + // cgroup v1 + parse_file("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") + .map(|x: i64| if x > 0 { Some(x) } else { None }) + .flatten() + } else { + // cgroup v2 + let cgroup_path = get_cgroupv2_path()?; + + let cpu_max_path = format!("{cgroup_path}/cpu.max"); + + parse_file(&cpu_max_path) + .map(|x: String| { + CGROUP_V2_CPU_RE + .captures(&x) + .map(|x| x.get(1).unwrap().as_str().parse::().ok()) + .flatten() }) .flatten() .map(|x| if x > 0 { Some(x) } else { None }) - .flatten(); + .flatten() } - - vcpus } pub fn get_memory() -> Option { - let mut memory = std::process::Command::new("cat") - .args(["/sys/fs/cgroup/memory.max"]) - .output() - .ok() - .map(|o| { - String::from_utf8_lossy(&o.stdout) - .to_string() - .trim() - .parse::() - .ok() - }) - .flatten(); + if Path::new("/sys/fs/cgroup/memory/memory.max").exists() { + // cgroup v1 + parse_file("/sys/fs/cgroup/memory/memory.max") + } else { + // cgroup v2 + let cgroup_path = get_cgroupv2_path()?; + let memory_max_path = format!("{cgroup_path}/memory.max"); - if memory.is_none() { - memory = std::process::Command::new("cat") - .args(["/sys/fs/cgroup/memory/memory.limit_in_bytes"]) - .output() - .ok() - .map(|o| { - String::from_utf8_lossy(&o.stdout) - .to_string() - .trim() - .parse::() - .ok() - }) - .flatten() + parse_file(&memory_max_path) } - - memory } pub fn get_worker_memory_usage() -> Option { - let mut total_memory_usage = std::process::Command::new("cat") - .args(["/sys/fs/cgroup/memory.current"]) - .output() - .ok() - .map(|o| { - String::from_utf8_lossy(&o.stdout) - .to_string() - .trim() - .parse::() - .ok() - }) - .flatten(); + if Path::new("/sys/fs/cgroup/memory/memory.usage_in_bytes").exists() { + // cgroup v1 + let total_memory_usage: i64 = parse_file("/sys/fs/cgroup/memory/memory.usage_in_bytes")?; - if total_memory_usage.is_none() { - total_memory_usage = std::process::Command::new("cat") - .args(["/sys/fs/cgroup/memory/memory.usage_in_bytes"]) - .output() - .ok() - .map(|o| { - String::from_utf8_lossy(&o.stdout) - .to_string() - .trim() - .parse::() - .ok() + let inactive_file = parse_file("/sys/fs/cgroup/memory/memory.stat") + .map(|x: String| { + CGROUP_V1_INACTIVE_FILE_RE + .captures(&x) + .map(|x| x.get(1).unwrap().as_str().parse::().ok()) + .flatten() }) - .flatten() - } + .flatten()?; - match total_memory_usage { - Some(total_memory_usage) => { - let mut inactive_file = std::process::Command::new("cat") - .args(["/sys/fs/cgroup/memory.stat"]) - .output() - .ok() - .map(|o| { - String::from_utf8_lossy(&o.stdout) - .to_string() - .split("\n") - .find(|s| s.starts_with("inactive_file")) - .map(|s| { - s.split(" ") - .collect::>() - .get(1) - .map(|s| s.trim().parse::().ok()) - .flatten() - }) - .flatten() - }) - .flatten(); + Some(total_memory_usage - inactive_file) + } else { + // cgroup v2 + let cgroup_path = get_cgroupv2_path()?; + let memory_current_path = format!("{cgroup_path}/memory.current"); - if inactive_file.is_none() { - inactive_file = std::process::Command::new("cat") - .args(["/sys/fs/cgroup/memory/memory.stat"]) - .output() - .ok() - .map(|o| { - String::from_utf8_lossy(&o.stdout) - .to_string() - .split("\n") - .find(|s| s.starts_with("total_inactive_file")) - .map(|s| { - s.split(" ") - .collect::>() - .get(1) - .map(|s| s.trim().parse::().ok()) - .flatten() - }) - .flatten() - }) - .flatten(); - } + let total_memory_usage: i64 = parse_file(&memory_current_path)?; - match inactive_file { - Some(inactive_file) => Some(total_memory_usage - inactive_file), - None => None, - } - } - None => None, + let memory_stat_path = format!("{cgroup_path}/memory.stat"); + + let inactive_file = parse_file(&memory_stat_path) + .map(|x: String| { + CGROUP_V2_INACTIVE_FILE_RE + .captures(&x) + .map(|x| x.get(1).unwrap().as_str().parse::().ok()) + .flatten() + }) + .flatten()?; + + Some(total_memory_usage - inactive_file) } }