OXYGENHUB
# c95446a17fdd oxygenhub / crates / hub / src / images.rs
4.4 KiB · Rust · last changed Images by hash on the hub: chunks, verified manifests, and placement for runners that fetch 25 days ago source blame history raw
1 
//! Images by hash (oxygenflow#39, OxygenFlow docs/jetstream.md §8): rootfs
2 
//! images kept site-wide for runners to fetch, so nobody copies an image
3 
//! between machines. An image arrives as chunks (Oxygen's FastCDC, sent
4 
//! only when this hub lacks them) and then its manifest, which is accepted
5 
//! only when every chunk is here and the image they make has the hash it
6 
//! is named by. Chunks are kept zstd-compressed under
7 
//! `<data>/images/chunks`, shared by every image that has them.
8
9 
use crate::db::{now_ms, Pool};
10 
use jet::images::{is_hash, Manifest};
11 
use std::io::Write;
12 
use std::path::{Path, PathBuf};
13
14 
/// Where this hub keeps images.
15 
#[derive(Clone, Debug)]
16 
pub struct Store {
17 
    root: PathBuf,
18 
}
19
20 
impl Store {
21 
    pub fn new(data: &Path) -> Store {
22 
        Store { root: data.join("images") }
23 
    }
24
25 
    fn chunk_path(&self, hash: &str) -> PathBuf {
26 
        self.root.join("chunks").join(&hash[..2]).join(hash)
27 
    }
28
29 
    fn manifest_path(&self, hash: &str) -> PathBuf {
30 
        self.root.join("manifests").join(format!("{hash}.json"))
31 
    }
32
33 
    pub fn has_chunk(&self, hash: &str) -> bool {
34 
        is_hash(hash) && self.chunk_path(hash).exists()
35 
    }
36
37 
    /// Keep chunk `hash`, whose bytes the caller has checked.
38 
    pub fn put_chunk(&self, hash: &str, bytes: &[u8]) -> std::io::Result<()> {
39 
        let path = self.chunk_path(hash);
40 
        if path.exists() {
41 
            return Ok(());
42 
        }
43 
        std::fs::create_dir_all(path.parent().expect("a chunk's directory"))?;
44 
        let tmp = path.with_extension(format!("tmp{}", std::process::id()));
45 
        let mut f = std::fs::File::create(&tmp)?;
46 
        f.write_all(&zstd::encode_all(bytes, 3)?)?;
47 
        f.sync_all()?;
48 
        std::fs::rename(&tmp, &path)
49 
    }
50
51 
    pub fn chunk(&self, hash: &str) -> std::io::Result<Vec<u8>> {
52 
        if !is_hash(hash) {
53 
            return Err(std::io::Error::new(std::io::ErrorKind::NotFound, "not a chunk name"));
54 
        }
55 
        zstd::decode_all(std::fs::File::open(self.chunk_path(hash))?)
56 
    }
57
58 
    pub fn manifest(&self, hash: &str) -> Option<Manifest> {
59 
        if !is_hash(hash) {
60 
            return None;
61 
        }
62 
        serde_json::from_slice(&std::fs::read(self.manifest_path(hash)).ok()?).ok()
63 
    }
64
65 
    /// Accept manifest `m`: well formed, every chunk here, and the image
66 
    /// they make hashing to its name. Reads the whole image once.
67 
    pub fn verify_and_keep(&self, m: &Manifest) -> Result<(), String> {
68 
        m.check().map_err(|e| e.to_string())?;
69 
        let mut whole = blake3::Hasher::new();
70 
        for c in &m.chunks {
71 
            if c.is_zeros() {
72 
                whole.update(&vec![0u8; c.len()]);
73 
                continue;
74 
            }
75 
            let bytes = self.chunk(c.hash()).map_err(|_| format!("chunk {} was not sent; send it first", &c.hash()[..12]))?;
76 
            if bytes.len() != c.len() {
77 
                return Err(format!("chunk {} is {} bytes, not {}", &c.hash()[..12], bytes.len(), c.len()));
78 
            }
79 
            whole.update(&bytes);
80 
        }
81 
        let got = whole.finalize().to_hex().to_string();
82 
        if got != m.hash {
83 
            return Err(format!("the chunks make an image named {}, not {}", &got[..12], &m.hash[..12]));
84 
        }
85 
        let path = self.manifest_path(&m.hash);
86 
        std::fs::create_dir_all(path.parent().expect("the manifests directory")).map_err(|e| e.to_string())?;
87 
        let tmp = path.with_extension("json.tmp");
88 
        std::fs::write(&tmp, serde_json::to_vec(m).map_err(|e| e.to_string())?).map_err(|e| e.to_string())?;
89 
        std::fs::rename(&tmp, &path).map_err(|e| e.to_string())
90 
    }
91 
}
92
93 
/// Record an image this hub now keeps.
94 
pub async fn record(pool: &Pool, m: &Manifest, uploader: &str) -> sqlx::Result<()> {
95 
    sqlx::query(
96 
        "INSERT INTO jet_images (hash, size, arch, chunks, uploader, created_ms) VALUES (?, ?, ?, ?, ?, ?) \
97 
         ON CONFLICT (hash) DO NOTHING",
98 
    )
99 
    .bind(&m.hash)
100 
    .bind(m.size as i64)
101 
    .bind(m.arch.clone().unwrap_or_default())
102 
    .bind(m.chunks.len() as i64)
103 
    .bind(uploader)
104 
    .bind(now_ms())
105 
    .execute(pool)
106 
    .await?;
107 
    if let Some(a) = m.arch.as_ref().filter(|a| !a.is_empty()) {
108 
        crate::db::flow::record_image_archs(pool, &[(m.hash.clone(), a.clone())].into_iter().collect()).await?;
109 
    }
110 
    Ok(())
111 
}
112
113 
/// Every image this hub keeps, by hash.
114 
pub async fn held(pool: &Pool) -> sqlx::Result<std::collections::HashSet<String>> {
115 
    Ok(sqlx::query_scalar::<_, String>("SELECT hash FROM jet_images").fetch_all(pool).await?.into_iter().collect())
116 
}