Images by hash: flow rootfs push and pull, and runners fetch what they lack 4cf0c99a1486

//! Images by hash with a hub (oxygenflow#39, docs/jetstream.md §8):

//! `flow rootfs push` sends an image's missing chunks and its manifest;

//! `flow rootfs pull`, and a runner that lacks a job's image, fetch the

//! manifest and the chunks they lack and write the image back, checked

//! against its hash. Fetched chunks stay in the home's chunk cache

//! (`chunks/`), so a rebuilt image costs only the chunks that changed.

use crate::home::Home;

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

use jet::images::{self, Chunk, Manifest};

use std::collections::HashSet;

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

use std::sync::Mutex;

use std::time::{Duration, Instant};

/// Requests in flight at once.

const PARALLEL: usize = 8;

/// A hub's image endpoints, with a token.

pub struct Client {

    base: String,

    token: String,

    agent: ureq::Agent,

}

/// What a push or a fetch moved.

pub struct Moved {

    pub chunks: usize,

    pub bytes: u64,

    pub elapsed: Duration,

}

impl Client {

    pub fn new(hub: &str, token: &str, ca: Option<&Path>) -> Result<Client> {

        let tls = oxy_core::protocol::tls_client_config(ca).map_err(|e| anyhow::anyhow!("{e}"))?;

        let agent = ureq::AgentBuilder::new().tls_config(tls).timeout(Duration::from_secs(120)).build();

        Ok(Client { base: format!("{}/api/v1", hub.trim_end_matches('/')), token: token.into(), agent })

    }

    fn auth(&self) -> String {

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

    }

    fn fail(url: &str, e: ureq::Error) -> anyhow::Error {

        match e {

            ureq::Error::Status(code, r) => {

                let text = r.into_string().unwrap_or_default();

                let detail = serde_json::from_str::<serde_json::Value>(&text)

                    .ok()

                    .and_then(|v| v["detail"].as_str().map(String::from))

                    .unwrap_or(text);

                anyhow::anyhow!("{url}: HTTP {code}: {detail}")

            }

            e => anyhow::anyhow!("{url}: {e}"),

        }

    }

    /// The hub's manifest of image `hash`, `None` when it does not hold it.

    pub fn manifest(&self, hash: &str) -> Result<Option<Manifest>> {

        let url = format!("{}/flow/images/{hash}", self.base);

        match self.agent.get(&url).set("Authorization", &self.auth()).call() {

            Ok(r) => Ok(Some(r.into_json()?)),

            Err(ureq::Error::Status(404, _)) => Ok(None),

            Err(e) => Err(Self::fail(&url, e)),

        }

    }

    fn chunk(&self, hash: &str) -> Result<Vec<u8>> {

        let url = format!("{}/flow/images/chunks/{hash}", self.base);

        let r = self.agent.get(&url).set("Authorization", &self.auth()).call().map_err(|e| Self::fail(&url, e))?;

        let mut out = Vec::new();

        std::io::Read::read_to_end(&mut r.into_reader(), &mut out)?;

        Ok(out)

    }

    /// Send image `hash` from `home`: only the chunks the hub lacks, then

    /// the manifest, which the hub checks against the hash.

    pub fn push(&self, home: &Home, hash: &str) -> Result<Moved> {

        let start = Instant::now();

        let path = home.rootfs(hash);

        if !path.is_file() {

            bail!("{} is not in this Flow home ({}); import it first", &hash[..12.min(hash.len())], path.display());

        }

        let m = images::scan_file(&path, home.image_arch(hash), |_, _| Ok(()))?;

        if m.hash != hash {

            bail!("{} holds an image named {}, not {}: the file was changed", path.display(), &m.hash[..12], &hash[..12]);

        }

        let mut missing: HashSet<String> = HashSet::new();

        for ask in m.needed().chunks(50_000) {

            let url = format!("{}/flow/images/missing", self.base);

            let r: serde_json::Value = self

                .agent

                .post(&url)

                .set("Authorization", &self.auth())

                .send_json(serde_json::json!({ "chunks": ask }))

                .map_err(|e| Self::fail(&url, e))?

                .into_json()?;

            missing.extend(r["missing"].as_array().into_iter().flatten().filter_map(|h| h.as_str().map(String::from)));

        }

        // Read the image again, sending the missing chunks as they come.

        let (count, bytes) = (missing.len(), Mutex::new(0u64));

        let queue: Mutex<Vec<(String, Vec<u8>)>> = Mutex::new(Vec::new());

        let failed: Mutex<Option<anyhow::Error>> = Mutex::new(None);

        let sent = Mutex::new(HashSet::new());

        let send = |c: &str, b: &[u8]| -> Result<()> {

            let url = format!("{}/flow/images/chunks/{c}", self.base);

            self.agent.put(&url).set("Authorization", &self.auth()).send_bytes(b).map_err(|e| Self::fail(&url, e))?;

            *bytes.lock().unwrap() += b.len() as u64;

            Ok(())

        };

        images::scan_file(&path, None, |c: &Chunk, b: &[u8]| {

            if missing.contains(c.hash()) && sent.lock().unwrap().insert(c.hash().to_string()) {

                queue.lock().unwrap().push((c.hash().to_string(), b.to_vec()));

            }

            if queue.lock().unwrap().len() >= PARALLEL * 4 {

                let batch: Vec<_> = std::mem::take(&mut *queue.lock().unwrap());

                send_batch(&batch, &send, &failed);

            }

            match failed.lock().unwrap().take() {

                Some(e) => Err(images::Error::Bad(format!("{e:#}"))),

                None => Ok(()),

            }

        })?;

        let rest: Vec<_> = std::mem::take(&mut *queue.lock().unwrap());

        send_batch(&rest, &send, &failed);

        if let Some(e) = failed.lock().unwrap().take() {

            return Err(e);

        }

        let url = format!("{}/flow/images/{hash}", self.base);

        self.agent.put(&url).set("Authorization", &self.auth()).send_json(&m).map_err(|e| Self::fail(&url, e))?;

        let bytes = *bytes.lock().unwrap();

        Ok(Moved { chunks: count, bytes, elapsed: start.elapsed() })

    }

    /// Fetch image `hash` into `home`: the chunks its cache lacks, the

    /// image written and checked, its architecture recorded. `Ok(None)`

    /// when the hub does not hold it.

    pub fn fetch(&self, home: &Home, hash: &str) -> Result<Option<Moved>> {

        let start = Instant::now();

        let Some(m) = self.manifest(hash)? else { return Ok(None) };

        if m.hash != hash {

            bail!("the hub's manifest for {} names {}", &hash[..12], &m.hash[..12]);

        }

        let cache = chunk_cache(home);

        let wanted: Vec<String> = m.needed().into_iter().filter(|h| !cache_path(&cache, h).is_file()).map(String::from).collect();

        let (bytes, failed) = (Mutex::new(0u64), Mutex::new(None::<anyhow::Error>));

        let next = Mutex::new(wanted.iter());

        std::thread::scope(|s| {

            for _ in 0..PARALLEL {

                s.spawn(|| loop {

                    let Some(h) = next.lock().unwrap().next() else { break };

                    if failed.lock().unwrap().is_some() {

                        break;

                    }

                    let got = self.chunk(h).and_then(|b| {

                        if blake3::hash(&b).to_hex().as_str() != h {

                            bail!("chunk {} came back with other bytes", &h[..12]);

                        }

                        let p = cache_path(&cache, h);

                        std::fs::create_dir_all(p.parent().expect("a chunk directory"))?;

                        let tmp = p.with_extension("part");

                        std::fs::write(&tmp, &b)?;

                        std::fs::rename(&tmp, &p)?;

                        *bytes.lock().unwrap() += b.len() as u64;

                        Ok(())

                    });

                    if let Err(e) = got {

                        *failed.lock().unwrap() = Some(e);

                        break;

                    }

                });

            }

        });

        if let Some(e) = failed.into_inner().unwrap() {

            return Err(e);

        }

        let out = home.rootfs(hash);

        std::fs::create_dir_all(out.parent().expect("the cas directory"))?;

        images::assemble(&m, &out, |c| Ok(std::fs::read(cache_path(&cache, c.hash()))?))

            .with_context(|| format!("write image {}", &hash[..12]))?;

        if let Some(arch) = &m.arch {

            let manifest = serde_json::json!({ "arch": arch, "imported_ms": std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_millis() as u64).unwrap_or(0), "from": "hub" });

            std::fs::write(home.image_manifest(hash), serde_json::to_vec(&manifest)?)?;

        }

        Ok(Some(Moved { chunks: wanted.len(), bytes: bytes.into_inner().unwrap(), elapsed: start.elapsed() }))

    }

}

fn send_batch(batch: &[(String, Vec<u8>)], send: &(dyn Fn(&str, &[u8]) -> Result<()> + Sync), failed: &Mutex<Option<anyhow::Error>>) {

    std::thread::scope(|s| {

        for part in batch.chunks(batch.len().div_ceil(PARALLEL).max(1)) {

            s.spawn(move || {

                for (h, b) in part {

                    if let Err(e) = send(h, b) {

                        *failed.lock().unwrap() = Some(e);

                        return;

                    }

                }

            });

        }

    });

}

/// The home's chunk cache.

pub fn chunk_cache(home: &Home) -> PathBuf {

    home.root.join("chunks")

}

fn cache_path(cache: &Path, hash: &str) -> PathBuf {

    cache.join(&hash[..2]).join(hash)

}

/// A size for people: 991 MiB, 1.2 GiB.

pub fn human(bytes: u64) -> String {

    if bytes >= 1 << 30 {

        format!("{:.1} GiB", bytes as f64 / (1u64 << 30) as f64)

    } else {

        format!("{} MiB", bytes >> 20)

    }

}

 "jet",

    /// Whether the runner fetches a job's image from the hub when it lacks

    /// it (oxygenflow#39): the hub may then offer it jobs whose image only

    /// the hub holds.

    #[serde(default, skip_serializing_if = "std::ops::Not::not")]

    pub fetches_images: bool,

jet.workspace = true

mod images;

    /// Send an image to a hub, so its runners fetch it instead of having it

    /// copied to them: only the chunks the hub lacks travel. Uses

    /// $OXY_TOKEN (a token that may write repositories) and $OXY_CA

    Push {

        /// The image's hash, as pipelines pin it (cas://blob/HASH)

        hash: String,

        /// The hub (default: the runner's, from runner.toml or flow.toml)

        #[arg(long)]

        hub: Option<String>,

    },

    /// Fetch an image a hub holds into this home, as a runner does for a

    /// job whose image it lacks. Uses $OXY_TOKEN and $OXY_CA

    Pull {

        hash: String,

        #[arg(long)]

        hub: Option<String>,

    },

                RootfsCmd::Push { hash, hub } => {

                    let client = image_client(&home, hub)?;

                    eprintln!("Sending image {} …", &hash[..12.min(hash.len())]);

                    let m = client.push(&home, &hash)?;

                    println!(

                        "Sent image {}: {} chunks the hub lacked, {} in {:.1} s",

                        &hash[..12],

                        m.chunks,

                        images::human(m.bytes),

                        m.elapsed.as_secs_f64()

                    );

                    Ok(())

                }

                RootfsCmd::Pull { hash, hub } => {

                    let client = image_client(&home, hub)?;

                    match client.fetch(&home, &hash)? {

                        Some(m) => {

                            println!(

                                "Fetched image {}: {} chunks, {} in {:.1} s",

                                &hash[..12],

                                m.chunks,

                                images::human(m.bytes),

                                m.elapsed.as_secs_f64()

                            );

                            Ok(())

                        }

                        None => anyhow::bail!("the hub does not hold image {}; send it with flow rootfs push", &hash[..12.min(hash.len())]),

                    }

                }

}

/// The hub image commands talk to: `--hub`, else the runner's; with

/// $OXY_TOKEN, else the runner's own token.

fn image_client(home: &Home, hub: Option<String>) -> Result<images::Client> {

    let runner = home.config.runner.clone();

    let hub = hub.or_else(|| runner.as_ref().map(|r| r.hub.clone())).context("name the hub with --hub https://HUB")?;

    let token = std::env::var("OXY_TOKEN")

        .ok()

        .or_else(|| runner.as_ref().map(|r| r.token.clone()))

        .context("set OXY_TOKEN to a token of yours on the hub")?;

    let ca = std::env::var_os("OXY_CA").map(PathBuf::from).or_else(|| runner.as_ref().and_then(|r| r.ca.clone()));

    images::Client::new(&hub, &token, ca.as_deref())

/// Fetch `claim`'s image from the hub when this home lacks it

/// (oxygenflow#39), saying so in the job's log. A hub that does not hold it

/// leaves it missing, and the preflight declines the job as before.

fn fetch_image(home: &Home, hub: &Hub, cfg: &RunnerConfig, claim: &Claim) {

    let hash = &claim.sandbox.rootfs;

    if hash.is_empty() || home.rootfs(hash).is_file() || !jet::images::is_hash(hash) {

        return;

    }

    let note = |text: String| {

        eprintln!("job {}: {text}", claim.job_id);

        hub.log(claim.job_id, vec![Record { t: 0, stream: "flow".into(), text }]).ok();

    };

    // One fetch of an image at a time: two jobs wanting it wait for one.

    static FETCHING: std::sync::Mutex<()> = std::sync::Mutex::new(());

    let _one = FETCHING.lock().unwrap_or_else(|p| p.into_inner());

    if home.rootfs(hash).is_file() {

        return;

    }

    let client = match crate::images::Client::new(&cfg.hub, &cfg.token, cfg.ca.as_deref()) {

        Ok(c) => c,

        Err(e) => return note(format!("image {}: cannot reach the hub: {e:#}", &hash[..12])),

    };

    note(format!("image {} is not on this runner: fetching it from the hub", &hash[..12]));

    match client.fetch(home, hash) {

        Ok(Some(m)) => note(format!(

            "image {}: fetched {} chunks, {} in {:.1} s",

            &hash[..12],

            m.chunks,

            crate::images::human(m.bytes),

            m.elapsed.as_secs_f64()

        )),

        Ok(None) => note(format!("image {}: the hub does not hold it either", &hash[..12])),

        Err(e) => note(format!("image {}: fetch failed: {e:#}", &hash[..12])),

    }

}

                fetches_images: true,

                        fetch_image(home, hub, cfg, &claim);