Skip to main content

start_command/
snapshot_resume_options.rs

1//! Options for the guarded Docker snapshot resume (issue #224):
2//!
3//! ```text
4//! --snapshot-wait <seconds>          wait for the host snapshot lock
5//! --snapshot-disk-wait <seconds>     wait for free disk instead of refusing
6//! --remove-original-after <seconds>  remove the original only if the
7//!                                    replacement still runs after the delay
8//! --snapshot-size-hint <bytes>       skip the `docker inspect --size` walk
9//! --snapshot-prepare <command>       `docker exec` a command in the original
10//!                                    before it is committed
11//! ```
12//!
13//! They only apply to `--resume`. The two waits also read
14//! START_SNAPSHOT_LOCK_WAIT and START_SNAPSHOT_DISK_WAIT when the flags are
15//! absent, so a supervisor can set them once for every resume it runs.
16
17use crate::args_parser::WrapperOptions;
18use serde_json::{json, Value};
19
20/// Parsed snapshot resume options; `None` when the flag was not given.
21#[derive(Debug, Clone, Default, PartialEq)]
22pub struct SnapshotResumeOptions {
23    pub wait: Option<f64>,
24    pub disk_wait: Option<f64>,
25    pub remove_original_after: Option<f64>,
26    pub size_hint: Option<u64>,
27    pub prepare: Option<String>,
28}
29
30const FLAGS: [&str; 5] = [
31    "--snapshot-wait",
32    "--snapshot-disk-wait",
33    "--remove-original-after",
34    "--snapshot-size-hint",
35    "--snapshot-prepare",
36];
37
38/// JavaScript's `Number.MAX_SAFE_INTEGER`, so both editions accept the same hints.
39const MAX_SAFE_INTEGER: u64 = (1 << 53) - 1;
40
41/// Parse a non-negative number of seconds; fractions are allowed.
42pub fn parse_seconds(value: &str, name: &str) -> Result<f64, String> {
43    let text = value.trim();
44    let (whole, fraction) = text.split_once('.').unwrap_or((text, "0"));
45    let digits = |s: &str| !s.is_empty() && s.chars().all(|c| c.is_ascii_digit());
46    match text.parse::<f64>() {
47        Ok(seconds) if digits(whole) && digits(fraction) && seconds.is_finite() => Ok(seconds),
48        _ => Err(format!(
49            "Invalid {} value: \"{}\". Expected a non-negative number of seconds.",
50            name, value
51        )),
52    }
53}
54
55/// Parse a size in bytes; k, m, g and t suffixes are 1024-based.
56pub fn parse_size_hint(value: &str) -> Result<u64, String> {
57    let text = value.trim().to_lowercase();
58    let (number, unit) = match text.char_indices().last() {
59        Some((index, suffix)) if "bkmgt".contains(suffix) => (&text[..index], Some(suffix)),
60        _ => (text.as_str(), None),
61    };
62    let multiplier: u64 = match unit {
63        None | Some('b') => 1,
64        Some('k') => 1 << 10,
65        Some('m') => 1 << 20,
66        Some('g') => 1 << 30,
67        _ => 1 << 40,
68    };
69    number
70        .parse::<u64>()
71        .ok()
72        .filter(|_| !number.is_empty() && number.chars().all(|c| c.is_ascii_digit()))
73        .and_then(|n| n.checked_mul(multiplier))
74        .filter(|bytes| *bytes <= MAX_SAFE_INTEGER)
75        .ok_or_else(|| {
76            format!(
77                "Invalid --snapshot-size-hint value: \"{}\". Expected a number of bytes (k, m, g, t suffixes allowed).",
78                value
79            )
80        })
81}
82
83fn store(flag: &str, value: &str, options: &mut SnapshotResumeOptions) -> Result<(), String> {
84    match flag {
85        "--snapshot-size-hint" => options.size_hint = Some(parse_size_hint(value)?),
86        "--snapshot-prepare" => {
87            if value.trim().is_empty() {
88                return Err("--snapshot-prepare requires a non-empty command".to_string());
89            }
90            options.prepare = Some(value.to_string());
91        }
92        "--snapshot-wait" => options.wait = Some(parse_seconds(value, flag)?),
93        "--snapshot-disk-wait" => options.disk_wait = Some(parse_seconds(value, flag)?),
94        _ => options.remove_original_after = Some(parse_seconds(value, flag)?),
95    }
96    Ok(())
97}
98
99/// Parse one snapshot option; returns the arguments consumed, 0 when
100/// `args[index]` is not one.
101pub fn parse(args: &[String], index: usize, options: &mut WrapperOptions) -> Result<usize, String> {
102    let arg = args[index].as_str();
103    let (flag, inline) = match arg.split_once('=') {
104        Some((flag, value)) => (flag, Some(value)),
105        None => (arg, None),
106    };
107    if !FLAGS.contains(&flag) {
108        return Ok(0);
109    }
110    if let Some(value) = inline {
111        store(flag, value, &mut options.snapshot)?;
112        return Ok(1);
113    }
114    // A prepare command may itself start with a dash only through `=`.
115    let Some(value) = args.get(index + 1).filter(|value| !value.starts_with('-')) else {
116        let kind = match flag {
117            "--snapshot-size-hint" => "bytes",
118            "--snapshot-prepare" => "command",
119            _ => "seconds",
120        };
121        return Err(format!("Option {} requires a {} argument", flag, kind));
122    };
123    store(flag, value, &mut options.snapshot)?;
124    Ok(2)
125}
126
127/// Reject snapshot options outside `--resume` and fill the waits from the
128/// environment.
129pub fn validate(options: &mut WrapperOptions) -> Result<(), String> {
130    validate_with_env(options, &|name| std::env::var(name).ok())
131}
132
133pub fn validate_with_env(
134    options: &mut WrapperOptions,
135    env: &dyn Fn(&str) -> Option<String>,
136) -> Result<(), String> {
137    let snapshot = &mut options.snapshot;
138    if options.resume.is_none() {
139        let given = [
140            snapshot.wait.is_some(),
141            snapshot.disk_wait.is_some(),
142            snapshot.remove_original_after.is_some(),
143            snapshot.size_hint.is_some(),
144            snapshot.prepare.is_some(),
145        ];
146        return match FLAGS.iter().zip(given).find(|(_, given)| *given) {
147            Some((flag, _)) => Err(format!("{} is only valid with --resume", flag)),
148            None => Ok(()),
149        };
150    }
151    for (field, name, flag) in [
152        (
153            &mut snapshot.wait,
154            "START_SNAPSHOT_LOCK_WAIT",
155            "--snapshot-wait",
156        ),
157        (
158            &mut snapshot.disk_wait,
159            "START_SNAPSHOT_DISK_WAIT",
160            "--snapshot-disk-wait",
161        ),
162    ] {
163        if field.is_none() {
164            if let Some(value) = env(name).filter(|value| !value.is_empty()) {
165                *field = Some(parse_seconds(&value, &format!("{} ({})", name, flag))?);
166            }
167        }
168    }
169    if snapshot.remove_original_after.is_some() {
170        options.remove_original = true;
171    }
172    Ok(())
173}
174
175/// Add the options to the resume `overrides` object.
176pub fn insert_overrides(options: &SnapshotResumeOptions, overrides: &mut Value) {
177    overrides["snapshotWait"] = json!(options.wait);
178    overrides["snapshotDiskWait"] = json!(options.disk_wait);
179    overrides["removeOriginalAfter"] = json!(options.remove_original_after);
180    overrides["snapshotSizeHint"] = json!(options.size_hint);
181    overrides["snapshotPrepare"] = json!(options.prepare);
182}
183
184/// What a snapshot resume reads from its `overrides`.
185#[derive(Debug, Clone, Default, PartialEq)]
186pub struct SnapshotSettings {
187    pub wait_seconds: f64,
188    pub disk_wait_seconds: f64,
189    pub size_hint: Option<u64>,
190    pub prepare: Option<String>,
191    pub remove_original_after: f64,
192    /// Lock file instead of START_SNAPSHOT_LOCK (tests).
193    pub lock_path: Option<std::path::PathBuf>,
194    /// Disk poll interval instead of START_SNAPSHOT_DISK_POLL (tests).
195    pub disk_poll_seconds: Option<f64>,
196    /// Progress interval instead of START_SNAPSHOT_PROGRESS_INTERVAL (tests).
197    pub progress_interval_seconds: Option<f64>,
198}
199
200impl SnapshotSettings {
201    pub fn from_overrides(overrides: &Value) -> Self {
202        let seconds = |key: &str| overrides.get(key).and_then(Value::as_f64);
203        Self {
204            wait_seconds: seconds("snapshotWait").unwrap_or(0.0),
205            disk_wait_seconds: seconds("snapshotDiskWait").unwrap_or(0.0),
206            size_hint: overrides.get("snapshotSizeHint").and_then(Value::as_u64),
207            prepare: overrides
208                .get("snapshotPrepare")
209                .and_then(Value::as_str)
210                .map(String::from),
211            remove_original_after: seconds("removeOriginalAfter").unwrap_or(0.0),
212            lock_path: overrides
213                .get("snapshotLockPath")
214                .and_then(Value::as_str)
215                .map(std::path::PathBuf::from),
216            disk_poll_seconds: seconds("snapshotDiskPoll"),
217            progress_interval_seconds: seconds("snapshotProgressInterval"),
218        }
219    }
220}