Builder: support --prefetch-policy option for merge

Signed-off-by: Yan Song <[email protected]>
dragonflyoss · Mar 28, 2022 · 3438c8d · 3438c8d
1 parent 40c4319
commit 3438c8d
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 8 deletions.
diff --git a/src/bin/nydus-image/main.rs b/src/bin/nydus-image/main.rs
@@ -349,6 +349,16 @@ fn prepare_cmd_args(bti_string: String) -> ArgMatches<'static> {
                         .help("Specify a chunk dictionary for chunk deduplication")
                         .takes_value(true)
                 )
+                .arg(
+                    Arg::with_name("prefetch-policy")
+                        .long("prefetch-policy")
+                        .short("P")
+                        .help("prefetch policy:")
+                        .takes_value(true)
+                        .required(false)
+                        .default_value("none")
+                        .possible_values(&["fs", "blob", "none"]),
+                )
                 .arg(
                     Arg::with_name("SOURCE")
                         .help("bootstrap paths (allow one or more)")
@@ -606,11 +616,7 @@ impl Command {
             }
         }
 
-        let prefetch_policy = matches
-            .value_of("prefetch-policy")
-            .unwrap_or_default()
-            .parse()?;
-        let prefetch = Prefetch::new(prefetch_policy)?;
+        let prefetch = Self::get_prefetch(matches)?;
 
         let mut build_ctx = BuildContext::new(
             blob_id,
@@ -699,7 +705,11 @@ impl Command {
         } else {
             None
         };
+        let prefetch = Self::get_prefetch(matches)?;
+        let mut ctx = BuildContext::default();
+        ctx.prefetch = prefetch;
         Merger::merge(
+            &mut ctx,
             source_bootstrap_paths,
             target_bootstrap_path.to_path_buf(),
             chunk_dict_path,
@@ -943,6 +953,14 @@ impl Command {
         }
     }
 
+    fn get_prefetch(matches: &clap::ArgMatches) -> Result<Prefetch> {
+        let prefetch_policy = matches
+            .value_of("prefetch-policy")
+            .unwrap_or_default()
+            .parse()?;
+        Prefetch::new(prefetch_policy)
+    }
+
     fn get_blob_offset(matches: &clap::ArgMatches) -> Result<u64> {
         match matches.value_of("blob-offset") {
             None => Ok(0),

diff --git a/src/bin/nydus-image/merge.rs b/src/bin/nydus-image/merge.rs
@@ -21,6 +21,7 @@ pub struct Merger {}
 
 impl Merger {
     pub fn merge(
+        ctx: &mut BuildContext,
         sources: Vec<PathBuf>,
         target: PathBuf,
         chunk_dict: Option<PathBuf>,
@@ -97,10 +98,9 @@ impl Merger {
         let mut bootstrap = Bootstrap::new()?;
         let storage = ArtifactStorage::SingleFile(target.clone());
         let mut bootstrap_ctx = BootstrapContext::new(storage, false)?;
-        let mut ctx = BuildContext::default();
-        bootstrap.build(&mut ctx, &mut bootstrap_ctx, &mut tree)?;
+        bootstrap.build(ctx, &mut bootstrap_ctx, &mut tree)?;
         let blob_table = blob_mgr.to_blob_table(&ctx)?;
-        bootstrap.dump(&mut ctx, &mut bootstrap_ctx, &blob_table)?;
+        bootstrap.dump(ctx, &mut bootstrap_ctx, &blob_table)?;
 
         Ok(())
     }