hypha: implement image cache

hypha/Cargo.toml

@@ -12,7 +12,11 @@ log = "0.4.20"
 env_logger = "0.10.1"
 flate2 = "1.0.28"
 tar = "0.4.40"
+directories = "5.0.1"
 walkdir = "2"
+serde = "1.0.195"
+serde_json = "1.0.111"
+sha256 = "1.5.0"
 
 [dependencies.clap]
 version = "4.4.18"

hypha/bin/controller.rs

@@ -1,6 +1,6 @@
 use clap::Parser;
 use hypha::ctl::Controller;
-use hypha::error::Result;
+use hypha::error::{HyphaError, Result};
 
 #[derive(Parser, Debug)]
 #[command(version, about)]
@@ -19,16 +19,38 @@ struct ControllerArgs {
 
     #[arg(short, long, default_value_t = 512)]
     mem: u64,
+
+    #[arg(short = 'C', long, default_value = "auto")]
+    cache: String,
 }
 
 fn main() -> Result<()> {
     env_logger::init();
 
     let args = ControllerArgs::parse();
-    let mut controller =
-        Controller::new(args.kernel, args.initrd, args.image, args.cpus, args.mem)?;
-    controller.compile()?;
+    let cache_path = if args.cache == "auto" {
+        default_cache_path()
+            .ok_or_else(|| HyphaError::new("unable to determine default cache path"))
+    } else {
+        Ok(args.cache)
+    }?;
+
+    let mut controller = Controller::new(
+        cache_path,
+        args.kernel,
+        args.initrd,
+        args.image,
+        args.cpus,
+        args.mem,
+    )?;
     let domid = controller.launch()?;
     println!("launched domain: {}", domid);
     Ok(())
 }
+
+fn default_cache_path() -> Option<String> {
+    let user_dirs = directories::UserDirs::new()?;
+    let mut path = user_dirs.home_dir().to_path_buf();
+    path.push(".hypha/cache");
+    Some(path.to_str()?.to_string())
+}

hypha/src/ctl/mod.rs

@@ -1,44 +1,56 @@
 use crate::error::Result;
-use crate::image::ImageCompiler;
+use crate::image::cache::ImageCache;
+use crate::image::{ImageCompiler, ImageInfo};
 use ocipkg::ImageName;
+use std::fs;
+use std::path::PathBuf;
 use xenclient::{DomainConfig, XenClient};
 
 pub struct Controller {
+    image_cache: ImageCache,
+    image: String,
     client: XenClient,
     kernel_path: String,
     initrd_path: String,
     vcpus: u32,
     mem: u64,
-    image: String,
 }
 
 impl Controller {
     pub fn new(
+        cache_path: String,
         kernel_path: String,
         initrd_path: String,
         image: String,
         vcpus: u32,
         mem: u64,
     ) -> Result<Controller> {
+        fs::create_dir_all(&cache_path)?;
+
         let client = XenClient::open()?;
+        let mut image_cache_path = PathBuf::from(cache_path);
+        image_cache_path.push("image");
+        fs::create_dir_all(&image_cache_path)?;
+        let image_cache = ImageCache::new(&image_cache_path)?;
         Ok(Controller {
+            image_cache,
+            image,
             client,
             kernel_path,
             initrd_path,
-            image,
             vcpus,
             mem,
         })
     }
 
-    pub fn compile(&mut self) -> Result<()> {
+    fn compile(&mut self) -> Result<ImageInfo> {
         let image = ImageName::parse(&self.image)?;
-        let compiler = ImageCompiler::new()?;
-        let _squashfs = compiler.compile(&image)?;
-        Ok(())
+        let compiler = ImageCompiler::new(&self.image_cache)?;
+        compiler.compile(&image)
     }
 
     pub fn launch(&mut self) -> Result<u32> {
+        let _image_info = self.compile()?;
         let config = DomainConfig {
             max_vcpus: self.vcpus,
             mem_mb: self.mem,

hypha/src/error.rs

@@ -66,3 +66,9 @@ impl From<BackhandError> for HyphaError {
         HyphaError::new(value.to_string().as_str())
     }
 }
+
+impl From<serde_json::Error> for HyphaError {
+    fn from(value: serde_json::Error) -> Self {
+        HyphaError::new(value.to_string().as_str())
+    }
+}

hypha/src/image/cache.rs

@@ -0,0 +1,51 @@
+use crate::image::{ImageInfo, Result};
+use log::debug;
+use oci_spec::image::ImageManifest;
+use std::fs;
+use std::path::{Path, PathBuf};
+
+pub struct ImageCache {
+    cache_dir: PathBuf,
+}
+
+impl ImageCache {
+    pub fn new(cache_dir: &Path) -> Result<ImageCache> {
+        Ok(ImageCache {
+            cache_dir: cache_dir.to_path_buf(),
+        })
+    }
+
+    pub fn recall(&self, digest: &str) -> Result<Option<ImageInfo>> {
+        let mut squashfs_path = self.cache_dir.clone();
+        let mut manifest_path = self.cache_dir.clone();
+        squashfs_path.push(format!("{}.squashfs", digest));
+        manifest_path.push(format!("{}.json", digest));
+        Ok(if squashfs_path.exists() && manifest_path.exists() {
+            let squashfs_metadata = fs::metadata(&squashfs_path)?;
+            let manifest_metadata = fs::metadata(&manifest_path)?;
+            if squashfs_metadata.is_file() && manifest_metadata.is_file() {
+                let manifest_text = fs::read_to_string(&manifest_path)?;
+                let manifest: ImageManifest = serde_json::from_str(manifest_text.as_str())?;
+                debug!("cache hit digest={}", digest);
+                Some(ImageInfo::new(squashfs_path.clone(), manifest)?)
+            } else {
+                None
+            }
+        } else {
+            debug!("cache miss digest={}", digest);
+            None
+        })
+    }
+
+    pub fn store(&self, digest: &str, info: &ImageInfo) -> Result<ImageInfo> {
+        debug!("cache store digest={}", digest);
+        let mut squashfs_path = self.cache_dir.clone();
+        let mut manifest_path = self.cache_dir.clone();
+        squashfs_path.push(format!("{}.squashfs", digest));
+        manifest_path.push(format!("{}.json", digest));
+        fs::copy(&info.squashfs, &squashfs_path)?;
+        let manifest_text = serde_json::to_string_pretty(&info.manifest)?;
+        fs::write(&manifest_path, manifest_text)?;
+        ImageInfo::new(squashfs_path.clone(), info.manifest.clone())
+    }
+}

hypha/src/image/mod.rs

@@ -1,7 +1,10 @@
+pub mod cache;
+
 use crate::error::{HyphaError, Result};
+use crate::image::cache::ImageCache;
 use backhand::{FilesystemWriter, NodeHeader};
 use log::{debug, trace};
-use oci_spec::image::MediaType;
+use oci_spec::image::{ImageManifest, MediaType};
 use ocipkg::distribution::Client;
 use ocipkg::error::Error;
 use ocipkg::{Digest, ImageName};
@@ -13,14 +16,27 @@ use std::path::PathBuf;
 use uuid::Uuid;
 use walkdir::WalkDir;
 
-pub struct ImageCompiler {}
-
-impl ImageCompiler {
-    pub fn new() -> Result<ImageCompiler> {
-        Ok(ImageCompiler {})
+pub struct ImageInfo {
+    pub squashfs: PathBuf,
+    pub manifest: ImageManifest,
 }
 
-    pub fn compile(&self, image: &ImageName) -> Result<String> {
+impl ImageInfo {
+    fn new(squashfs: PathBuf, manifest: ImageManifest) -> Result<ImageInfo> {
+        Ok(ImageInfo { squashfs, manifest })
+    }
+}
+
+pub struct ImageCompiler<'a> {
+    cache: &'a ImageCache,
+}
+
+impl ImageCompiler<'_> {
+    pub fn new(cache: &ImageCache) -> Result<ImageCompiler> {
+        Ok(ImageCompiler { cache })
+    }
+
+    pub fn compile(&self, image: &ImageName) -> Result<ImageInfo> {
         debug!("ImageCompiler compile image={image}");
         let mut tmp_dir = std::env::temp_dir().clone();
         tmp_dir.push(format!("hypha-compile-{}", Uuid::new_v4()));
@@ -29,11 +45,17 @@ impl ImageCompiler {
         fs::create_dir_all(&image_dir)?;
         let mut squash_file = tmp_dir.clone();
         squash_file.push("image.squashfs");
-        self.download(image, &image_dir)?;
-        self.squash(&image_dir, &squash_file)
+        let info = self.download_and_compile(image, &image_dir, &squash_file)?;
+        fs::remove_dir_all(tmp_dir)?;
+        Ok(info)
     }
 
-    fn download(&self, image: &ImageName, image_dir: &PathBuf) -> Result<()> {
+    fn download_and_compile(
+        &self,
+        image: &ImageName,
+        image_dir: &PathBuf,
+        squash_file: &PathBuf,
+    ) -> Result<ImageInfo> {
         debug!(
             "ImageCompiler download image={image}, image_dir={}",
             image_dir.to_str().unwrap()
@@ -43,12 +65,18 @@ impl ImageCompiler {
         } = image;
         let mut client = Client::new(image.registry_url()?, name.clone())?;
         let manifest = client.get_manifest(reference)?;
+        let manifest_serialized = serde_json::to_string(&manifest)?;
+        let manifest_digest = sha256::digest(manifest_serialized);
+        if let Some(cached) = self.cache.recall(&manifest_digest)? {
+            return Ok(cached);
+        }
         for layer in manifest.layers() {
             debug!(
                 "ImageCompiler download start digest={} size={}",
                 layer.digest(),
                 layer.size()
             );
+
             let blob = client.get_blob(&Digest::new(layer.digest())?)?;
             match layer.media_type() {
                 MediaType::ImageLayerGzip => {}
@@ -71,12 +99,14 @@ impl ImageCompiler {
                 layer.digest(),
                 layer.size()
             );
-            return Ok(());
+            self.squash(image_dir, squash_file)?;
+            let info = ImageInfo::new(squash_file.clone(), manifest.clone())?;
+            return self.cache.store(&manifest_digest, &info);
         }
         Err(Error::MissingLayer.into())
     }
 
-    fn squash(&self, image_dir: &PathBuf, squash_file: &PathBuf) -> Result<String> {
+    fn squash(&self, image_dir: &PathBuf, squash_file: &PathBuf) -> Result<()> {
         let mut writer = FilesystemWriter::default();
         let walk = WalkDir::new(image_dir).follow_links(false);
         for entry in walk {
@@ -137,9 +167,8 @@ impl ImageCompiler {
             .ok_or_else(|| HyphaError::new("failed to convert squashfs string"))?;
 
         let mut out = File::create(squash_file)?;
-        trace!("ImageCompiler squash generateI : {}", squash_file_path);
+        trace!("ImageCompiler squash generate: {}", squash_file_path);
         writer.write(&mut out)?;
-
-        Ok(squash_file_path.to_string())
+        Ok(())
     }
 }