Checkpoints: a handed-back job resumes (refs #77) f201e1bc5bcd

//! Checkpoints (oxygenflow#77, docs/handover.md §7): the directories a job

//! declares with `checkpoint "/src/…"`, saved when its runner hands it back

//! and restored before the next attempt runs the same exec, so a tool that

//! skips finished work (cargo, make) carries on where it stopped.

//!

//! The format is the runner's own, so CI builds need no new crate: a magic

//! line, then entries (a kind byte, the path under /src, mode, modification

//! time and, for a file, its bytes). Files keep their times, since cargo's

//! fingerprints read them. Links are left out. Unpacking accepts only paths

//! under the job's own checkpoints, without `.` or `..`.

use anyhow::{bail, Context, Result};

use flow_core::pipeline::{Job, MountSource};

use std::path::{Path, PathBuf};

use std::time::{Duration, SystemTime, UNIX_EPOCH};

/// At most this much is saved: the hub keeps it as it keeps an artifact.

pub const MAX_BYTES: usize = 512 << 20;

const MAGIC: &[u8] = b"OXYFLOW-CHECKPOINT 1\n";

const DIR: u8 = 0;

const FILE: u8 = 1;

/// What was saved or restored.

#[derive(Debug, Default, PartialEq, Eq)]

pub struct Counts {

    pub files: usize,

    pub dirs: usize,

    pub bytes: u64,

}

/// Where checkpoint `path` (under /src) is on this host, for a job of a

/// directory tier (WASM, syscall): in a disk mount's directory when a mount

/// holds it, else in the job's tree. `None` for one under a tmpfs, which

/// the syscall tier keeps in memory.

pub fn host_dir(path: &str, job: &Job, tree: &Path, scratch: &Path, wasm: bool) -> Option<PathBuf> {

    let under = |t: &str| path == t || path.starts_with(&format!("{}/", t.trim_end_matches('/')));

    let mut disks = 0;

    let mut best: Option<(usize, PathBuf)> = None;

    for (i, m) in job.mounts.iter().enumerate() {

        let dir = match m.source {

            MountSource::Disk => {

                disks += 1;

                Some(crate::syscall::disk_dir(scratch, disks - 1))

            }

            MountSource::Tmpfs if wasm => Some(scratch.join(format!("tmpfs-{i}"))),

            MountSource::Tmpfs => None,

        };

        if under(&m.target) && best.as_ref().is_none_or(|(len, _)| m.target.len() > *len) {

            let rest = path[m.target.trim_end_matches('/').len()..].trim_start_matches('/');

            best = Some((m.target.len(), dir.map(|d| d.join(rest)).unwrap_or_default()));

        }

    }

    match best {

        Some((_, d)) if d.as_os_str().is_empty() => None,

        Some((_, d)) => Some(d),

        None => Some(tree.join(path.trim_start_matches("/src").trim_start_matches('/'))),

    }

}

/// Pack `dirs` (each a checkpoint path and where it is on this host).

pub fn pack(dirs: &[(String, PathBuf)]) -> Result<(Vec<u8>, Counts)> {

    let mut out = MAGIC.to_vec();

    let mut n = Counts::default();

    for (path, host) in dirs {

        if !host.is_dir() {

            continue;

        }

        let mut stack = vec![(host.clone(), path.trim_end_matches('/').to_string())];

        while let Some((at, name)) = stack.pop() {

            let meta = std::fs::symlink_metadata(&at)?;

            if meta.is_dir() {

                entry(&mut out, DIR, &name, &meta, &[]);

                n.dirs += 1;

                let mut children: Vec<_> = std::fs::read_dir(&at)?.collect::<std::io::Result<_>>()?;

                children.sort_by_key(|e| e.file_name());

                for c in children.into_iter().rev() {

                    let Some(child) = c.file_name().to_str().map(String::from) else { continue };

                    stack.push((c.path(), format!("{name}/{child}")));

                }

            } else if meta.is_file() {

                let bytes = std::fs::read(&at).with_context(|| format!("read {}", at.display()))?;

                n.files += 1;

                n.bytes += bytes.len() as u64;

                entry(&mut out, FILE, &name, &meta, &bytes);

            }

            if out.len() > MAX_BYTES {

                bail!("the checkpoint is over {} MiB", MAX_BYTES >> 20);

            }

        }

    }

    Ok((out, n))

}

fn entry(out: &mut Vec<u8>, kind: u8, path: &str, meta: &std::fs::Metadata, bytes: &[u8]) {

    use std::os::unix::fs::PermissionsExt;

    let since = meta.modified().ok().and_then(|t| t.duration_since(UNIX_EPOCH).ok()).unwrap_or_default();

    out.push(kind);

    out.extend((path.len() as u32).to_le_bytes());

    out.extend(path.as_bytes());

    out.extend(meta.permissions().mode().to_le_bytes());

    out.extend(since.as_secs().to_le_bytes());

    out.extend(since.subsec_nanos().to_le_bytes());

    out.extend((bytes.len() as u64).to_le_bytes());

    out.extend(bytes);

}

/// Restore a checkpoint: each entry at `place(checkpoint)` joined with the

/// rest of its path, for the job's own `checkpoints` only.

pub fn unpack(bytes: &[u8], checkpoints: &[String], place: impl Fn(&str) -> Option<PathBuf>) -> Result<Counts> {

    use std::os::unix::fs::PermissionsExt;

    let Some(mut rest) = bytes.strip_prefix(MAGIC) else { bail!("not a checkpoint") };

    let mut n = Counts::default();

    let take = |rest: &mut &[u8], k: usize| -> Result<Vec<u8>> {

        if rest.len() < k {

            bail!("the checkpoint is cut short");

        }

        let (a, b) = rest.split_at(k);

        *rest = b;

        Ok(a.to_vec())

    };

    while !rest.is_empty() {

        let kind = take(&mut rest, 1)?[0];

        let len = u32::from_le_bytes(take(&mut rest, 4)?.try_into().unwrap()) as usize;

        let path = String::from_utf8(take(&mut rest, len)?).context("a path that is not UTF-8")?;

        let mode = u32::from_le_bytes(take(&mut rest, 4)?.try_into().unwrap());

        let secs = u64::from_le_bytes(take(&mut rest, 8)?.try_into().unwrap());

        let nanos = u32::from_le_bytes(take(&mut rest, 4)?.try_into().unwrap());

        let size = u64::from_le_bytes(take(&mut rest, 8)?.try_into().unwrap()) as usize;

        let data = take(&mut rest, size)?;

        // Under one of the job's checkpoints, without . or .. on the way.

        let Some(root) = checkpoints.iter().map(|c| c.trim_end_matches('/')).find(|c| path == *c || path.starts_with(&format!("{c}/")))

        else {

            bail!("{path} is not under the job's checkpoints");

        };

        let below = path[root.len()..].trim_start_matches('/');

        if below.split('/').any(|p| p == "." || p == ".." || (p.is_empty() && !below.is_empty())) {

            bail!("{path}: . or .. in a checkpoint path");

        }

        let Some(base) = place(root) else { continue };

        let at = if below.is_empty() { base } else { base.join(below) };

        match kind {

            DIR => {

                std::fs::create_dir_all(&at)?;

                n.dirs += 1;

            }

            FILE => {

                if let Some(parent) = at.parent() {

                    std::fs::create_dir_all(parent)?;

                }

                // Never through a link already there.

                if std::fs::symlink_metadata(&at).is_ok_and(|m| !m.is_file()) {

                    bail!("{path}: something other than a file is in the way");

                }

                std::fs::write(&at, &data)?;

                let f = std::fs::File::options().write(true).open(&at)?;

                f.set_permissions(std::fs::Permissions::from_mode(mode & 0o777))?;

                f.set_modified(SystemTime::UNIX_EPOCH + Duration::new(secs, nanos))?;

                n.files += 1;

                n.bytes += data.len() as u64;

            }

            other => bail!("{path}: unknown entry kind {other}"),

        }

    }

    Ok(n)

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn a_checkpoint_comes_back_with_its_times_and_only_where_it_belongs() {

        let from = tempfile::tempdir().unwrap();

        let target = from.path().join("target");

        std::fs::create_dir_all(target.join("debug/deps")).unwrap();

        std::fs::write(target.join("debug/deps/libx.rlib"), b"built").unwrap();

        std::fs::write(target.join("debug/.fingerprint"), b"fp").unwrap();

        let old = SystemTime::UNIX_EPOCH + Duration::new(1_700_000_000, 123_000_000);

        std::fs::File::options().write(true).open(target.join("debug/deps/libx.rlib")).unwrap().set_modified(old).unwrap();

        let (bytes, n) = pack(&[("/src/target".into(), target.clone())]).unwrap();

        assert_eq!((n.files, n.dirs, n.bytes), (2, 3, 7));

        let to = tempfile::tempdir().unwrap();

        let place = |root: &str| (root == "/src/target").then(|| to.path().join("t"));

        let back = unpack(&bytes, &["/src/target".into()], place).unwrap();

        assert_eq!(back, n);

        let restored = to.path().join("t/debug/deps/libx.rlib");

        assert_eq!(std::fs::read(&restored).unwrap(), b"built");

        assert_eq!(std::fs::metadata(&restored).unwrap().modified().unwrap(), old, "cargo reads the times");

        // Not the job's checkpoint, or a path that climbs out: refused.

        assert!(unpack(&bytes, &["/src/other".into()], place).is_err());

        let mut evil = MAGIC.to_vec();

        let meta = std::fs::metadata(target.join("debug/.fingerprint")).unwrap();

        entry(&mut evil, FILE, "/src/target/../escape", &meta, b"x");

        assert!(unpack(&evil, &["/src/target".into()], place).is_err());

        assert!(!to.path().join("escape").exists());

        assert!(unpack(b"junk", &["/src/target".into()], place).is_err());

    }

    #[test]

    fn a_checkpoint_under_a_disk_mount_is_in_the_mount_s_directory() {

        let p = flow_core::Pipeline::parse(

            "schema-version 1\nsandbox \"w\" {\n    runtime \"wasm\"\n}\njob \"j\" {\n    sandbox \"w\"\n    mount {\n        source \"ephemeral://disk\"\n        target \"/src/target\"\n        size \"1GiB\"\n    }\n    checkpoint \"/src/target\" \"/src/state\"\n    exec \"x.wasm\"\n}\n",

        )

        .unwrap();

        let job = &p.jobs["j"];

        assert_eq!(job.checkpoints, ["/src/state", "/src/target"]);

        let (tree, scratch) = (Path::new("/w/tree"), Path::new("/w"));

        assert_eq!(host_dir("/src/target", job, tree, scratch, true), Some(crate::syscall::disk_dir(scratch, 0)));

        assert_eq!(host_dir("/src/state", job, tree, scratch, true), Some(PathBuf::from("/w/tree/state")));

    }

}

- **Checkpoints: a handed-back job resumes (#77):** `checkpoint "/src/…"` on a job names directories that hold its progress; a runner handing the job back packs and keeps them (`POST …/checkpoint`), and the next attempt's runner restores them before running the same exec, so it carries on rather than starting again. The receipt lists the checkpoint among its inputs. WASM and syscall tiers for now. Spec §11.5.2.

    /// The checkpoint an earlier attempt saved before its runner handed the

    /// job back (oxygenflow#77): the blob to fetch from

    /// `/api/v1/flow/jobs/{job_id}/checkpoint` and check, then restore

    /// before the job runs.

    #[serde(default, skip_serializing_if = "Option::is_none")]

    pub checkpoint: Option<String>,

    /// Directories under /src that hold the job's progress (`checkpoint

    /// "/src/target"`, oxygenflow#77), sorted: saved when its runner hands

    /// it back, restored before the next attempt runs the same exec.

    #[serde(default, skip_serializing_if = "Vec::is_empty")]

    pub checkpoints: Vec<String>,

        checkpoints: Vec::new(),

            && key != "checkpoint"

            }

            "checkpoint" => {

                ctx.no_children(c);

                for p in ctx.strings(c) {

                    if output_path_ok(&p) && p.trim_end_matches('/') != "/src" {

                        job.checkpoints.push(p);

                    } else {

                        ctx.err(c.span(), format!("checkpoint {p:?} must be a directory under /src, without . or .., quotes or spaces"));

                    }

                }

    job.checkpoints.sort();

    job.checkpoints.dedup();

        checkpoint: None,

pub mod checkpoint;

            checkpoints: Vec::new(),

    /// Keep a checkpoint of job `job` (oxygenflow#77) as its runner hands

    /// it back: the hub answers with its blob, which the next attempt's

    /// claim names.

    pub fn save_checkpoint(&self, job: i64, bytes: &[u8]) -> Result<String> {

        let url = format!("{}/api/v1/flow/jobs/{job}/checkpoint", self.base);

        let r = retrying(&url, &PATIENT_DELAYS, true, || {

            self.agent

                .post(&url)

                .timeout(Duration::from_secs(600))

                .set("Authorization", &format!("Bearer {}", self.token))

                .set("Content-Type", "application/octet-stream")

                .send_bytes(bytes)

                .map_err(Box::new)

        })?;

        let v: serde_json::Value = r.into_json()?;

        v["blob"].as_str().map(String::from).context("the hub named no blob")

    }

    /// The checkpoint an earlier attempt of job `job` saved (oxygenflow#77),

    /// which the caller checks against the blob its claim names.

    pub fn checkpoint(&self, job: i64) -> Result<Vec<u8>> {

        let url = format!("{}/api/v1/flow/jobs/{job}/checkpoint", self.base);

        let r = retrying(&url, &PATIENT_DELAYS, false, || {

            self.agent

                .get(&url)

                .timeout(Duration::from_secs(600))

                .set("Authorization", &format!("Bearer {}", self.token))

                .call()

                .map_err(Box::new)

        })?;

        let mut bytes = Vec::new();

        std::io::Read::read_to_end(&mut std::io::Read::take(r.into_reader(), crate::checkpoint::MAX_BYTES as u64 + 1), &mut bytes)?;

        Ok(bytes)

    }

}

/// Pack and keep a handed-back job's checkpoints (oxygenflow#77); what

/// happened goes to its log, and a failure only means the next attempt

/// starts again.

#[allow(clippy::too_many_arguments)]

fn save_checkpoint(hub: &Hub, id: i64, job: &flow_core::Job, scratch: &Path, dir_tier: bool, wasm: bool, mut say: impl FnMut(String)) {

    if !dir_tier {

        say("checkpoint: only the WASM and syscall tiers save one yet; the next attempt starts again".into());

        return;

    }

    let dirs: Vec<(String, PathBuf)> = job

        .checkpoints

        .iter()

        .filter_map(|c| crate::checkpoint::host_dir(c, job, &scratch.join("tree"), scratch, wasm).map(|d| (c.clone(), d)))

        .collect();

    match crate::checkpoint::pack(&dirs).and_then(|(bytes, n)| Ok((hub.save_checkpoint(id, &bytes)?, n))) {

        Ok((blob, n)) => {

            say(format!("checkpoint {}: {} files, {} bytes, kept for the next attempt", &blob[..12.min(blob.len())], n.files, n.bytes))

        }

        Err(e) => say(format!("checkpoint: not kept ({e:#}); the next attempt starts again")),

    }

        // An earlier attempt's checkpoint (oxygenflow#77): back in place,

        // over the cache, before the same exec runs again; the receipt

        // lists it among the job's inputs.

        if let Some(blob) = &claim.checkpoint {

            if dir_tier {

                let bytes = hub.checkpoint(claim.job_id).context("fetching the checkpoint")?;

                if Repo::blob_ref(&bytes).to_hex() != *blob {

                    bail!("the checkpoint fetched is not blob {}, as the claim names", &blob[..12.min(blob.len())]);

                }

                let place = |root: &str| crate::checkpoint::host_dir(root, &job, &scratch.join("tree"), &scratch, wasm_tier);

                let n = crate::checkpoint::unpack(&bytes, &job.checkpoints, place)?;

                note(

                    &mut sink,

                    format!("restored checkpoint {}: {} files, {} bytes, from the attempt handed back", &blob[..12], n.files, n.bytes),

                );

                input_blobs.push(Repo::blob_ref(&bytes).0);

            } else {

                note(&mut sink, "a checkpoint was saved, but only the WASM and syscall tiers restore one yet: starting again".into());

            }

        }

            // Handed back as the runner leaves: its checkpoints first, for

            // the next attempt (oxygenflow#77).

            if leaving().is_some() && !job.checkpoints.is_empty() {

                save_checkpoint(hub, claim.job_id, &job, &scratch, dir_tier, wasm_tier, |t| note(&mut sink, t));

            }

**Checkpoints (#77).** A job may name directories that hold its progress:

`checkpoint "/src/target"` (one or more, under /src, not /src itself).

When its runner hands it back, the runner packs them (files with their

modes and modification times, which cargo's fingerprints read; at most

512 MiB) and keeps them with `POST /api/v1/flow/jobs/{id}/checkpoint`

before the release; the claim of the next attempt names the blob

(`checkpoint`), and that runner fetches it (`GET` on the same path),

checks it, and restores it over the job's cache before running the same

exec, so a tool that skips finished work carries on. The receipt lists

the checkpoint among the job's inputs. As built: the WASM and syscall

tiers save and restore checkpoints, under a disk mount or in the tree;

the MicroVM tier, whose files are in an image, starts again for now.