Merge invoked subcomand in nydusify and add smoke test

Signed-off-by: Zhao Yuan <[email protected]>
dragonflyoss · Dec 8, 2023 · 196d707 · 196d707
1 parent bed13de
commit 196d707
Show file tree

Hide file tree

Showing 6 changed files with 106 additions and 241 deletions.
diff --git a/contrib/nydusify/pkg/build/builder.go b/contrib/nydusify/pkg/build/builder.go
@@ -41,13 +41,11 @@ type CompactOption struct {
 	CompactConfigPath   string
 }
 
-type SaveOption struct {
-	BootstrapPath string
-	DatabasePath  string
-}
-
 type GenerateOption struct {
-	DatabasePath string
+	BootstrapPaths         []string
+	DatabasePath           string
+	ChunkdictBootstrapPath string
+	OutputPath             string
 }
 
 type Builder struct {
@@ -153,31 +151,22 @@ func (builder *Builder) Run(option BuilderOption) error {
 	return builder.run(args, option.PrefetchPatterns)
 }
 
-// Save calls `nydus-image chunkdict save` to parse Nydus bootstrap
-func (builder *Builder) Save(option SaveOption) error {
-	args := []string{
-		"chunkdict",
-		"save",
-		"--log-level",
-		"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")
+	logrus.Infof("Invoking 'nydus-image chunkdict generate' command")
 	args := []string{
 		"chunkdict",
 		"generate",
 		"--log-level",
 		"warn",
+		"--bootstrap",
+		option.ChunkdictBootstrapPath,
 		"--database",
 		option.DatabasePath,
+		"--output-json",
+		option.OutputPath,
 	}
+	args = append(args, option.BootstrapPaths...)
+
 	return builder.run(args, "")
 }
diff --git a/contrib/nydusify/pkg/chunkdict/generator/generator.go b/contrib/nydusify/pkg/chunkdict/generator/generator.go
@@ -59,72 +59,66 @@ func New(opt Opt) (*Generator, error) {
 
 // Generate saves multiple Nydus bootstraps into the database one by one.
 func (generator *Generator) Generate(ctx context.Context) error {
-	for index := range generator.Sources {
-		if err := generator.save(ctx, index); err != nil {
-			if utils.RetryWithHTTP(err) {
+	var bootstrapPaths []string
+	bootstrapPaths, err := generator.pull(ctx)
+
+	if err != nil {
+		if utils.RetryWithHTTP(err) {
+			for index := range generator.Sources {
 				generator.sourcesParser[index].Remote.MaybeWithHTTP(err)
 			}
-			if err := generator.save(ctx, index); err != nil {
-				return err
-			}
+		}
+		bootstrapPaths, err = generator.pull(ctx)
+		if err != nil {
+			return err
 		}
 	}
-	if err := generator.deduplicating(ctx); err != nil {
+
+	if err := generator.generate(ctx, bootstrapPaths); err != nil {
 		return err
 	}
 	return nil
 }
 
-// "save" stores information of chunk and blob of a Nydus Image in the database
-func (generator *Generator) save(ctx context.Context, index int) error {
-	currentDir, _ := os.Getwd()
-	sourceParsed, err := generator.sourcesParser[index].Parse(ctx)
-	if err != nil {
-		return errors.Wrap(err, "parse Nydus image")
-	}
-
-	// Create a directory to store the image bootstrap
-	nydusImageName := strings.Replace(generator.Sources[index], "/", ":", -1)
-	bootstrapFolderPath := filepath.Join(currentDir, 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, bootstrapFolderPath, index); err != nil {
-		return errors.Wrap(err, "output image information")
-	}
-
-	databaseName := "chunkdict.db"
-	databaseType := "sqlite"
-	DatabasePath := databaseType + "://" + filepath.Join(currentDir, generator.WorkDir, databaseName)
+// Pull the bootstrap of nydus image
+func (generator *Generator) pull(ctx context.Context) ([]string, error) {
+	var bootstrapPaths []string
+	for index := range generator.Sources {
+		sourceParsed, err := generator.sourcesParser[index].Parse(ctx)
+		if err != nil {
+			return nil, errors.Wrap(err, "parse Nydus image")
+		}
 
-	// Invoke "nydus-image save" command
-	builder := build.NewBuilder(generator.NydusImagePath)
-	if err := builder.Save(build.SaveOption{
-		BootstrapPath: filepath.Join(bootstrapFolderPath, "nydus_bootstrap"),
-		DatabasePath:  DatabasePath,
-	}); err != nil {
-		return errors.Wrap(err, "invalid nydus bootstrap format")
+		// Create a directory to store the image bootstrap
+		nydusImageName := strings.Replace(generator.Sources[index], "/", ":", -1)
+		bootstrapDirPath := filepath.Join(generator.WorkDir, nydusImageName)
+		if err := os.MkdirAll(bootstrapDirPath, fs.ModePerm); err != nil {
+			return nil, errors.Wrap(err, "creat work directory")
+		}
+		if err := generator.Output(ctx, sourceParsed, bootstrapDirPath, index); err != nil {
+			return nil, errors.Wrap(err, "output image information")
+		}
+		bootstrapPath := filepath.Join(bootstrapDirPath, "nydus_bootstrap")
+		bootstrapPaths = append(bootstrapPaths, bootstrapPath)
 	}
-
-	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
+	return bootstrapPaths, nil
 }
 
-func (generator *Generator) deduplicating(ctx context.Context) error {
-	builder := build.NewBuilder(generator.NydusImagePath)
+func (generator *Generator) generate(ctx context.Context, bootstrapPaths []string) error {
+	// Invoke "nydus-image generate" command
 	currentDir, _ := os.Getwd()
-
-	databaseName := "chunkdict.db"
+	builder := build.NewBuilder(generator.NydusImagePath)
 	databaseType := "sqlite"
-	DatabasePath := databaseType + "://" + filepath.Join(currentDir, generator.WorkDir, databaseName)
 	if err := builder.Generate(build.GenerateOption{
-		DatabasePath: DatabasePath,
+		BootstrapPaths:         bootstrapPaths,
+		ChunkdictBootstrapPath: filepath.Join(generator.WorkDir, "chunkdict_bootstrap"),
+		DatabasePath:           databaseType + "://" + filepath.Join(currentDir, generator.WorkDir, "database.db"),
+		OutputPath:             filepath.Join(generator.WorkDir, "nydus_bootstrap_output.json"),
 	}); err != nil {
 		return errors.Wrap(err, "invalid nydus bootstrap format")
 	}
+
+	logrus.Infof("Successfully generate image chunk dictionary")
+
 	return nil
 }
diff --git a/smoke/tests/image_test.go b/smoke/tests/image_test.go
@@ -140,48 +140,50 @@ func (i *ImageTestSuite) TestConvertAndCopyImage(t *testing.T, ctx tool.Context,
 	tool.RunWithoutOutput(t, checkCmd)
 }
 
-func (i *ImageTestSuite) TestGenerate() test.Generator {
-
-	scenarios := tool.DescartesIterator{}
-	scenarios.Dimension(paramImage, []interface{}{"nginx:latest"})
-
-	return func() (name string, testCase test.Case) {
-		if !scenarios.HasNext() {
-			return
-		}
-		scenario := scenarios.Next()
-		ctx := tool.DefaultContext(i.T)
-
-		image := i.prepareImage(i.T, scenario.GetString(paramImage))
-		return scenario.Str(), func(t *testing.T) {
-			i.TestGenerateChunkDict(t, *ctx, image)
-		}
-	}
-}
-
-func (i *ImageTestSuite) TestGenerateChunkDict(t *testing.T, ctx tool.Context, source string) {
-
-	// Prepare work directory
-	ctx.PrepareWorkDir(t)
-	defer ctx.Destroy(t)
-
-	target := fmt.Sprintf("%s-nydus-%s", source, uuid.NewString())
-	logLevel := "--log-level warn"
-
-	// Convert image
-	convertCmd := fmt.Sprintf(
-		"%s %s convert --source %s --target %s --nydus-image %s --work-dir %s",
-		ctx.Binary.Nydusify, logLevel, source, target, ctx.Binary.Builder, ctx.Env.WorkDir,
-	)
-	tool.RunWithoutOutput(t, convertCmd)
-
-	nydusifyPath := ctx.Binary.Nydusify
-	generateCmd := fmt.Sprintf(
-		"%s %s chunkdict generate --sources %s --nydus-image %s --work-dir %s",
-		nydusifyPath, logLevel, target, ctx.Binary.Builder, ctx.Env.WorkDir,
-	)
-	tool.RunWithoutOutput(t, generateCmd)
-
+func (i *ImageTestSuite) TestGenerateChunkdict() test.Generator {
+    return func() (name string, testCase test.Case) {
+		imagename1 := "redis:7.0.1"
+		imagename2 := "redis:7.0.2"
+		imagename3 := "redis:7.0.3"
+		image1 := i.prepareImage(i.T, imagename1)
+		image2 := i.prepareImage(i.T, imagename2)
+		image3 := i.prepareImage(i.T, imagename3)
+        ctx := tool.DefaultContext(i.T)
+
+        // Prepare work directory
+        ctx.PrepareWorkDir(i.T)
+        defer ctx.Destroy(i.T)
+
+        logLevel := "--log-level warn"
+        nydusifyPath := ctx.Binary.Nydusify
+
+		target1 := fmt.Sprintf("%s-nydus-%s", image1, uuid.NewString())
+		target2 := fmt.Sprintf("%s-nydus-%s", image2, uuid.NewString())
+		target3 := fmt.Sprintf("%s-nydus-%s", image3, uuid.NewString())
+        convertCmd1 := fmt.Sprintf(
+			"%s %s convert --source %s --target %s --nydus-image %s --work-dir %s",
+			ctx.Binary.Nydusify, logLevel, image1, target1, ctx.Binary.Builder, ctx.Env.TempDir,
+		)
+		tool.RunWithoutOutput(i.T, convertCmd1)
+		convertCmd2 := fmt.Sprintf(
+			"%s %s convert --source %s --target %s --nydus-image %s --work-dir %s",
+			ctx.Binary.Nydusify, logLevel, image1, target2, ctx.Binary.Builder, ctx.Env.TempDir,
+		)
+		tool.RunWithoutOutput(i.T, convertCmd2)
+		convertCmd3 := fmt.Sprintf(
+			"%s %s convert --source %s --target %s --nydus-image %s --work-dir %s",
+			ctx.Binary.Nydusify, logLevel, image1, target3, ctx.Binary.Builder, ctx.Env.TempDir,
+		)
+		tool.RunWithoutOutput(i.T, convertCmd3)
+        target := fmt.Sprintf("%s,%s,%s", target1, target2, target3)
+
+		generateCmd := fmt.Sprintf(
+			"%s %s chunkdict generate --sources %s --nydus-image %s --work-dir %s",
+			nydusifyPath, logLevel, target, ctx.Binary.Builder, ctx.Env.TempDir,
+		)
+		tool.RunWithoutOutput(i.T, generateCmd)
+        return "generateChunkdict", nil
+    }
 }
 
 func (i *ImageTestSuite) prepareImage(t *testing.T, image string) string {

diff --git a/smoke/tests/tool/context.go b/smoke/tests/tool/context.go
@@ -42,6 +42,7 @@ type RuntimeContext struct {
 }
 
 type EnvContext struct {
+	TempDir       string
 	WorkDir       string
 	BlobDir       string
 	CacheDir      string
@@ -99,6 +100,7 @@ func (ctx *Context) PrepareWorkDir(t *testing.T) {
 	require.NoError(t, err)
 
 	ctx.Env = EnvContext{
+		TempDir:  tempDir,
 		WorkDir:  workDir,
 		BlobDir:  blobDir,
 		CacheDir: cacheDir,

diff --git a/src/bin/nydus-image/deduplicate.rs b/src/bin/nydus-image/deduplicate.rs
@@ -229,9 +229,9 @@ pub struct Algorithm<D: Database + Send + Sync> {
 }
 
 // Generate deduplicated chunkdict by exponential_smoothing algorithm
-type Versiondic = HashMap<String, Vec<ChunkdictChunkInfo>>;
+type VersionMap = HashMap<String, Vec<ChunkdictChunkInfo>>;
 // Generate deduplicated chunkdict by cluster algorithm
-type Imagedic = Vec<HashMap<Vec<String>, Vec<ChunkdictChunkInfo>>>;
+type ImageMap = Vec<HashMap<Vec<String>, Vec<ChunkdictChunkInfo>>>;
 
 impl Algorithm<SqliteDatabase> {
     pub fn new(algorithm: String, db_url: &str) -> anyhow::Result<Self> {
@@ -679,7 +679,7 @@ impl Algorithm<SqliteDatabase> {
 
     pub fn deduplicate_version(
         all_chunks: &[ChunkdictChunkInfo],
-    ) -> anyhow::Result<(Versiondic, Imagedic)> {
+    ) -> anyhow::Result<(VersionMap, ImageMap)> {
         let mut all_chunks_size = 0;
         for i in all_chunks {
             all_chunks_size += i.chunk_compressed_size;