support generating new image with separate blob for prefetched files

builder/src/optimize_prefetch.rs

@@ -8,6 +8,7 @@ use crate::BlobManager;
 use crate::Bootstrap;
 use crate::BootstrapManager;
 use crate::BuildContext;
+use crate::BuildOutput;
 use crate::ChunkSource;
 use crate::ConversionType;
 use crate::NodeChunk;
@@ -18,7 +19,6 @@ use crate::TreeNode;
 use anyhow::Context;
 use anyhow::{Ok, Result};
 use nydus_api::ConfigV2;
-use nydus_rafs::metadata::layout::v6::RafsV6BlobTable;
 use nydus_rafs::metadata::layout::RafsBlobTable;
 use nydus_rafs::metadata::RafsSuper;
 use nydus_rafs::metadata::RafsVersion;
@@ -28,8 +28,7 @@ use nydus_storage::meta::BlobChunkInfoV1Ondisk;
 use nydus_utils::compress;
 use sha2::Digest;
 use std::fs::File;
-use std::io::Read;
-use std::io::Seek;
+use std::io::{Read, Seek, Write};
 use std::mem::size_of;
 use std::sync::Arc;
 pub struct OptimizePrefetch {}
@@ -52,6 +51,7 @@ impl PrefetchBlobState {
             ctx.blob_features,
         );
         blob_info.set_compressor(ctx.compressor);
+        blob_info.set_separated_with_prefetch_files_feature(true);
         let mut blob_ctx = BlobContext::from(ctx, &blob_info, ChunkSource::Build)?;
         blob_ctx.blob_meta_info_enabled = true;
         let blob_writer = ArtifactWriter::new(crate::ArtifactStorage::FileDir(
@@ -72,13 +72,16 @@ impl OptimizePrefetch {
         tree: &mut Tree,
         ctx: &mut BuildContext,
         bootstrap_mgr: &mut BootstrapManager,
-        blob_table: &mut RafsV6BlobTable,
+        blob_table: &mut RafsBlobTable,
         blobs_dir_path: PathBuf,
         prefetch_nodes: Vec<TreeNode>,
-    ) -> Result<()> {
+    ) -> Result<BuildOutput> {
         // create a new blob for prefetch layer
-        let blob_layer_num = blob_table.entries.len();
 
+        let blob_layer_num = match blob_table {
+            RafsBlobTable::V5(table) => table.get_all().len(),
+            RafsBlobTable::V6(table) => table.get_all().len(),
+        };
         let mut blob_state = PrefetchBlobState::new(&ctx, blob_layer_num as u32, &blobs_dir_path)?;
         let mut batch = BatchContextGenerator::new(0)?;
         for node in &prefetch_nodes {
@@ -92,66 +95,53 @@ impl OptimizePrefetch {
             )?;
         }
 
-        Self::dump_blob(ctx, blob_table, &mut blob_state)?;
+        let blob_mgr = Self::dump_blob(ctx, blob_table, &mut blob_state)?;
 
         debug!("prefetch blob id: {}", ctx.blob_id);
 
         Self::build_dump_bootstrap(tree, ctx, bootstrap_mgr, blob_table)?;
-        Ok(())
+        BuildOutput::new(&blob_mgr, &bootstrap_mgr.bootstrap_storage)
     }
 
     fn build_dump_bootstrap(
         tree: &mut Tree,
         ctx: &mut BuildContext,
         bootstrap_mgr: &mut BootstrapManager,
-        blob_table: &mut RafsV6BlobTable,
+        blob_table: &mut RafsBlobTable,
     ) -> Result<()> {
         let mut bootstrap_ctx = bootstrap_mgr.create_ctx()?;
         let mut bootstrap = Bootstrap::new(tree.clone())?;
 
         // Build bootstrap
         bootstrap.build(ctx, &mut bootstrap_ctx)?;
 
-        // Verify and update prefetch blob
-        assert!(
-            blob_table
-                .entries
-                .iter()
-                .filter(|blob| blob.blob_id() == "prefetch-blob")
-                .count()
-                == 1,
-            "Expected exactly one prefetch-blob"
-        );
-
-        // Rewrite prefetch blob id
-        blob_table
-            .entries
-            .iter_mut()
-            .filter(|blob| blob.blob_id() == "prefetch-blob")
-            .for_each(|blob| {
-                let mut info = (**blob).clone();
-                info.set_blob_id(ctx.blob_id.clone());
-                *blob = Arc::new(info);
-            });
-
-        // Dump bootstrap
-        let blob_table_withprefetch = RafsBlobTable::V6(blob_table.clone());
+        let blob_table_withprefetch = match blob_table {
+            RafsBlobTable::V5(table) => RafsBlobTable::V5(table.clone()),
+            RafsBlobTable::V6(table) => RafsBlobTable::V6(table.clone()),
+        };
         bootstrap.dump(
             ctx,
             &mut bootstrap_mgr.bootstrap_storage,
             &mut bootstrap_ctx,
             &blob_table_withprefetch,
         )?;
-
         Ok(())
     }
 
     fn dump_blob(
         ctx: &mut BuildContext,
-        blob_table: &mut RafsV6BlobTable,
+        blob_table: &mut RafsBlobTable,
         blob_state: &mut PrefetchBlobState,
-    ) -> Result<()> {
-        blob_table.entries.push(blob_state.blob_info.clone().into());
+    ) -> Result<BlobManager> {
+        match blob_table {
+            RafsBlobTable::V5(table) => {
+                table.entries.push(blob_state.blob_info.clone().into());
+            }
+            RafsBlobTable::V6(table) => {
+                table.entries.push(blob_state.blob_info.clone().into());
+            }
+        }
+
         let mut blob_mgr = BlobManager::new(ctx.digester);
         blob_mgr.add_blob(blob_state.blob_ctx.clone());
         blob_mgr.set_current_blob_index(0);
@@ -168,30 +158,57 @@ impl OptimizePrefetch {
             .1
             .blob_id
             .clone();
-        Ok(())
+
+        let entries = match blob_table {
+            RafsBlobTable::V5(table) => table.get_all(),
+            RafsBlobTable::V6(table) => table.get_all(),
+        };
+
+        // Verify and update prefetch blob
+        assert!(
+            entries
+                .iter()
+                .filter(|blob| blob.blob_id() == "prefetch-blob")
+                .count()
+                == 1,
+            "Expected exactly one prefetch-blob"
+        );
+        // Rewrite prefetch blob id
+        match blob_table {
+            RafsBlobTable::V5(table) => {
+                rewrite_blob_id(&mut table.entries, "prefetch-blob", ctx.blob_id.clone())
+            }
+            RafsBlobTable::V6(table) => {
+                rewrite_blob_id(&mut table.entries, "prefetch-blob", ctx.blob_id.clone())
+            }
+        }
+        Ok(blob_mgr)
     }
 
     fn process_prefetch_node(
         tree: &mut Tree,
         node: &TreeNode,
         prefetch_state: &mut PrefetchBlobState,
         batch: &mut BatchContextGenerator,
-        blob_table: &RafsV6BlobTable,
+        blob_table: &RafsBlobTable,
         blobs_dir_path: &Path,
     ) -> Result<()> {
         let tree_node = tree
             .get_node_mut(&node.borrow().path())
             .ok_or(anyhow!("failed to get node"))?
             .node
             .as_ref();
+        let entries = match blob_table {
+            RafsBlobTable::V5(table) => table.get_all(),
+            RafsBlobTable::V6(table) => table.get_all(),
+        };
         let blob_id = tree_node
             .borrow()
             .chunks
             .first()
-            .and_then(|chunk| blob_table.entries.get(chunk.inner.blob_index() as usize))
+            .and_then(|chunk| entries.get(chunk.inner.blob_index() as usize).cloned())
             .map(|entry| entry.blob_id())
             .ok_or(anyhow!("failed to get blob id"))?;
-
         let mut blob_file = Arc::new(File::open(blobs_dir_path.join(blob_id))?);
 
         tree_node.borrow_mut().layer_idx = prefetch_state.blob_info.blob_index() as u16;
@@ -247,6 +264,17 @@ impl OptimizePrefetch {
     }
 }
 
+fn rewrite_blob_id(entries: &mut [Arc<BlobInfo>], blob_id: &str, new_blob_id: String) {
+    entries
+        .iter_mut()
+        .filter(|blob| blob.blob_id() == blob_id)
+        .for_each(|blob| {
+            let mut info = (**blob).clone();
+            info.set_blob_id(new_blob_id.clone());
+            *blob = Arc::new(info);
+        });
+}
+
 pub fn update_ctx_from_bootstrap(
     ctx: &mut BuildContext,
     config: Arc<ConfigV2>,

contrib/nydusify/cmd/nydusify.go

@@ -16,6 +16,8 @@ import (
 	"runtime"
 	"strings"
 
+	"github.com/dragonflyoss/nydus/contrib/nydusify/pkg/optimizer"
+
 	"github.com/containerd/containerd/reference/docker"
 	"github.com/distribution/reference"
 	"github.com/dustin/go-humanize"
@@ -1160,6 +1162,97 @@ func main() {
 				return copier.Copy(context.Background(), opt)
 			},
 		},
+		{
+			Name:  "optimize",
+			Usage: "Optimize a source nydus image and push to the target",
+			Flags: []cli.Flag{
+				&cli.StringFlag{
+					Name:     "source",
+					Required: true,
+					Usage:    "Source (Nydus) image reference",
+					EnvVars:  []string{"SOURCE"},
+				},
+				&cli.StringFlag{
+					Name:     "target",
+					Required: true,
+					Usage:    "Target (Nydus) image reference",
+					EnvVars:  []string{"TARGET"},
+				},
+				&cli.BoolFlag{
+					Name:     "source-insecure",
+					Required: false,
+					Usage:    "Skip verifying server certs for HTTPS source registry",
+					EnvVars:  []string{"SOURCE_INSECURE"},
+				},
+				&cli.BoolFlag{
+					Name:     "target-insecure",
+					Required: false,
+					Usage:    "Skip verifying server certs for HTTPS target registry",
+					EnvVars:  []string{"TARGET_INSECURE"},
+				},
+
+				&cli.StringFlag{
+					Name:    "policy",
+					Value:   "separated-blob-with-prefetch-files",
+					Usage:   "Specify the optimizing way",
+					EnvVars: []string{"OPTIMIZE_POLICY"},
+				},
+				&cli.StringFlag{
+					Name:     "prefetch-files",
+					Required: false,
+					Usage:    "File path to include prefetch files for optimization",
+					EnvVars:  []string{"PREFETCH_FILES"},
+				},
+
+				&cli.StringFlag{
+					Name:    "work-dir",
+					Value:   "./tmp",
+					Usage:   "Working directory for image optimization",
+					EnvVars: []string{"WORK_DIR"},
+				},
+
+				&cli.StringFlag{
+					Name:    "nydus-image",
+					Value:   "nydus-image",
+					Usage:   "Path to the nydus-image binary, default to search in PATH",
+					EnvVars: []string{"NYDUS_IMAGE"},
+				},
+
+				&cli.StringFlag{
+					Name:  "push-chunk-size",
+					Value: "0MB",
+					Usage: "Chunk size for pushing a blob layer in chunked",
+				},
+			},
+			Action: func(c *cli.Context) error {
+				setupLogLevel(c)
+
+				pushChunkSize, err := humanize.ParseBytes(c.String("push-chunk-size"))
+				if err != nil {
+					return errors.Wrap(err, "invalid --push-chunk-size option")
+				}
+				if pushChunkSize > 0 {
+					logrus.Infof("will push layer with chunk size %s", c.String("push-chunk-size"))
+				}
+				opt := optimizer.Opt{
+					WorkDir:        c.String("work-dir"),
+					NydusImagePath: c.String("nydus-image"),
+
+					Source:         c.String("source"),
+					Target:         c.String("target"),
+					SourceInsecure: c.Bool("source-insecure"),
+					TargetInsecure: c.Bool("target-insecure"),
+
+					AllPlatforms: c.Bool("all-platforms"),
+					Platforms:    c.String("platform"),
+
+					PushChunkSize:     int64(pushChunkSize),
+					PrefetchFilesPath: c.String("prefetch-files"),
+				}
+
+				return optimizer.Optimize(context.Background(), opt)
+			},
+		},
 		{
 			Name:  "commit",
 			Usage: "Create and push a new nydus image from a container's changes that use a nydus image",