use serde_json::Value;
use crate::docker_cleanup::docker_command;
use crate::execution_control::CommandRunner;
pub const DEFAULT_SHM_SIZE: i64 = 64 * 1024 * 1024;
const KIB: i64 = 1024;
const MIB: i64 = KIB * 1024;
const GIB: i64 = MIB * 1024;
pub fn format_docker_bytes(bytes: i64) -> String {
for (unit, size) in [("g", GIB), ("m", MIB), ("k", KIB)] {
if bytes >= size && bytes % size == 0 {
return format!("{}{}", bytes / size, unit);
}
}
bytes.to_string()
}
pub fn format_docker_cpus(nano_cpus: i64) -> String {
let text = format!("{:.9}", nano_cpus as f64 / 1e9);
let text = text.trim_end_matches('0').trim_end_matches('.');
text.to_string()
}
fn number(value: Option<&Value>) -> Option<i64> {
match value? {
Value::Number(n) => n.as_i64().or_else(|| n.as_f64().map(|f| f as i64)),
Value::String(s) => s.trim().parse().ok(),
_ => None,
}
}
fn positive(value: Option<&Value>) -> Option<i64> {
number(value).filter(|n| *n > 0)
}
fn non_empty(value: Option<&Value>) -> Option<String> {
value
.and_then(Value::as_str)
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::to_string)
}
fn scalar_text(value: Option<&Value>) -> String {
match value {
Some(Value::String(s)) => s.clone(),
Some(Value::Number(n)) => n.to_string(),
Some(Value::Bool(b)) => b.to_string(),
_ => String::new(),
}
}
pub fn parse_docker_resource_limits(host_config: &Value) -> Vec<String> {
let Some(h) = host_config.as_object() else {
return Vec::new();
};
let mut limits = Vec::new();
if let Some(memory) = positive(h.get("Memory")) {
limits.push(format!("--memory={}", format_docker_bytes(memory)));
match number(h.get("MemorySwap")) {
Some(-1) => limits.push("--memory-swap=-1".to_string()),
Some(swap) if swap > 0 => {
limits.push(format!("--memory-swap={}", format_docker_bytes(swap)))
}
_ => {}
}
}
if let Some(reservation) = positive(h.get("MemoryReservation")) {
limits.push(format!(
"--memory-reservation={}",
format_docker_bytes(reservation)
));
}
if let Some(nano_cpus) = positive(h.get("NanoCpus")) {
limits.push(format!("--cpus={}", format_docker_cpus(nano_cpus)));
} else {
if let Some(quota) = positive(h.get("CpuQuota")) {
limits.push(format!("--cpu-quota={}", quota));
}
if let Some(period) = positive(h.get("CpuPeriod")) {
limits.push(format!("--cpu-period={}", period));
}
}
if let Some(shares) = positive(h.get("CpuShares")) {
limits.push(format!("--cpu-shares={}", shares));
}
if let Some(cpus) = non_empty(h.get("CpusetCpus")) {
limits.push(format!("--cpuset-cpus={}", cpus));
}
if let Some(mems) = non_empty(h.get("CpusetMems")) {
limits.push(format!("--cpuset-mems={}", mems));
}
if let Some(pids) = positive(h.get("PidsLimit")) {
limits.push(format!("--pids-limit={}", pids));
}
if let Some(shm) = positive(h.get("ShmSize")) {
if shm != DEFAULT_SHM_SIZE {
limits.push(format!("--shm-size={}", format_docker_bytes(shm)));
}
}
if let Some(storage) = h.get("StorageOpt").and_then(Value::as_object) {
let mut keys: Vec<&String> = storage.keys().collect();
keys.sort();
for key in keys {
let value = scalar_text(storage.get(key));
if !value.trim().is_empty() {
limits.push(format!("--storage-opt={}={}", key, value.trim()));
}
}
}
if let Some(ulimits) = h.get("Ulimits").and_then(Value::as_array) {
for ulimit in ulimits {
if let Some(name) = non_empty(ulimit.get("Name")) {
limits.push(format!(
"--ulimit={}={}:{}",
name,
scalar_text(ulimit.get("Soft")),
scalar_text(ulimit.get("Hard"))
));
}
}
}
limits
}
pub fn read_docker_resource_limits<R: CommandRunner + ?Sized>(
container_name: &str,
runner: &R,
) -> Option<Vec<String>> {
let docker = docker_command().to_string_lossy().to_string();
let args: Vec<String> = [
"inspect",
"--format",
"{{json .HostConfig}}",
container_name,
]
.iter()
.map(|s| s.to_string())
.collect();
let result = runner.run(&docker, &args);
if !result.success {
return None;
}
let host_config: Value = serde_json::from_str(result.stdout.trim()).ok()?;
Some(parse_docker_resource_limits(&host_config))
}
pub fn normalize_resource_limits(value: Option<&Value>) -> Vec<String> {
match value {
Some(Value::Array(items)) => items
.iter()
.map(|item| match item {
Value::String(s) => s.clone(),
other => other.to_string(),
})
.filter(|flag| flag.starts_with("--"))
.collect(),
Some(Value::String(text)) => text
.split_whitespace()
.filter(|flag| flag.starts_with("--"))
.map(str::to_string)
.collect(),
_ => Vec::new(),
}
}
pub fn build_resource_limits_status_line(limits: &[String]) -> Option<String> {
if limits.is_empty() {
None
} else {
Some(format!("[Isolation] Resource limits: {}", limits.join(" ")))
}
}
pub fn cpu_count(limits: &[String]) -> Option<f64> {
let get = |prefix: &str| {
limits.iter().find_map(|s| {
s.strip_prefix(prefix)
.and_then(|n| n.parse::<f64>().ok())
.filter(|n| *n > 0.0)
})
};
get("--cpus=")
.or_else(|| get("--cpu-quota=").map(|q| q / get("--cpu-period=").unwrap_or(100000.0)))
}
pub fn cpu_update_args(limits: &[String], cpus: f64) -> Vec<String> {
if !limits.iter().any(|s| s.starts_with("--cpus="))
&& limits
.iter()
.any(|s| s.starts_with("--cpu-quota=") || s.starts_with("--cpu-period="))
{
vec![
"--cpu-period=100000".into(),
format!("--cpu-quota={}", (cpus * 100000.0).ceil() as i64),
]
} else {
vec!["--cpus".into(), cpus.to_string()]
}
}
pub fn with_cpu_limit(mut limits: Vec<String>, cpus: f64) -> Vec<String> {
let args = cpu_update_args(&limits, cpus);
limits.retain(|s| {
!["--cpus=", "--cpu-quota=", "--cpu-period="]
.iter()
.any(|p| s.starts_with(p))
});
if args[0] == "--cpus" {
limits.push(format!("--cpus={}", cpus));
} else {
limits.extend(args);
}
limits
}