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        "--memory-reservation" => &mut options.memory_reservation,
66        "--cpus" => &mut options.cpus,
67        "--on-kill-resume-memory" => &mut options.on_kill_resume_memory,
68        _ => return Ok(0),
69    };
70    let value = inline
71        .or_else(|| args.get(index + 1).map(String::as_str))
72        .filter(|v| !v.is_empty() && !v.starts_with('-'))
73        .ok_or_else(|| format!("Option {} requires a resource limit", flag))?;
74    parse_spec(value, flag == "--cpus")?;
75    *target = Some(value.into());
76    Ok(if inline.is_some() { 1 } else { 2 })
77}
78
79pub fn validate(options: &WrapperOptions) -> Result<(), String> {
80    if options.memory.is_none()
81        && options.memory_swap.is_none()
82        && options.memory_reservation.is_none()
83        && options.cpus.is_none()
84        && options.on_kill_resume_memory.is_none()
85    {
86        return Ok(());
87    }
88    if options.resume.is_none() && options.isolated.as_deref() != Some("docker") {
89        return Err("Resource options require Docker isolation or --resume".into());
90    }
91    if options.memory_swap.is_some() && options.memory.is_none() && options.resume.is_none() {
92        return Err("--memory-swap requires --memory".into());
93    }
94    if options.on_kill_resume_memory.is_some()
95        && options.on_kill_resume.is_none()
96        && options.recovery_command.is_none()
97    {
98        return Err(
99            "--on-kill-resume-memory requires --on-kill-resume or --recovery-command".into(),
100        );
101    }
102    Ok(())
103}
104
105pub fn random_fraction() -> f64 {
106    (uuid::Uuid::new_v4().as_u128() as u64 & ((1_u64 << 53) - 1)) as f64 / (1_u64 << 53) as f64
107}
108
109pub fn resolve<R: CommandRunner + ?Sized>(
110    specs: &Value,
111    base: &[String],
112    runner: &R,
113    random: impl Fn() -> f64,
114) -> Result<(Vec<String>, Value), String> {
115    let mut parsed = Vec::new();
116    for field in ["memory", "memorySwap", "memoryReservation", "cpus"] {
117        if let Some(value) = specs.get(field).and_then(Value::as_str) {
118            parsed.push((field, parse_spec(value, field == "cpus")?));
119        }
120    }
121    let capacity = if parsed.iter().any(|(_, s)| s.2) {
122        let info = runner.run(
123            &docker_command().to_string_lossy(),
124            &["info".into(), "--format".into(), "{{json .}}".into()],
125        );
126        if !info.success {
127            return Err("Cannot read resource capacity from Docker daemon".into());
128        }
129        serde_json::from_str::<Value>(&info.stdout).map_err(|e| e.to_string())?
130    } else {
131        Value::Null
132    };
133    let mut values = json!({});
134    let mut limits = base.to_vec();
135    for (field, (min, max, percent)) in parsed {
136        let amount = if percent {
137            capacity
138                .get(if field == "cpus" { "NCPU" } else { "MemTotal" })
139                .and_then(Value::as_f64)
140                .filter(|n| *n > 0.0)
141                .ok_or("Docker daemon did not report MemTotal/NCPU")?
142                * (min + (max - min) * random())
143                / 100.0
144        } else {
145            min
146        };
147        let amount = if field == "cpus" {
148            (amount * 1e9).floor() / 1e9
149        } else {
150            amount.floor()
151        };
152        if amount
153            < if field == "cpus" {
154                0.01
155            } else {
156                6.0 * 1024.0 * 1024.0
157            }
158        {
159            return Err(format!("Resolved {} is below Docker's minimum", field));
160        }
161        values[field] = if field == "cpus" {
162            json!(amount)
163        } else {
164            json!(amount as u64)
165        };
166    }
167    if values.get("memory").is_some() && values.get("memorySwap").is_none() {
168        values["memorySwap"] = values["memory"].clone();
169    }
170    let memory = values.get("memory").and_then(Value::as_f64).or_else(|| {
171        limits
172            .iter()
173            .find_map(|s| s.strip_prefix("--memory="))
174            .and_then(|s| parse_spec(s, false).ok())
175            .map(|s| s.0)
176    });
177    if let Some(swap) = values.get("memorySwap").and_then(Value::as_f64) {
178        if memory.is_none_or(|m| swap < m) {
179            return Err("memory-swap must be at least memory".into());
180        }
181    }
182    let reservation = values
183        .get("memoryReservation")
184        .and_then(Value::as_f64)
185        .or_else(|| {
186            limits
187                .iter()
188                .find_map(|s| s.strip_prefix("--memory-reservation="))
189                .and_then(|s| parse_spec(s, false).ok())
190                .map(|s| s.0)
191        });
192    if memory.zip(reservation).is_some_and(|(m, r)| r >= m) {
193        return Err("memory-reservation must be less than memory".into());
194    }
195    for field in ["memory", "memorySwap", "memoryReservation", "cpus"] {
196        if let Some(value) = values.get(field) {
197            let flag = if field == "memorySwap" {
198                "--memory-swap"
199            } else if field == "memoryReservation" {
200                "--memory-reservation"
201            } else if field == "memory" {
202                "--memory"
203            } else {
204                "--cpus"
205            };
206            limits.retain(|s| {
207                !s.starts_with(&format!("{}=", flag))
208                    && !(field == "cpus"
209                        && (s.starts_with("--cpu-quota=") || s.starts_with("--cpu-period=")))
210            });
211            limits.push(format!("{}={}", flag, value.as_f64().unwrap()));
212        }
213    }
214    Ok((limits, values))
215}
216
217pub fn specs(options: &WrapperOptions) -> Value {
218    json!({"memory":options.memory,"memorySwap":options.memory_swap,"memoryReservation":options.memory_reservation,"cpus":options.cpus})
219}
220
221pub fn prepare(options: &mut WrapperOptions) -> Result<(), String> {
222    let (limits, values) = resolve(
223        &specs(options),
224        &options.resource_limits,
225        &crate::execution_control::SystemCommandRunner,
226        random_fraction,
227    )?;
228    options.resource_limits = limits;
229    options.resolved_limits = values;
230    Ok(())
231}
232
233pub fn log_line(specs: &Value, values: &Value) -> Option<String> {
234    let entries = values.as_object()?;
235    if entries.is_empty() {
236        return None;
237    }
238    Some(format!(
239        "Limits: {}",
240        entries
241            .iter()
242            .map(|(k, v)| format!(
243                "{}={} ({} -> {})",
244                k,
245                v,
246                specs.get(k).and_then(Value::as_str).unwrap_or("default"),
247                v
248            ))
249            .collect::<Vec<_>>()
250            .join(" ")
251    ))
252}
start_command/docker_resource_options.rs

start_command/
docker_resource_options.rs
