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

start_command/
resume_resources.rs
