start_command/
docker_resource_options.rs

1//! Resolve launch/resume resource specs using Docker daemon capacity, never local RAM.
2use crate::args_parser::WrapperOptions;
3use crate::docker_cleanup::docker_command;
4use crate::execution_control::CommandRunner;
5use serde_json::{json, Value};
6
7pub fn parse_spec(text: &str, cpu: bool) -> Result<(f64, f64, bool), String> {
8    let text = text.trim().to_lowercase();
9    let invalid = || format!("Invalid resource limit: {}", text);
10    if text.contains('%') {
11        let parts = text.split('-').collect::<Vec<_>>();
12        if parts.is_empty() || parts.len() > 2 {
13            return Err(invalid());
14        }
15        let values = parts
16            .iter()
17            .map(|s| {
18                s.strip_suffix('%')
19                    .ok_or_else(invalid)?
20                    .parse::<f64>()
21                    .map_err(|_| invalid())
22            })
23            .collect::<Result<Vec<_>, _>>()?;
24        let min = values[0];
25        let max = *values.last().unwrap();
26        return if min.is_finite() && max.is_finite() && min > 0.0 && max <= 100.0 && min <= max {
27            Ok((min, max, true))
28        } else {
29            Err(invalid())
30        };
31    }
32    let end = text
33        .find(|c: char| !c.is_ascii_digit() && c != '.')
34        .unwrap_or(text.len());
35    let number = text[..end].parse::<f64>().map_err(|_| invalid())?;
36    let suffix = &text[end..];
37    let multiplier = match suffix {
38        "" | "b" => 1.0,
39        "k" | "kb" => 1024.0,
40        "m" | "mb" => 1024.0_f64.powi(2),
41        "g" | "gb" => 1024.0_f64.powi(3),
42        "t" | "tb" => 1024.0_f64.powi(4),
43        _ => return Err(invalid()),
44    };
45    let amount = number * multiplier;
46    if (cpu && !suffix.is_empty())
47        || !amount.is_finite()
48        || amount <= 0.0
49        || amount > 9007199254740991.0
50    {
51        return Err(invalid());
52    }
53    Ok((amount, amount, false))
54}
55
56pub fn parse(args: &[String], index: usize, options: &mut WrapperOptions) -> Result<usize, String> {
57    let arg = &args[index];
58    let (flag, inline) = arg
59        .split_once('=')
60        .map(|(f, v)| (f, Some(v)))
61        .unwrap_or((arg, None));
62    let target = match flag {
63        "--memory" => &mut options.memory,
64        "--memory-swap" => &mut options.memory_swap,
65        "--cpus" => &mut options.cpus,
66        "--on-kill-resume-memory" => &mut options.on_kill_resume_memory,
67        _ => return Ok(0),
68    };
69    let value = inline
70        .or_else(|| args.get(index + 1).map(String::as_str))
71        .filter(|v| !v.is_empty() && !v.starts_with('-'))
72        .ok_or_else(|| format!("Option {} requires a resource limit", flag))?;
73    parse_spec(value, flag == "--cpus")?;
74    *target = Some(value.into());
75    Ok(if inline.is_some() { 1 } else { 2 })
76}
77
78pub fn validate(options: &WrapperOptions) -> Result<(), String> {
79    if options.memory.is_none()
80        && options.memory_swap.is_none()
81        && options.cpus.is_none()
82        && options.on_kill_resume_memory.is_none()
83    {
84        return Ok(());
85    }
86    if options.resume.is_none() && options.isolated.as_deref() != Some("docker") {
87        return Err("Resource options require Docker isolation or --resume".into());
88    }
89    if options.memory_swap.is_some() && options.memory.is_none() && options.resume.is_none() {
90        return Err("--memory-swap requires --memory".into());
91    }
92    if options.on_kill_resume_memory.is_some()
93        && options.on_kill_resume.is_none()
94        && options.recovery_command.is_none()
95    {
96        return Err(
97            "--on-kill-resume-memory requires --on-kill-resume or --recovery-command".into(),
98        );
99    }
100    Ok(())
101}
102
103pub fn random_fraction() -> f64 {
104    (uuid::Uuid::new_v4().as_u128() as u64 & ((1_u64 << 53) - 1)) as f64 / (1_u64 << 53) as f64
105}
106
107pub fn resolve<R: CommandRunner + ?Sized>(
108    specs: &Value,
109    base: &[String],
110    runner: &R,
111    random: impl Fn() -> f64,
112) -> Result<(Vec<String>, Value), String> {
113    let mut parsed = Vec::new();
114    for field in ["memory", "memorySwap", "cpus"] {
115        if let Some(value) = specs.get(field).and_then(Value::as_str) {
116            parsed.push((field, parse_spec(value, field == "cpus")?));
117        }
118    }
119    let capacity = if parsed.iter().any(|(_, s)| s.2) {
120        let info = runner.run(
121            &docker_command().to_string_lossy(),
122            &["info".into(), "--format".into(), "{{json .}}".into()],
123        );
124        if !info.success {
125            return Err("Cannot read resource capacity from Docker daemon".into());
126        }
127        serde_json::from_str::<Value>(&info.stdout).map_err(|e| e.to_string())?
128    } else {
129        Value::Null
130    };
131    let mut values = json!({});
132    let mut limits = base.to_vec();
133    for (field, (min, max, percent)) in parsed {
134        let amount = if percent {
135            capacity
136                .get(if field == "cpus" { "NCPU" } else { "MemTotal" })
137                .and_then(Value::as_f64)
138                .filter(|n| *n > 0.0)
139                .ok_or("Docker daemon did not report MemTotal/NCPU")?
140                * (min + (max - min) * random())
141                / 100.0
142        } else {
143            min
144        };
145        let amount = if field == "cpus" {
146            (amount * 1e9).floor() / 1e9
147        } else {
148            amount.floor()
149        };
150        if amount
151            < if field == "cpus" {
152                0.01
153            } else {
154                6.0 * 1024.0 * 1024.0
155            }
156        {
157            return Err(format!("Resolved {} is below Docker's minimum", field));
158        }
159        values[field] = if field == "cpus" {
160            json!(amount)
161        } else {
162            json!(amount as u64)
163        };
164    }
165    if values.get("memory").is_some() && values.get("memorySwap").is_none() {
166        values["memorySwap"] = values["memory"].clone();
167    }
168    let memory = values.get("memory").and_then(Value::as_f64).or_else(|| {
169        limits
170            .iter()
171            .find_map(|s| s.strip_prefix("--memory="))
172            .and_then(|s| parse_spec(s, false).ok())
173            .map(|s| s.0)
174    });
175    if let Some(swap) = values.get("memorySwap").and_then(Value::as_f64) {
176        if memory.is_none_or(|m| swap < m) {
177            return Err("memory-swap must be at least memory".into());
178        }
179    }
180    for field in ["memory", "memorySwap", "cpus"] {
181        if let Some(value) = values.get(field) {
182            let flag = if field == "memorySwap" {
183                "--memory-swap"
184            } else if field == "memory" {
185                "--memory"
186            } else {
187                "--cpus"
188            };
189            limits.retain(|s| {
190                !s.starts_with(&format!("{}=", flag))
191                    && !(field == "cpus"
192                        && (s.starts_with("--cpu-quota=") || s.starts_with("--cpu-period=")))
193            });
194            limits.push(format!("{}={}", flag, value.as_f64().unwrap()));
195        }
196    }
197    Ok((limits, values))
198}
199
200pub fn specs(options: &WrapperOptions) -> Value {
201    json!({"memory":options.memory,"memorySwap":options.memory_swap,"cpus":options.cpus})
202}
203
204pub fn prepare(options: &mut WrapperOptions) -> Result<(), String> {
205    let (limits, values) = resolve(
206        &specs(options),
207        &options.resource_limits,
208        &crate::execution_control::SystemCommandRunner,
209        random_fraction,
210    )?;
211    options.resource_limits = limits;
212    options.resolved_limits = values;
213    Ok(())
214}
215
216pub fn log_line(specs: &Value, values: &Value) -> Option<String> {
217    let entries = values.as_object()?;
218    if entries.is_empty() {
219        return None;
220    }
221    Some(format!(
222        "Limits: {}",
223        entries
224            .iter()
225            .map(|(k, v)| format!(
226                "{}={} ({} -> {})",
227                k,
228                v,
229                specs.get(k).and_then(Value::as_str).unwrap_or("default"),
230                v
231            ))
232            .collect::<Vec<_>>()
233            .join(" ")
234    ))
235}
start_command/docker_resource_options.rs

start_command/
docker_resource_options.rs
