From 80545863a7c0c651b36b5c252756c55ee88d34cf Mon Sep 17 00:00:00 2001
From: Lin Wang <l.wang@mail.dlut.edu.cn>
Date: Tue, 19 Sep 2023 11:19:37 +0800
Subject: [PATCH] =?UTF-8?q?nydus-image=EF=BC=9AOptimize=20Chunkdict=20Save?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Refactor the Deduplicate implementation to only
initialize config when inserting chunk data.
Simplify code for better maintainability.

Signed-off-by: Lin Wang <l.wang@mail.dlut.edu.cn>
---
 src/bin/nydus-image/deduplicate.rs | 27 ++++++++++++++++++++-------
 src/bin/nydus-image/main.rs        | 28 ++++++----------------------
 2 files changed, 26 insertions(+), 29 deletions(-)

diff --git a/src/bin/nydus-image/deduplicate.rs b/src/bin/nydus-image/deduplicate.rs
index f2319162e52..1ad9d0c9c93 100644
--- a/src/bin/nydus-image/deduplicate.rs
+++ b/src/bin/nydus-image/deduplicate.rs
@@ -128,26 +128,35 @@ impl Database for SqliteDatabase {
 }
 
 pub struct Deduplicate<D: Database + Send + Sync> {
-    sb: RafsSuper,
+    sb: Option<RafsSuper>,
     db: D,
 }
 
 const IN_MEMORY_DB_URL: &str = ":memory:";
 
 impl Deduplicate<SqliteDatabase> {
-    pub fn new(bootstrap_path: &Path, config: Arc<ConfigV2>, db_url: &str) -> anyhow::Result<Self> {
-        let (sb, _) = RafsSuper::load_from_file(bootstrap_path, config, false)?;
+    pub fn new(db_url: &str) -> anyhow::Result<Self> {
         let db = if db_url == IN_MEMORY_DB_URL {
             SqliteDatabase::new_in_memory()?
         } else {
             SqliteDatabase::new(db_url)?
         };
-        Ok(Self { sb, db })
+        Ok(Self { sb: None, db })
     }
 
-    pub fn save_metadata(&mut self) -> anyhow::Result<Vec<Arc<BlobInfo>>> {
+    pub fn save_metadata(
+        &mut self,
+        bootstrap_path: &Path,
+        config: Arc<ConfigV2>,
+    ) -> anyhow::Result<Vec<Arc<BlobInfo>>> {
+        let (sb, _) = RafsSuper::load_from_file(bootstrap_path, config, false)?;
+        self.sb = Some(sb);
         self.create_tables()?;
-        let blob_infos = self.sb.superblock.get_blob_infos();
+        let sb = self
+            .sb
+            .as_ref()
+            .ok_or_else(|| anyhow!("RafsSuper is not initialized"))?;
+        let blob_infos = sb.superblock.get_blob_infos();
         self.insert_blobs(&blob_infos)?;
         self.insert_chunks(&blob_infos)?;
         Ok(blob_infos)
@@ -195,7 +204,11 @@ impl Deduplicate<SqliteDatabase> {
             }
             Ok(())
         };
-        let tree = Tree::from_bootstrap(&self.sb, &mut ())
+        let sb = self
+            .sb
+            .as_ref()
+            .ok_or_else(|| anyhow!("RafsSuper is not initialized"))?;
+        let tree = Tree::from_bootstrap(sb, &mut ())
             .context("Failed to load bootstrap for deduplication.")?;
         tree.walk_dfs_pre(process_chunk)?;
         Ok(())
diff --git a/src/bin/nydus-image/main.rs b/src/bin/nydus-image/main.rs
index 200b5cc3784..b42c110b3b6 100644
--- a/src/bin/nydus-image/main.rs
+++ b/src/bin/nydus-image/main.rs
@@ -34,7 +34,7 @@ use nydus_builder::{
 use nydus_rafs::metadata::{RafsSuper, RafsSuperConfig, RafsVersion};
 use nydus_storage::backend::localfs::LocalFs;
 use nydus_storage::backend::BlobBackend;
-use nydus_storage::device::{BlobFeatures, BlobInfo};
+use nydus_storage::device::BlobFeatures;
 use nydus_storage::factory::BlobFactory;
 use nydus_storage::meta::{format_blob_features, BatchContextGenerator};
 use nydus_storage::{RAFS_DEFAULT_CHUNK_SIZE, RAFS_MAX_CHUNK_SIZE};
@@ -369,7 +369,6 @@ fn prepare_cmd_args(bti_string: &'static str) -> App {
                             .short('B')
                             .long("bootstrap")
                             .help("File path of RAFS meta blob/bootstrap")
-                            .conflicts_with("BOOTSTRAP")
                             .required(false),
                         )
                     .arg(
@@ -399,7 +398,7 @@ fn prepare_cmd_args(bti_string: &'static str) -> App {
                     )
                     .arg(arg_output_json.clone())
             )
-    );
+                );
 
     let app = app.subcommand(
         App::new("merge")
@@ -1151,32 +1150,17 @@ impl Command {
             bail!("Invalid database URL: {}", db_url);
         }
 
-        let blobs: Vec<Arc<BlobInfo>> = match db_strs[0] {
+        match db_strs[0] {
             "sqlite" => {
                 let mut deduplicate: Deduplicate<SqliteDatabase> =
-                    Deduplicate::<SqliteDatabase>::new(bootstrap_path, config, db_strs[1])?;
-                deduplicate.save_metadata()?
+                    Deduplicate::<SqliteDatabase>::new(db_strs[1])?;
+                deduplicate.save_metadata(bootstrap_path, config)?
             }
             _ => {
                 bail!("Unsupported database type: {}, please use a valid database URI, such as 'sqlite:///path/to/database.db'.", db_strs[0])
             }
         };
-        info!("RAFS filesystem metadata is saved:");
-
-        let mut blob_ids = Vec::new();
-        for (idx, blob) in blobs.iter().enumerate() {
-            info!(
-                "\t {}: {}, compressed data size 0x{:x}, compressed file size 0x{:x}, uncompressed file size 0x{:x}, chunks: 0x{:x}, features: {}.",
-                idx,
-                blob.blob_id(),
-                blob.compressed_data_size(),
-                blob.compressed_size(),
-                blob.uncompressed_size(),
-                blob.chunk_count(),
-                format_blob_features(blob.features()),
-            );
-            blob_ids.push(blob.blob_id().to_string());
-        }
+        info!("Chunkdict metadata is saved at: {:?}", db_url);
         Ok(())
     }