From d3bf9fb4d51217802999c363572f0ffff06c2715 Mon Sep 17 00:00:00 2001
From: Zhao Yuan <1627990440@qq.com>
Date: Mon, 20 Nov 2023 13:25:43 +0000
Subject: [PATCH] Add function that nydusify invoke "nydus-image chunkdict
 generate"

Signed-off-by: Zhao Yuan <1627990440@qq.com>
---
 contrib/nydusify/pkg/build/builder.go         | 21 ++++++++++
 .../pkg/chunkdict/generator/generator.go      | 38 +++++++++++++++----
 2 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/contrib/nydusify/pkg/build/builder.go b/contrib/nydusify/pkg/build/builder.go
index 564f57454ae..4c8ec68e13f 100644
--- a/contrib/nydusify/pkg/build/builder.go
+++ b/contrib/nydusify/pkg/build/builder.go
@@ -43,6 +43,11 @@ type CompactOption struct {
 
 type SaveOption struct {
 	BootstrapPath string
+	DatabasePath  string
+}
+
+type GenerateOption struct {
+	DatabasePath string
 }
 
 type Builder struct {
@@ -157,6 +162,22 @@ func (builder *Builder) Save(option SaveOption) error {
 		"warn",
 		"--bootstrap",
 		option.BootstrapPath,
+		"--database",
+		option.DatabasePath,
+	}
+	return builder.run(args, "")
+}
+
+// Generate calls `nydus-image chunkdict generate` to get chunkdict
+func (builder *Builder) Generate(option GenerateOption) error {
+	logrus.Infof("Invoking 'nydus-image chunkdict generate' subcommand")
+	args := []string{
+		"chunkdict",
+		"generate",
+		"--log-level",
+		"warn",
+		"--database",
+		option.DatabasePath,
 	}
 	return builder.run(args, "")
 }
diff --git a/contrib/nydusify/pkg/chunkdict/generator/generator.go b/contrib/nydusify/pkg/chunkdict/generator/generator.go
index e9bcb4ae6b5..5b2b13bd94b 100644
--- a/contrib/nydusify/pkg/chunkdict/generator/generator.go
+++ b/contrib/nydusify/pkg/chunkdict/generator/generator.go
@@ -69,6 +69,9 @@ func (generator *Generator) Generate(ctx context.Context) error {
 			}
 		}
 	}
+	if err := generator.deduplicating(ctx); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -81,26 +84,47 @@ func (generator *Generator) save(ctx context.Context, index int) error {
 
 	// Create a directory to store the image bootstrap
 	nydusImageName := strings.Replace(generator.Sources[index], "/", ":", -1)
-	folderPath := filepath.Join(generator.WorkDir, nydusImageName)
-	if err := os.MkdirAll(folderPath, fs.ModePerm); err != nil {
+	bootstrapFolderPath := filepath.Join(generator.WorkDir, nydusImageName)
+	if err := os.MkdirAll(bootstrapFolderPath, fs.ModePerm); err != nil {
 		return errors.Wrap(err, "creat work directory")
 	}
-	if err := generator.Output(ctx, sourceParsed, folderPath, index); err != nil {
+	if err := generator.Output(ctx, sourceParsed, bootstrapFolderPath, index); err != nil {
 		return errors.Wrap(err, "output image information")
 	}
 
+	databaseName := "chunkdict.db"
+	databaseType := "sqlite"
+	currentDir, _ := os.Getwd()
+	DatabasePath := databaseType + "://" + filepath.Join(currentDir, generator.WorkDir, databaseName)
+
 	// Invoke "nydus-image save" command
 	builder := build.NewBuilder(generator.NydusImagePath)
 	if err := builder.Save(build.SaveOption{
-		BootstrapPath: filepath.Join(folderPath, "nydus_bootstrap"),
+		BootstrapPath: filepath.Join(bootstrapFolderPath, "nydus_bootstrap"),
+		DatabasePath:  DatabasePath,
 	}); err != nil {
 		return errors.Wrap(err, "invalid nydus bootstrap format")
 	}
 
-	logrus.Infof("Save chunk information from image %s", generator.sourcesParser[index].Remote.Ref)
+	logrus.Infof("Saving chunk information from image %s", generator.sourcesParser[index].Remote.Ref)
+
+	// if err := os.RemoveAll(folderPath); err != nil {
+	// 	return errors.Wrap(err, "remove work directory")
+	// }
+	return nil
+}
 
-	if err := os.RemoveAll(folderPath); err != nil {
-		return errors.Wrap(err, "remove work directory")
+func (generator *Generator) deduplicating(ctx context.Context) error {
+	builder := build.NewBuilder(generator.NydusImagePath)
+	currentDir, _ := os.Getwd()
+
+	databaseName := "chunkdict.db"
+	databaseType := "sqlite"
+	DatabasePath := databaseType + "://" + filepath.Join(currentDir, generator.WorkDir, databaseName)
+	if err := builder.Generate(build.GenerateOption{
+		DatabasePath: DatabasePath,
+	}); err != nil {
+		return errors.Wrap(err, "invalid nydus bootstrap format")
 	}
 	return nil
 }