OXYGENHUB
◆ v0.4.2 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
92a797b2 Andy McCall ↶ 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
}