start_command/
resume_resources.rs

1use crate::docker_resource_limits::{normalize_resource_limits, read_docker_resource_limits};
2use crate::execution_control::CommandRunner;
3use crate::execution_store::ExecutionRecord;
4use crate::session_probe::{SessionProbe, SessionState};
5use serde_json::{json, Value};
6
7pub fn restore_base_cpu(mut limits: Vec<String>, record: &ExecutionRecord) -> Vec<String> {
8    if record.options.contains_key("cpuPenaltyConfig") {
9        let base = normalize_resource_limits(
10            record
11                .options
12                .get("baseResourceLimits")
13                .or_else(|| record.options.get("resourceLimits")),
14        );
15        let cpu = |s: &String| {
16            s.starts_with("--cpus=")
17                || s.starts_with("--cpu-quota=")
18                || s.starts_with("--cpu-period=")
19        };
20        limits.retain(|s| !cpu(s));
21        limits.extend(base.into_iter().filter(cpu));
22    }
23    limits
24}
25
26#[allow(clippy::type_complexity)]
27pub fn prepare<R: CommandRunner>(
28    record: &ExecutionRecord,
29    probe: &SessionProbe,
30    overrides: &Value,
31    runner: &R,
32) -> Result<(Option<Vec<String>>, Option<Value>, Value), String> {
33    if record.options.get("isolated").and_then(Value::as_str) != Some("docker") {
34        if ["memory", "memorySwap", "cpus"]
35            .iter()
36            .any(|k| overrides.get(k).and_then(Value::as_str).is_some())
37        {
38            return Err("Resource overrides require a Docker execution".into());
39        }
40        return Ok((None, None, Value::Null));
41    }
42    let live = if probe.state == SessionState::Stopped {
43        record
44            .options
45            .get("sessionName")
46            .and_then(Value::as_str)
47            .and_then(|name| read_docker_resource_limits(name, runner))
48    } else {
49        None
50    };
51    let base = restore_base_cpu(
52        live.unwrap_or_else(|| normalize_resource_limits(record.options.get("resourceLimits"))),
53        record,
54    );
55    let old_memory = base
56        .iter()
57        .find_map(|s| {
58            s.strip_prefix("--memory=")
59                .and_then(|s| crate::docker_resource_options::parse_spec(s, false).ok())
60                .map(|v| json!(v.0 as u64))
61        })
62        .unwrap_or(json!("unlimited"));
63    let (mut limits, values) = crate::docker_resource_options::resolve(
64        overrides,
65        &base,
66        runner,
67        crate::docker_resource_options::random_fraction,
68    )?;
69    if overrides.get("cpus").and_then(Value::as_str).is_some()
70        && probe.state == SessionState::Stopped
71    {
72        // Preserve the existing CPU representation for an in-place update.
73        let cpu =
74            crate::docker_resource_limits::with_cpu_limit(base, values["cpus"].as_f64().unwrap());
75        let is_cpu = |s: &String| {
76            ["--cpus=", "--cpu-quota=", "--cpu-period="]
77                .iter()
78                .any(|p| s.starts_with(p))
79        };
80        limits.retain(|s| !is_cpu(s));
81        limits.extend(cpu.into_iter().filter(is_cpu));
82    }
83    if record.options.contains_key("cpuPenaltyConfig")
84        && overrides.get("cpus").and_then(Value::as_str).is_none()
85    {
86        limits = clamp_cpu_to_daemon(limits, runner)?;
87    }
88    let mut resolved = record
89        .options
90        .get("resolvedLimits")
91        .filter(|value| value.is_object())
92        .cloned()
93        .unwrap_or_else(|| json!({}));
94    for (k, v) in values.as_object().unwrap() {
95        resolved[k] = v.clone();
96    }
97    if record.options.contains_key("cpuPenaltyConfig") {
98        if let Some(cpu) = crate::docker_resource_limits::cpu_count(&limits) {
99            resolved["cpus"] = json!(cpu);
100        }
101    }
102
103    Ok((Some(limits), Some(resolved), old_memory))
104}
105
106pub fn clamp_cpu_to_daemon<R: CommandRunner + ?Sized>(
107    limits: Vec<String>,
108    runner: &R,
109) -> Result<Vec<String>, String> {
110    let info = runner.run(
111        &crate::docker_cleanup::docker_command().to_string_lossy(),
112        &["info".into(), "--format".into(), "{{json .}}".into()],
113    );
114    let capacity = info
115        .success
116        .then(|| serde_json::from_str::<Value>(&info.stdout).ok())
117        .flatten()
118        .and_then(|v| v["NCPU"].as_f64())
119        .filter(|n| *n > 0.0)
120        .ok_or("Cannot restore base CPU capacity")?;
121    let base = crate::docker_resource_limits::cpu_count(&limits).unwrap_or(capacity);
122    Ok(crate::docker_resource_limits::with_cpu_limit(
123        limits,
124        base.min(capacity),
125    ))
126}
start_command/resume_resources.rs

start_command/
resume_resources.rs
