Optimize file copying and stage saving between stages.

This change calculates the exact files and directories needed between
stages used in the COPY command. Instead of saving the entire
stage as a tarball, we now save only the necessary files.

pkg/commands/add.go

@@ -47,7 +47,7 @@ type AddCommand struct {
 func (a *AddCommand) ExecuteCommand(config *v1.Config, buildArgs *dockerfile.BuildArgs) error {
 	replacementEnvs := buildArgs.ReplacementEnvs(config.Env)
 
-	srcs, dest, err := resolveEnvAndWildcards(a.cmd.SourcesAndDest, a.buildcontext, replacementEnvs)
+	srcs, dest, err := util.ResolveEnvAndWildcards(a.cmd.SourcesAndDest, a.buildcontext, replacementEnvs)
 	if err != nil {
 		return err
 	}
@@ -114,7 +114,7 @@ func (a *AddCommand) String() string {
 func (a *AddCommand) FilesUsedFromContext(config *v1.Config, buildArgs *dockerfile.BuildArgs) ([]string, error) {
 	replacementEnvs := buildArgs.ReplacementEnvs(config.Env)
 
-	srcs, _, err := resolveEnvAndWildcards(a.cmd.SourcesAndDest, a.buildcontext, replacementEnvs)
+	srcs, _, err := util.ResolveEnvAndWildcards(a.cmd.SourcesAndDest, a.buildcontext, replacementEnvs)
 	if err != nil {
 		return nil, err
 	}

pkg/commands/copy.go

@@ -45,7 +45,7 @@ func (c *CopyCommand) ExecuteCommand(config *v1.Config, buildArgs *dockerfile.Bu
 
 	replacementEnvs := buildArgs.ReplacementEnvs(config.Env)
 
-	srcs, dest, err := resolveEnvAndWildcards(c.cmd.SourcesAndDest, c.buildcontext, replacementEnvs)
+	srcs, dest, err := util.ResolveEnvAndWildcards(c.cmd.SourcesAndDest, c.buildcontext, replacementEnvs)
 	if err != nil {
 		return err
 	}
@@ -100,18 +100,6 @@ func (c *CopyCommand) ExecuteCommand(config *v1.Config, buildArgs *dockerfile.Bu
 	return nil
 }
 
-func resolveEnvAndWildcards(sd instructions.SourcesAndDest, buildcontext string, envs []string) ([]string, string, error) {
-	// First, resolve any environment replacement
-	resolvedEnvs, err := util.ResolveEnvironmentReplacementList(sd, envs, true)
-	if err != nil {
-		return nil, "", err
-	}
-	dest := resolvedEnvs[len(resolvedEnvs)-1]
-	// Resolve wildcards and get a list of resolved sources
-	srcs, err := util.ResolveSources(resolvedEnvs, buildcontext)
-	return srcs, dest, err
-}
-
 // FilesToSnapshot should return an empty array if still nil; no files were changed
 func (c *CopyCommand) FilesToSnapshot() []string {
 	return c.snapshotFiles
@@ -129,7 +117,7 @@ func (c *CopyCommand) FilesUsedFromContext(config *v1.Config, buildArgs *dockerf
 	}
 
 	replacementEnvs := buildArgs.ReplacementEnvs(config.Env)
-	srcs, _, err := resolveEnvAndWildcards(c.cmd.SourcesAndDest, c.buildcontext, replacementEnvs)
+	srcs, _, err := util.ResolveEnvAndWildcards(c.cmd.SourcesAndDest, c.buildcontext, replacementEnvs)
 	if err != nil {
 		return nil, err
 	}

pkg/config/stage.go

@@ -26,4 +26,5 @@ type KanikoStage struct {
 	BaseImageStoredLocally bool
 	SaveStage              bool
 	MetaArgs               []instructions.ArgCommand
+	Index                  int
 }

pkg/dockerfile/dockerfile.go

@@ -25,6 +25,8 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/sirupsen/logrus"
+
 	"github.com/GoogleContainerTools/kaniko/pkg/config"
 	"github.com/GoogleContainerTools/kaniko/pkg/util"
 	"github.com/moby/buildkit/frontend/dockerfile/instructions"
@@ -67,13 +69,15 @@ func Stages(opts *config.KanikoOptions) ([]config.KanikoStage, error) {
 			return nil, errors.Wrap(err, "resolving base name")
 		}
 		stage.Name = resolvedBaseName
+		logrus.Infof("Resolved base name %s to %s", stage.BaseName, stage.Name)
 		kanikoStages = append(kanikoStages, config.KanikoStage{
 			Stage:                  stage,
 			BaseImageIndex:         baseImageIndex(index, stages),
 			BaseImageStoredLocally: (baseImageIndex(index, stages) != -1),
 			SaveStage:              saveStage(index, stages),
 			Final:                  index == targetStage,
 			MetaArgs:               metaArgs,
+			Index:                  index,
 		})
 		if index == targetStage {
 			break
@@ -175,14 +179,6 @@ func saveStage(index int, stages []instructions.Stage) bool {
 				return true
 			}
 		}
-		for _, cmd := range stage.Commands {
-			switch c := cmd.(type) {
-			case *instructions.CopyCommand:
-				if c.From == strconv.Itoa(index) {
-					return true
-				}
-			}
-		}
 	}
 	return false
 }

pkg/dockerfile/dockerfile_test.go

@@ -114,7 +114,7 @@ func Test_SaveStage(t *testing.T) {
 		{
 			name:     "reference stage in later copy command",
 			index:    0,
-			expected: true,
+			expected: false,
 		},
 		{
 			name:     "reference stage in later from command",

pkg/executor/build.go

@@ -23,6 +23,8 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/otiai10/copy"
+
 	"github.com/google/go-containerregistry/pkg/v1/partial"
 
 	"github.com/moby/buildkit/frontend/dockerfile/instructions"
@@ -60,10 +62,11 @@ type stageBuilder struct {
 	opts            *config.KanikoOptions
 	cmds            []commands.DockerCommand
 	args            *dockerfile.BuildArgs
+	crossStageDeps  map[int][]string
 }
 
 // newStageBuilder returns a new type stageBuilder which contains all the information required to build the stage
-func newStageBuilder(opts *config.KanikoOptions, stage config.KanikoStage) (*stageBuilder, error) {
+func newStageBuilder(opts *config.KanikoOptions, stage config.KanikoStage, crossStageDeps map[int][]string) (*stageBuilder, error) {
 	sourceImage, err := util.RetrieveSourceImage(stage, opts)
 	if err != nil {
 		return nil, err
@@ -96,6 +99,7 @@ func newStageBuilder(opts *config.KanikoOptions, stage config.KanikoStage) (*sta
 		snapshotter:     snapshotter,
 		baseImageDigest: digest.String(),
 		opts:            opts,
+		crossStageDeps:  crossStageDeps,
 	}
 
 	for _, cmd := range s.stage.Commands {
@@ -207,6 +211,10 @@ func (s *stageBuilder) build() error {
 			break
 		}
 	}
+	if len(s.crossStageDeps[s.stage.Index]) > 0 {
+		shouldUnpack = true
+	}
+
 	if shouldUnpack {
 		t := timing.Start("FS Unpacking")
 		if _, err := util.GetFSFromImage(constants.RootDir, s.image); err != nil {
@@ -353,6 +361,63 @@ func (s *stageBuilder) saveSnapshotToImage(createdBy string, tarPath string) err
 
 }
 
+func CalculateDependencies(opts *config.KanikoOptions) (map[int][]string, error) {
+	stages, err := dockerfile.Stages(opts)
+	if err != nil {
+		return nil, err
+	}
+	images := []v1.Image{}
+	depGraph := map[int][]string{}
+	for _, s := range stages {
+		ba := dockerfile.NewBuildArgs(opts.BuildArgs)
+		ba.AddMetaArgs(s.MetaArgs)
+		var image v1.Image
+		var err error
+		if s.BaseImageStoredLocally {
+			image = images[s.BaseImageIndex]
+		} else if s.Name == constants.NoBaseImage {
+			image = empty.Image
+		} else {
+			image, err = util.RetrieveSourceImage(s, opts)
+			if err != nil {
+				return nil, err
+			}
+		}
+		initializeConfig(image)
+		cfg, err := image.ConfigFile()
+		if err != nil {
+			return nil, err
+		}
+		for _, c := range s.Commands {
+			switch cmd := c.(type) {
+			case *instructions.CopyCommand:
+				if cmd.From != "" {
+					i, err := strconv.Atoi(cmd.From)
+					if err != nil {
+						continue
+					}
+					resolved, err := util.ResolveEnvironmentReplacementList(cmd.SourcesAndDest, cfg.Config.Env, true)
+					if err != nil {
+						return nil, err
+					}
+
+					depGraph[i] = append(depGraph[i], resolved[0:len(resolved)-1]...)
+				}
+			case *instructions.EnvCommand:
+				if err := util.UpdateConfigEnv(cmd.Env, &cfg.Config, ba.ReplacementEnvs(cfg.Config.Env)); err != nil {
+					return nil, err
+				}
+				image, err = mutate.Config(image, cfg.Config)
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+		images = append(images, image)
+	}
+	return depGraph, nil
+}
+
 // DoBuild executes building the Dockerfile
 func DoBuild(opts *config.KanikoOptions) (v1.Image, error) {
 	t := timing.Start("Total Build Time")
@@ -369,8 +434,14 @@ func DoBuild(opts *config.KanikoOptions) (v1.Image, error) {
 		return nil, err
 	}
 
+	crossStageDependencies, err := CalculateDependencies(opts)
+	if err != nil {
+		return nil, err
+	}
+	logrus.Infof("Built cross stage deps: %v", crossStageDependencies)
+
 	for index, stage := range stages {
-		sb, err := newStageBuilder(opts, stage)
+		sb, err := newStageBuilder(opts, stage, crossStageDependencies)
 		if err != nil {
 			return nil, err
 		}
@@ -405,10 +476,21 @@ func DoBuild(opts *config.KanikoOptions) (v1.Image, error) {
 			if err := saveStageAsTarball(strconv.Itoa(index), sourceImage); err != nil {
 				return nil, err
 			}
-			if err := extractImageToDependecyDir(strconv.Itoa(index), sourceImage); err != nil {
-				return nil, err
-			}
 		}
+
+		filesToSave, err := filesToSave(crossStageDependencies[index])
+		if err != nil {
+			return nil, err
+		}
+		dstDir := filepath.Join(constants.KanikoDir, strconv.Itoa(index))
+		if err := os.MkdirAll(dstDir, 0644); err != nil {
+			return nil, err
+		}
+		for _, p := range filesToSave {
+			logrus.Infof("Saving file %s for later use.", p)
+			copy.Copy(p, filepath.Join(dstDir, p))
+		}
+
 		// Delete the filesystem
 		if err := util.DeleteFilesystem(); err != nil {
 			return nil, err
@@ -418,6 +500,18 @@ func DoBuild(opts *config.KanikoOptions) (v1.Image, error) {
 	return nil, err
 }
 
+func filesToSave(deps []string) ([]string, error) {
+	allFiles := []string{}
+	for _, src := range deps {
+		srcs, err := filepath.Glob(src)
+		if err != nil {
+			return nil, err
+		}
+		allFiles = append(allFiles, srcs...)
+	}
+	return allFiles, nil
+}
+
 func fetchExtraStages(stages []config.KanikoStage, opts *config.KanikoOptions) error {
 	t := timing.Start("Fetching Extra Stages")
 	defer timing.DefaultRun.Stop(t)