OXYGENHUB

Coverage / crates/hub/src/images.rs 92.8%

194 of 209 lines run, 32 of 76 functions · run #496 on main · run missed

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 {
2125×
    pub fn new(data: &Path) -> Store {
2225×
        Store { root: data.join("images") }
2325×
    }
24
2560036×
    fn chunk_path(&self, hash: &str) -> PathBuf {
2660036×
        self.root.join("chunks").join(&hash[..2]).join(hash)
2760036×
    }
28
2917×
    fn manifest_path(&self, hash: &str) -> PathBuf {
3017×
        self.root.join("manifests").join(format!("{hash}.json"))
3117×
    }
32
3360006×
    pub fn has_chunk(&self, hash: &str) -> bool {
3460006×
        is_hash(hash) && self.chunk_path(hash).exists()
3560006×
    }
36
37
    /// Keep chunk `hash`, whose bytes the caller has checked.
388×
    pub fn put_chunk(&self, hash: &str, bytes: &[u8]) -> std::io::Result<()> {
398×
        let path = self.chunk_path(hash);
408×
        if path.exists() {
410×
            return Ok(());
428×
        }
438×
        std::fs::create_dir_all(path.parent().expect("a chunk's directory"))?;
448×
        let tmp = path.with_extension(format!("tmp{}", std::process::id()));
458×
        let mut f = std::fs::File::create(&tmp)?;
468×
        f.write_all(&zstd::encode_all(bytes, 3)?)?;
478×
        f.sync_all()?;
488×
        std::fs::rename(&tmp, &path)
498×
    }
50
5122×
    pub fn chunk(&self, hash: &str) -> std::io::Result<Vec<u8>> {
5222×
        if !is_hash(hash) {
530×
            return Err(std::io::Error::new(std::io::ErrorKind::NotFound, "not a chunk name"));
5422×
        }
5522×
        zstd::decode_all(std::fs::File::open(self.chunk_path(hash))?)
5622×
    }
57
5812×
    pub fn manifest(&self, hash: &str) -> Option<Manifest> {
5912×
        if !is_hash(hash) {
600×
            return None;
6112×
        }
6212×
        serde_json::from_slice(&std::fs::read(self.manifest_path(hash)).ok()?).ok()
6312×
    }
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.
675×
    pub fn verify_and_keep(&self, m: &Manifest) -> Result<(), String> {
685×
        m.check().map_err(|e| e.to_string())?;
695×
        let mut whole = blake3::Hasher::new();
7055×
        for c in &m.chunks {
7155×
            if c.is_zeros() {
7243×
                whole.update(&vec![0u8; c.len()]);
7343×
                continue;
7412×
            }
7512×
            let bytes = self.chunk(c.hash()).map_err(|_| format!("chunk {} was not sent; send it first", &c.hash()[..12]))?;
7611×
            if bytes.len() != c.len() {
770×
                return Err(format!("chunk {} is {} bytes, not {}", &c.hash()[..12], bytes.len(), c.len()));
7811×
            }
7911×
            whole.update(&bytes);
80
        }
814×
        let got = whole.finalize().to_hex().to_string();
824×
        if got != m.hash {
831×
            return Err(format!("the chunks make an image named {}, not {}", &got[..12], &m.hash[..12]));
843×
        }
853×
        self.keep_manifest(m)
865×
    }
87
88
    /// Write manifest `m`, whose image is here.
895×
    fn keep_manifest(&self, m: &Manifest) -> Result<(), String> {
905×
        let path = self.manifest_path(&m.hash);
915×
        std::fs::create_dir_all(path.parent().expect("the manifests directory")).map_err(|e| e.to_string())?;
925×
        let tmp = path.with_extension("json.tmp");
935×
        std::fs::write(&tmp, serde_json::to_vec(m).map_err(|e| e.to_string())?).map_err(|e| e.to_string())?;
945×
        std::fs::rename(&tmp, &path).map_err(|e| e.to_string())
955×
    }
96
97
    /// What image `m` is for, read from its own files (OxygenFlow spec
98
    /// §11.2, oxygenhub#173): its shell's ELF header, found through its ext4
99
    /// out of the kept chunks, a few of them, not the whole image. `None`
100
    /// when it is not ext4 or has no shell to read.
1015×
    pub fn read_arch(&self, m: &Manifest) -> Option<&'static str> {
1025×
        let mut starts = Vec::with_capacity(m.chunks.len());
1035×
        let mut end = 0u64;
10478×
        for c in &m.chunks {
10578×
            starts.push(end);
10678×
            end += c.len() as u64;
10778×
        }
1085×
        let mut cache: std::collections::HashMap<usize, Vec<u8>> = std::collections::HashMap::new();
1095×
        let short = || std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "past the image's end");
11031×
        jet::rootfs::arch(&mut |at: u64, buf: &mut [u8]| {
11131×
            let mut done = 0usize;
11262×
            while done < buf.len() {
11331×
                let pos = at + done as u64;
114177×
                let i = starts.partition_point(|&s| s <= pos).checked_sub(1).ok_or_else(short)?;
11531×
                let c = &m.chunks[i];
11631×
                let within = (pos - starts[i]) as usize;
11731×
                if within >= c.len() {
1180×
                    return Err(short());
11931×
                }
12031×
                let n = (c.len() - within).min(buf.len() - done);
12131×
                if c.is_zeros() {
1220×
                    buf[done..done + n].fill(0);
1230×
                } else {
12431×
                    if !cache.contains_key(&i) {
1257×
                        if cache.len() >= 32 {
1260×
                            cache.clear();
1277×
                        }
1287×
                        cache.insert(i, self.chunk(c.hash())?);
12924×
                    }
13031×
                    let bytes = &cache[&i];
13131×
                    if bytes.len() != c.len() {
1320×
                        return Err(short());
13331×
                    }
13431×
                    buf[done..done + n].copy_from_slice(&bytes[within..within + n]);
135
                }
13631×
                done += n;
137
            }
13831×
            Ok(())
13931×
        })
1405×
    }
141
142
    /// Keep image `m`, checked as [`Store::verify_and_keep`], with the
143
    /// architecture its own files name: the manifest's word is taken only
144
    /// when the image cannot be read. The manifest as kept.
1455×
    pub fn keep_reading_arch(&self, m: &Manifest) -> Result<Manifest, String> {
1465×
        self.verify_and_keep(m)?;
1473×
        match self.read_arch(m) {
1481×
            Some(a) if m.arch.as_deref() != Some(a) => {
1491×
                let mut kept = m.clone();
1501×
                kept.arch = Some(a.to_string());
1511×
                self.keep_manifest(&kept)?;
1521×
                Ok(kept)
153
            }
1542×
            _ => Ok(m.clone()),
155
        }
1565×
    }
157
158
    /// Record `arch` in kept image `hash`'s manifest.
1591×
    fn set_arch(&self, hash: &str, arch: &str) -> Result<(), String> {
1601×
        let mut m = self.manifest(hash).ok_or("no such image")?;
1611×
        m.arch = Some(arch.to_string());
1621×
        self.keep_manifest(&m)
1631×
    }
164
}
165
166
/// Record an image this hub now keeps.
1674×
pub async fn record(pool: &Pool, m: &Manifest, uploader: &str) -> sqlx::Result<()> {
1684×
    sqlx::query(
1694×
        "INSERT INTO jet_images (hash, size, arch, chunks, uploader, created_ms) VALUES ($1, $2, $3, $4, $5, $6) \
1704×
         ON CONFLICT (hash) DO UPDATE SET arch = excluded.arch WHERE excluded.arch != ''",
1714×
    )
1724×
    .bind(&m.hash)
1734×
    .bind(m.size as i64)
1744×
    .bind(m.arch.clone().unwrap_or_default())
1754×
    .bind(m.chunks.len() as i64)
1764×
    .bind(uploader)
1774×
    .bind(now_ms())
1784×
    .execute(pool)
1794×
    .await?;
1804×
    if let Some(a) = m.arch.as_ref().filter(|a| !a.is_empty()) {
1813×
        crate::db::flow::record_image_archs(pool, &[(m.hash.clone(), a.clone())].into_iter().collect()).await?;
1821×
    }
1834×
    Ok(())
1844×
}
185
186
/// Learn what every image is for (oxygenhub#173), once at start: read
187
/// each kept image that has no architecture from its own files, then give
188
/// any image still unknown the architecture its jobs passed on. A runner's
189
/// own report, read from the image it holds, wins over what was learned.
190
/// How many images were read.
191192×
pub async fn learn_archs(app: &crate::Shared) -> sqlx::Result<usize> {
192192×
    let mut read = 0;
193192×
    for k in kept(&app.db).await?.into_iter().filter(|k| k.arch.is_empty()) {
1942×
        let store = app.images();
1952×
        let hash = k.hash.clone();
1962×
        let arch = tokio::task::spawn_blocking(move || {
1972×
            let a = store.read_arch(&store.manifest(&hash)?)?;
1981×
            store.set_arch(&hash, a).ok()?;
1991×
            Some(a)
2002×
        })
2012×
        .await
2022×
        .ok()
2032×
        .flatten();
2042×
        if let Some(a) = arch {
2051×
            sqlx::query("UPDATE jet_images SET arch = $1 WHERE hash = $2").bind(a).bind(&k.hash).execute(&app.db).await?;
2061×
            crate::db::flow::record_image_archs(&app.db, &[(k.hash.clone(), a.to_string())].into_iter().collect()).await?;
2071×
            read += 1;
2081×
        }
209
    }
210192×
    crate::db::flow::learn_archs_from_runs(&app.db).await?;
211192×
    Ok(read)
212192×
}
213
214
/// Every image this hub keeps, by hash.
215377×
pub async fn held(pool: &Pool) -> sqlx::Result<std::collections::HashSet<String>> {
216377×
    Ok(sqlx::query_scalar::<_, String>("SELECT hash FROM jet_images").fetch_all(pool).await?.into_iter().collect())
217377×
}
218
219
/// An image this hub keeps (`jet_images`).
220
#[derive(sqlx::FromRow)]
221
pub struct Kept {
222
    pub hash: String,
223
    pub size: i64,
224
    pub arch: String,
225
    pub chunks: i64,
226
    pub uploader: String,
227
    pub created_ms: i64,
228
}
229
230195×
pub async fn kept(pool: &Pool) -> sqlx::Result<Vec<Kept>> {
231195×
    sqlx::query_as("SELECT hash, size, arch, chunks, uploader, created_ms FROM jet_images").fetch_all(pool).await
232195×
}
233
234
/// One repository's use of one image: its jobs since `since_ms`, and when
235
/// it last ran one.
236
#[derive(sqlx::FromRow)]
237
pub struct Use {
238
    pub image: String,
239
    pub owner: String,
240
    pub name: String,
241
    pub jobs: i64,
242
    pub last_ms: i64,
243
}
244
245
/// What each repository's jobs ran on since `since_ms`, newest first.
2462×
pub async fn uses(pool: &Pool, since_ms: i64) -> sqlx::Result<Vec<Use>> {
2472×
    let rootfs = crate::db::json_text("j.sandbox", "rootfs");
2482×
    sqlx::query_as(&format!(
2492×
        "SELECT lower({rootfs}) AS image, p.owner, p.name, COUNT(*) AS jobs,
2502×
                MAX(COALESCE(j.started_ms, r.created_ms)) AS last_ms
2512×
         FROM flow_jobs j JOIN flow_runs r ON r.id = j.run_id JOIN repositories p ON p.id = r.repo_id
2522×
         WHERE r.created_ms >= $1 AND COALESCE({rootfs}, '') != ''
2532×
         GROUP BY lower({rootfs}), p.id, p.owner, p.name ORDER BY last_ms DESC"
2542×
    ))
2552×
    .bind(since_ms)
2562×
    .fetch_all(pool)
2572×
    .await
2582×
}
259
260
/// What one repository's jobs ran on since `since_ms`, newest first
261
/// (its Images page, oxygenhub#254).
2621×
pub async fn uses_of(pool: &Pool, repo_id: i64, since_ms: i64) -> sqlx::Result<Vec<Use>> {
2631×
    let rootfs = crate::db::json_text("j.sandbox", "rootfs");
2641×
    sqlx::query_as(&format!(
2651×
        "SELECT lower({rootfs}) AS image, p.owner, p.name, COUNT(*) AS jobs,
2661×
                MAX(COALESCE(j.started_ms, r.created_ms)) AS last_ms
2671×
         FROM flow_jobs j JOIN flow_runs r ON r.id = j.run_id JOIN repositories p ON p.id = r.repo_id
2681×
         WHERE r.created_ms >= $1 AND p.id = $2 AND COALESCE({rootfs}, '') != ''
2691×
         GROUP BY lower({rootfs}), p.id, p.owner, p.name ORDER BY last_ms DESC"
2701×
    ))
2711×
    .bind(since_ms)
2721×
    .bind(repo_id)
2731×
    .fetch_all(pool)
2741×
    .await
2751×
}
276
277
/// The architecture runners reported for each image (`flow_images`).
2782×
pub async fn reported_archs(pool: &Pool) -> sqlx::Result<std::collections::HashMap<String, String>> {
2792×
    Ok(sqlx::query_as::<_, (String, String)>("SELECT hash, arch FROM flow_images").fetch_all(pool).await?.into_iter().collect())
2802×
}