feat: oci tar format, bit-perfect disk storage for config and manifes…

…t, concurrent image pulls (#88) * oci: retain bit-perfect copies of manifest and config on disk * feat: oci tar format support * feat: concurrent image pulls
edera-dev · Apr 16, 2024 · e450ebd · e450ebd
1 parent 79f7742
commit e450ebd
Show file tree

Hide file tree

Showing 21 changed files with 492 additions and 143 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/ctl/src/cli/launch.rs b/crates/ctl/src/cli/launch.rs
@@ -1,13 +1,13 @@
 use std::collections::HashMap;
 
 use anyhow::Result;
-use clap::Parser;
+use clap::{Parser, ValueEnum};
 use krata::{
     events::EventStream,
     v1::{
         common::{
-            guest_image_spec::Image, GuestImageSpec, GuestOciImageFormat, GuestOciImageSpec,
-            GuestSpec, GuestStatus, GuestTaskSpec, GuestTaskSpecEnvVar,
+            guest_image_spec::Image, GuestImageSpec, GuestOciImageSpec, GuestSpec, GuestStatus,
+            GuestTaskSpec, GuestTaskSpecEnvVar, OciImageFormat,
         },
         control::{
             control_service_client::ControlServiceClient, watch_events_reply::Event,
@@ -21,13 +21,17 @@ use tonic::{transport::Channel, Request};
 
 use crate::{console::StdioConsoleStream, pull::pull_interactive_progress};
 
-use super::pull::PullImageFormat;
+#[derive(ValueEnum, Clone, Debug, PartialEq, Eq)]
+pub enum LaunchImageFormat {
+    Squashfs,
+    Erofs,
+}
 
 #[derive(Parser)]
 #[command(about = "Launch a new guest")]
 pub struct LauchCommand {
     #[arg(short = 'S', long, default_value = "squashfs", help = "Image format")]
-    image_format: PullImageFormat,
+    image_format: LaunchImageFormat,
     #[arg(short, long, help = "Name of the guest")]
     name: Option<String>,
     #[arg(
@@ -78,8 +82,8 @@ impl LauchCommand {
             .pull_image(PullImageRequest {
                 image: self.oci.clone(),
                 format: match self.image_format {
-                    PullImageFormat::Squashfs => GuestOciImageFormat::Squashfs.into(),
-                    PullImageFormat::Erofs => GuestOciImageFormat::Erofs.into(),
+                    LaunchImageFormat::Squashfs => OciImageFormat::Squashfs.into(),
+                    LaunchImageFormat::Erofs => OciImageFormat::Erofs.into(),
                 },
             })
             .await?;

diff --git a/crates/ctl/src/cli/pull.rs b/crates/ctl/src/cli/pull.rs
@@ -1,7 +1,7 @@
 use anyhow::Result;
 use clap::{Parser, ValueEnum};
 use krata::v1::{
-    common::GuestOciImageFormat,
+    common::OciImageFormat,
     control::{control_service_client::ControlServiceClient, PullImageRequest},
 };
 
@@ -13,6 +13,7 @@ use crate::pull::pull_interactive_progress;
 pub enum PullImageFormat {
     Squashfs,
     Erofs,
+    Tar,
 }
 
 #[derive(Parser)]
@@ -30,8 +31,9 @@ impl PullCommand {
             .pull_image(PullImageRequest {
                 image: self.image.clone(),
                 format: match self.image_format {
-                    PullImageFormat::Squashfs => GuestOciImageFormat::Squashfs.into(),
-                    PullImageFormat::Erofs => GuestOciImageFormat::Erofs.into(),
+                    PullImageFormat::Squashfs => OciImageFormat::Squashfs.into(),
+                    PullImageFormat::Erofs => OciImageFormat::Erofs.into(),
+                    PullImageFormat::Tar => OciImageFormat::Tar.into(),
                 },
             })
             .await?;

diff --git a/crates/daemon/Cargo.toml b/crates/daemon/Cargo.toml
@@ -23,6 +23,7 @@ krata-runtime = { path = "../runtime", version = "^0.0.9" }
 log = { workspace = true }
 prost = { workspace = true }
 redb = { workspace = true }
+scopeguard = { workspace = true }
 signal-hook = { workspace = true }
 tokio = { workspace = true }
 tokio-stream = { workspace = true }

diff --git a/crates/daemon/src/control.rs b/crates/daemon/src/control.rs
@@ -6,7 +6,7 @@ use krata::{
         IdmMetricsRequest,
     },
     v1::{
-        common::{Guest, GuestOciImageFormat, GuestState, GuestStatus},
+        common::{Guest, GuestState, GuestStatus, OciImageFormat},
         control::{
             control_service_server::ControlService, ConsoleDataReply, ConsoleDataRequest,
             CreateGuestReply, CreateGuestRequest, DestroyGuestReply, DestroyGuestRequest,
@@ -18,7 +18,7 @@ use krata::{
 };
 use krataoci::{
     name::ImageName,
-    packer::{service::OciPackerService, OciImagePacked, OciPackedFormat},
+    packer::{service::OciPackerService, OciPackedFormat, OciPackedImage},
     progress::{OciProgress, OciProgressContext},
 };
 use std::{pin::Pin, str::FromStr};
@@ -90,8 +90,8 @@ enum ConsoleDataSelect {
 }
 
 enum PullImageSelect {
-    Progress(usize),
-    Completed(Result<Result<OciImagePacked, anyhow::Error>, JoinError>),
+    Progress(Option<OciProgress>),
+    Completed(Result<Result<OciPackedImage, anyhow::Error>, JoinError>),
 }
 
 #[tonic::async_trait]
@@ -362,36 +362,51 @@ impl ControlService for DaemonControlService {
             message: err.to_string(),
         })?;
         let format = match request.format() {
-            GuestOciImageFormat::Unknown => OciPackedFormat::Squashfs,
-            GuestOciImageFormat::Squashfs => OciPackedFormat::Squashfs,
-            GuestOciImageFormat::Erofs => OciPackedFormat::Erofs,
+            OciImageFormat::Unknown => OciPackedFormat::Squashfs,
+            OciImageFormat::Squashfs => OciPackedFormat::Squashfs,
+            OciImageFormat::Erofs => OciPackedFormat::Erofs,
+            OciImageFormat::Tar => OciPackedFormat::Tar,
         };
-        let (sender, mut receiver) = channel::<OciProgress>(100);
-        let context = OciProgressContext::new(sender);
-
+        let (context, mut receiver) = OciProgressContext::create();
         let our_packer = self.packer.clone();
 
         let output = try_stream! {
             let mut task = tokio::task::spawn(async move {
                 our_packer.request(name, format, context).await
             });
+            let abort_handle = task.abort_handle();
+            let _task_cancel_guard = scopeguard::guard(abort_handle, |handle| {
+                handle.abort();
+            });
+
             loop {
-                let mut progresses = Vec::new();
                 let what = select! {
-                    x = receiver.recv_many(&mut progresses, 10) => PullImageSelect::Progress(x),
+                    x = receiver.recv() => PullImageSelect::Progress(x.ok()),
                     x = &mut task => PullImageSelect::Completed(x),
                 };
                 match what {
-                    PullImageSelect::Progress(count) => {
-                        if count > 0 {
-                            let progress = progresses.remove(progresses.len() - 1);
-                            let reply = PullImageReply {
-                                progress: Some(convert_oci_progress(progress)),
-                                digest: String::new(),
-                                format: GuestOciImageFormat::Unknown.into(),
-                            };
-                            yield reply;
+                    PullImageSelect::Progress(Some(mut progress)) => {
+                        let mut drain = 0;
+                        loop {
+                            if drain >= 10 {
+                                break;
+                            }
+
+                            if let Ok(latest) = receiver.try_recv() {
+                                progress = latest;
+                            } else {
+                                break;
+                            }
+
+                            drain += 1;
                         }
+
+                        let reply = PullImageReply {
+                            progress: Some(convert_oci_progress(progress)),
+                            digest: String::new(),
+                            format: OciImageFormat::Unknown.into(),
+                        };
+                        yield reply;
                     },
 
                     PullImageSelect::Completed(result) => {
@@ -405,13 +420,18 @@ impl ControlService for DaemonControlService {
                             progress: None,
                             digest: packed.digest,
                             format: match packed.format {
-                                OciPackedFormat::Squashfs => GuestOciImageFormat::Squashfs.into(),
-                                OciPackedFormat::Erofs => GuestOciImageFormat::Erofs.into(),
+                                OciPackedFormat::Squashfs => OciImageFormat::Squashfs.into(),
+                                OciPackedFormat::Erofs => OciImageFormat::Erofs.into(),
+                                _ => OciImageFormat::Unknown.into(),
                             },
                         };
                         yield reply;
                         break;
                     },
+
+                    _ => {
+                        continue;
+                    }
                 }
             }
         };

diff --git a/crates/daemon/src/reconcile/guest.rs b/crates/daemon/src/reconcile/guest.rs
@@ -9,7 +9,7 @@ use krata::launchcfg::LaunchPackedFormat;
 use krata::v1::{
     common::{
         guest_image_spec::Image, Guest, GuestErrorInfo, GuestExitInfo, GuestNetworkState,
-        GuestOciImageFormat, GuestState, GuestStatus,
+        GuestState, GuestStatus, OciImageFormat,
     },
     control::GuestChangedEvent,
 };
@@ -244,9 +244,12 @@ impl GuestReconciler {
             .recall(
                 &oci.digest,
                 match oci.format() {
-                    GuestOciImageFormat::Unknown => OciPackedFormat::Squashfs,
-                    GuestOciImageFormat::Squashfs => OciPackedFormat::Squashfs,
-                    GuestOciImageFormat::Erofs => OciPackedFormat::Erofs,
+                    OciImageFormat::Unknown => OciPackedFormat::Squashfs,
+                    OciImageFormat::Squashfs => OciPackedFormat::Squashfs,
+                    OciImageFormat::Erofs => OciPackedFormat::Erofs,
+                    OciImageFormat::Tar => {
+                        return Err(anyhow!("tar image format is not supported for guests"));
+                    }
                 },
             )
             .await?;

diff --git a/crates/krata/proto/krata/v1/common.proto b/crates/krata/proto/krata/v1/common.proto
@@ -29,15 +29,17 @@ message GuestImageSpec {
     }
 }
 
-enum GuestOciImageFormat {
-    GUEST_OCI_IMAGE_FORMAT_UNKNOWN = 0;
-    GUEST_OCI_IMAGE_FORMAT_SQUASHFS = 1;
-    GUEST_OCI_IMAGE_FORMAT_EROFS = 2;
+enum OciImageFormat {
+    OCI_IMAGE_FORMAT_UNKNOWN = 0;
+    OCI_IMAGE_FORMAT_SQUASHFS = 1;
+    OCI_IMAGE_FORMAT_EROFS = 2;
+    // Tar format is not launchable, and is intended for kernel images.
+    OCI_IMAGE_FORMAT_TAR = 3;
 }
 
 message GuestOciImageSpec {
     string digest = 1;
-    GuestOciImageFormat format = 2;
+    OciImageFormat format = 2;
 }
 
 message GuestTaskSpec {

diff --git a/crates/krata/proto/krata/v1/control.proto b/crates/krata/proto/krata/v1/control.proto
@@ -124,11 +124,11 @@ message PullImageProgress {
 
 message PullImageRequest {
     string image = 1;
-    krata.v1.common.GuestOciImageFormat format = 2;
+    krata.v1.common.OciImageFormat format = 2;
 }
 
 message PullImageReply {
     PullImageProgress progress = 1;
     string digest = 2;
-    krata.v1.common.GuestOciImageFormat format = 3;
+    krata.v1.common.OciImageFormat format = 3;
 }
diff --git a/crates/oci/examples/squashify.rs b/crates/oci/examples/squashify.rs
@@ -5,10 +5,10 @@ use env_logger::Env;
 use krataoci::{
     name::ImageName,
     packer::{service::OciPackerService, OciPackedFormat},
-    progress::{OciProgress, OciProgressContext},
+    progress::OciProgressContext,
     registry::OciPlatform,
 };
-use tokio::{fs, sync::mpsc::channel};
+use tokio::fs;
 
 #[tokio::main]
 async fn main() -> Result<()> {
@@ -22,14 +22,28 @@ async fn main() -> Result<()> {
         fs::create_dir(&cache_dir).await?;
     }
 
-    let (sender, mut receiver) = channel::<OciProgress>(100);
+    let (context, mut receiver) = OciProgressContext::create();
     tokio::task::spawn(async move {
         loop {
-            let mut progresses = Vec::new();
-            let _ = receiver.recv_many(&mut progresses, 100).await;
-            let Some(progress) = progresses.last() else {
-                continue;
+            let Ok(mut progress) = receiver.recv().await else {
+                return;
             };
+
+            let mut drain = 0;
+            loop {
+                if drain >= 10 {
+                    break;
+                }
+
+                if let Ok(latest) = receiver.try_recv() {
+                    progress = latest;
+                } else {
+                    break;
+                }
+
+                drain += 1;
+            }
+
             println!("phase {:?}", progress.phase);
             for (id, layer) in &progress.layers {
                 println!(
@@ -39,7 +53,6 @@ async fn main() -> Result<()> {
             }
         }
     });
-    let context = OciProgressContext::new(sender);
     let service = OciPackerService::new(seed, &cache_dir, OciPlatform::current())?;
     let packed = service
         .request(image.clone(), OciPackedFormat::Squashfs, context)