OXYGENHUB

Coverage / crates/flow/src/checkpoint.rs 92.2%

141 of 153 lines run, 16 of 34 functions · run #216 on main · run missed

1
//! Checkpoints (oxygenflow#77, docs/handover.md §7): the directories a job
2
//! declares with `checkpoint "/src/…"`, saved when its runner hands it back
3
//! and restored before the next attempt runs the same exec, so a tool that
4
//! skips finished work (cargo, make) carries on where it stopped.
5
//!
6
//! The format is the runner's own, so CI builds need no new crate: a magic
7
//! line, then entries (a kind byte, the path under /src, mode, modification
8
//! time and, for a file, its bytes). Files keep their times, since cargo's
9
//! fingerprints read them. Links are left out. Unpacking accepts only paths
10
//! under the job's own checkpoints, without `.` or `..`.
11
12
use anyhow::{bail, Context, Result};
13
use flow_core::pipeline::{Job, MountSource};
14
use std::path::{Path, PathBuf};
15
use std::time::{Duration, SystemTime, UNIX_EPOCH};
16
17
/// At most this much is saved: the hub keeps it as it keeps an artifact.
18
pub const MAX_BYTES: usize = 512 << 20;
19
20
const MAGIC: &[u8] = b"OXYFLOW-CHECKPOINT 1\n";
21
const DIR: u8 = 0;
22
const FILE: u8 = 1;
23
24
/// What was saved or restored.
25
#[derive(Debug, Default, PartialEq, Eq)]
26
pub struct Counts {
27
    pub files: usize,
28
    pub dirs: usize,
29
    pub bytes: u64,
30
}
31
32
/// Where checkpoint `path` (under /src) is on this host, for a job of a
33
/// directory tier (WASM, syscall): in a disk mount's directory when a mount
34
/// holds it, else in the job's tree. `None` for one under a tmpfs, which
35
/// the syscall tier keeps in memory.
362×
pub fn host_dir(path: &str, job: &Job, tree: &Path, scratch: &Path, wasm: bool) -> Option<PathBuf> {
372×
    let under = |t: &str| path == t || path.starts_with(&format!("{}/", t.trim_end_matches('/')));
382×
    let mut disks = 0;
392×
    let mut best: Option<(usize, PathBuf)> = None;
402×
    for (i, m) in job.mounts.iter().enumerate() {
412×
        let dir = match m.source {
42
            MountSource::Disk => {
432×
                disks += 1;
442×
                Some(crate::syscall::disk_dir(scratch, disks - 1))
45
            }
460×
            MountSource::Tmpfs if wasm => Some(scratch.join(format!("tmpfs-{i}"))),
470×
            MountSource::Tmpfs => None,
48
        };
492×
        if under(&m.target) && best.as_ref().is_none_or(|(len, _)| m.target.len() > *len) {
501×
            let rest = path[m.target.trim_end_matches('/').len()..].trim_start_matches('/');
511×
            best = Some((m.target.len(), dir.map(|d| d.join(rest)).unwrap_or_default()));
521×
        }
53
    }
541×
    match best {
551×
        Some((_, d)) if d.as_os_str().is_empty() => None,
561×
        Some((_, d)) => Some(d),
571×
        None => Some(tree.join(path.trim_start_matches("/src").trim_start_matches('/'))),
58
    }
592×
}
60
61
/// Pack `dirs` (each a checkpoint path and where it is on this host).
621×
pub fn pack(dirs: &[(String, PathBuf)]) -> Result<(Vec<u8>, Counts)> {
631×
    let mut out = MAGIC.to_vec();
641×
    let mut n = Counts::default();
651×
    for (path, host) in dirs {
661×
        if !host.is_dir() {
670×
            continue;
681×
        }
691×
        let mut stack = vec![(host.clone(), path.trim_end_matches('/').to_string())];
706×
        while let Some((at, name)) = stack.pop() {
715×
            let meta = std::fs::symlink_metadata(&at)?;
725×
            if meta.is_dir() {
733×
                entry(&mut out, DIR, &name, &meta, &[]);
743×
                n.dirs += 1;
753×
                let mut children: Vec<_> = std::fs::read_dir(&at)?.collect::<std::io::Result<_>>()?;
763×
                children.sort_by_key(|e| e.file_name());
774×
                for c in children.into_iter().rev() {
784×
                    let Some(child) = c.file_name().to_str().map(String::from) else { continue };
794×
                    stack.push((c.path(), format!("{name}/{child}")));
80
                }
812×
            } else if meta.is_file() {
822×
                let bytes = std::fs::read(&at).with_context(|| format!("read {}", at.display()))?;
832×
                n.files += 1;
842×
                n.bytes += bytes.len() as u64;
852×
                entry(&mut out, FILE, &name, &meta, &bytes);
860×
            }
875×
            if out.len() > MAX_BYTES {
880×
                bail!("the checkpoint is over {} MiB", MAX_BYTES >> 20);
895×
            }
90
        }
91
    }
921×
    Ok((out, n))
931×
}
94
956×
fn entry(out: &mut Vec<u8>, kind: u8, path: &str, meta: &std::fs::Metadata, bytes: &[u8]) {
966×
    let since = meta.modified().ok().and_then(|t| t.duration_since(UNIX_EPOCH).ok()).unwrap_or_default();
976×
    out.push(kind);
986×
    out.extend((path.len() as u32).to_le_bytes());
996×
    out.extend(path.as_bytes());
1006×
    out.extend(crate::portable::mode(meta).to_le_bytes());
1016×
    out.extend(since.as_secs().to_le_bytes());
1026×
    out.extend(since.subsec_nanos().to_le_bytes());
1036×
    out.extend((bytes.len() as u64).to_le_bytes());
1046×
    out.extend(bytes);
1056×
}
106
107
/// Restore a checkpoint: each entry at `place(checkpoint)` joined with the
108
/// rest of its path, for the job's own `checkpoints` only.
1094×
pub fn unpack(bytes: &[u8], checkpoints: &[String], place: impl Fn(&str) -> Option<PathBuf>) -> Result<Counts> {
1104×
    let Some(mut rest) = bytes.strip_prefix(MAGIC) else { bail!("not a checkpoint") };
1113×
    let mut n = Counts::default();
11256×
    let take = |rest: &mut &[u8], k: usize| -> Result<Vec<u8>> {
11356×
        if rest.len() < k {
1140×
            bail!("the checkpoint is cut short");
11556×
        }
11656×
        let (a, b) = rest.split_at(k);
11756×
        *rest = b;
11856×
        Ok(a.to_vec())
11956×
    };
1208×
    while !rest.is_empty() {
1217×
        let kind = take(&mut rest, 1)?[0];
1227×
        let len = u32::from_le_bytes(take(&mut rest, 4)?.try_into().unwrap()) as usize;
1237×
        let path = String::from_utf8(take(&mut rest, len)?).context("a path that is not UTF-8")?;
1247×
        let mode = u32::from_le_bytes(take(&mut rest, 4)?.try_into().unwrap());
1257×
        let secs = u64::from_le_bytes(take(&mut rest, 8)?.try_into().unwrap());
1267×
        let nanos = u32::from_le_bytes(take(&mut rest, 4)?.try_into().unwrap());
1277×
        let size = u64::from_le_bytes(take(&mut rest, 8)?.try_into().unwrap()) as usize;
1287×
        let data = take(&mut rest, size)?;
129
        // Under one of the job's checkpoints, without . or .. on the way.
1307×
        let Some(root) = checkpoints.iter().map(|c| c.trim_end_matches('/')).find(|c| path == *c || path.starts_with(&format!("{c}/")))
131
        else {
1321×
            bail!("{path} is not under the job's checkpoints");
133
        };
1346×
        let below = path[root.len()..].trim_start_matches('/');
13510×
        if below.split('/').any(|p| p == "." || p == ".." || (p.is_empty() && !below.is_empty())) {
1361×
            bail!("{path}: . or .. in a checkpoint path");
1375×
        }
1385×
        let Some(base) = place(root) else { continue };
1395×
        let at = if below.is_empty() { base } else { base.join(below) };
1405×
        match kind {
141
            DIR => {
1423×
                std::fs::create_dir_all(&at)?;
1433×
                n.dirs += 1;
144
            }
145
            FILE => {
1462×
                if let Some(parent) = at.parent() {
1472×
                    std::fs::create_dir_all(parent)?;
1480×
                }
149
                // Never through a link already there.
1502×
                if std::fs::symlink_metadata(&at).is_ok_and(|m| !m.is_file()) {
1510×
                    bail!("{path}: something other than a file is in the way");
1522×
                }
1532×
                std::fs::write(&at, &data)?;
1542×
                let f = std::fs::File::options().write(true).open(&at)?;
1552×
                crate::portable::set_mode(&at, mode & 0o777)?;
1562×
                f.set_modified(SystemTime::UNIX_EPOCH + Duration::new(secs, nanos))?;
1572×
                n.files += 1;
1582×
                n.bytes += data.len() as u64;
159
            }
1600×
            other => bail!("{path}: unknown entry kind {other}"),
161
        }
162
    }
1631×
    Ok(n)
1644×
}
165
166
#[cfg(test)]
167
mod tests {
168
    use super::*;
169
170
    #[test]
1711×
    fn a_checkpoint_comes_back_with_its_times_and_only_where_it_belongs() {
1721×
        let from = tempfile::tempdir().unwrap();
1731×
        let target = from.path().join("target");
1741×
        std::fs::create_dir_all(target.join("debug/deps")).unwrap();
1751×
        std::fs::write(target.join("debug/deps/libx.rlib"), b"built").unwrap();
1761×
        std::fs::write(target.join("debug/.fingerprint"), b"fp").unwrap();
1771×
        let old = SystemTime::UNIX_EPOCH + Duration::new(1_700_000_000, 123_000_000);
1781×
        std::fs::File::options().write(true).open(target.join("debug/deps/libx.rlib")).unwrap().set_modified(old).unwrap();
1791×
        let (bytes, n) = pack(&[("/src/target".into(), target.clone())]).unwrap();
1801×
        assert_eq!((n.files, n.dirs, n.bytes), (2, 3, 7));
181
1821×
        let to = tempfile::tempdir().unwrap();
1835×
        let place = |root: &str| (root == "/src/target").then(|| to.path().join("t"));
1841×
        let back = unpack(&bytes, &["/src/target".into()], place).unwrap();
1851×
        assert_eq!(back, n);
1861×
        let restored = to.path().join("t/debug/deps/libx.rlib");
1871×
        assert_eq!(std::fs::read(&restored).unwrap(), b"built");
1881×
        assert_eq!(std::fs::metadata(&restored).unwrap().modified().unwrap(), old, "cargo reads the times");
189
190
        // Not the job's checkpoint, or a path that climbs out: refused.
1911×
        assert!(unpack(&bytes, &["/src/other".into()], place).is_err());
1921×
        let mut evil = MAGIC.to_vec();
1931×
        let meta = std::fs::metadata(target.join("debug/.fingerprint")).unwrap();
1941×
        entry(&mut evil, FILE, "/src/target/../escape", &meta, b"x");
1951×
        assert!(unpack(&evil, &["/src/target".into()], place).is_err());
1961×
        assert!(!to.path().join("escape").exists());
1971×
        assert!(unpack(b"junk", &["/src/target".into()], place).is_err());
1981×
    }
199
200
    #[test]
2011×
    fn a_checkpoint_under_a_disk_mount_is_in_the_mount_s_directory() {
2021×
        let p = flow_core::Pipeline::parse(
2031×
            "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",
204
        )
2051×
        .unwrap();
2061×
        let job = &p.jobs["j"];
2071×
        assert_eq!(job.checkpoints, ["/src/state", "/src/target"]);
2081×
        let (tree, scratch) = (Path::new("/w/tree"), Path::new("/w"));
2091×
        assert_eq!(host_dir("/src/target", job, tree, scratch, true), Some(crate::syscall::disk_dir(scratch, 0)));
2101×
        assert_eq!(host_dir("/src/state", job, tree, scratch, true), Some(PathBuf::from("/w/tree/state")));
2111×
    }
212
}