Repository: uber/makisu
Pull request #279 (2019-10-13 03:14:40)

Correctly cache "COPY --from" (2nd attempt)
ref #72

Sharma-Rajat: This works. I have tested it with all my failing images. :) 
yiranwang52: > This works. I have tested it with all my failing images. :)

nice! thanks!

Patch:
diff --git a/lib/builder/build_plan.go b/lib/builder/build_plan.go
index 3fd2e19..759e832 100644
--- a/lib/builder/build_plan.go
+++ b/lib/builder/build_plan.go
@@ -16,6 +16,7 @@ package builder
 
 import (
 	"fmt"
+	"hash/crc32"
 	"strconv"
 
 	"github.com/uber/makisu/lib/cache"
@@ -23,6 +24,7 @@ import (
 	"github.com/uber/makisu/lib/docker/image"
 	"github.com/uber/makisu/lib/log"
 	"github.com/uber/makisu/lib/parser/dockerfile"
+	"github.com/uber/makisu/lib/utils"
 	"github.com/uber/makisu/lib/utils/stringset"
 )
 
@@ -39,14 +41,20 @@ type BuildPlan struct {
 	target       image.Name
 	replicas     []image.Name
 	cacheMgr     cache.Manager
-	stages       []*buildStage
 
-	// Aliases of stages. For stages missing `AS`, an index number will be used
-	// as alias.
-	// Note: this doesn't include those `COPY --from=<alias>` if that alias is
-	// name of another image.
+	// stages contains the build stages defined in dockerfile.
+	stages []*buildStage
+
+	// TODO: this is not used for now.
+	// Aliases of stages.
 	stageAliases map[string]struct{}
 
+	// TODO: this is not used for now.
+	// Index aliases of stages.
+	// This extra index is needed because shadow stages could be inserted into
+	// stages list to support `COPY --from=<image>`.
+	stageIndexAliases map[string]*buildStage
+
 	opts *buildPlanOptions
 }
 
@@ -56,80 +64,99 @@ func NewBuildPlan(
 	ctx *context.BuildContext, target image.Name, replicas []image.Name, cacheMgr cache.Manager,
 	parsedStages []*dockerfile.Stage, allowModifyFS, forceCommit bool) (*BuildPlan, error) {
 
-	aliases, err := buildAliases(parsedStages)
-	if err != nil {
-		return nil, fmt.Errorf("build alias list: %s", err)
-	}
-
 	plan := &BuildPlan{
-		baseCtx:      ctx,
-		copyFromDirs: make(map[string][]string),
-		target:       target,
-		replicas:     replicas,
-		cacheMgr:     cacheMgr,
-		stages:       make([]*buildStage, len(parsedStages)),
-		stageAliases: aliases,
+		baseCtx:           ctx,
+		copyFromDirs:      make(map[string][]string),
+		target:            target,
+		replicas:          replicas,
+		cacheMgr:          cacheMgr,
+		stages:            make([]*buildStage, 0),
+		stageAliases:      make(map[string]struct{}),
+		stageIndexAliases: make(map[string]*buildStage),
 		opts: &buildPlanOptions{
 			forceCommit:   forceCommit,
 			allowModifyFS: allowModifyFS,
 		},
 	}
 
+	if err := plan.processStagesAndAliases(ctx, parsedStages); err != nil {
+		return nil, fmt.Errorf("process stages and aliases: %s", err)
+	}
+
+	return plan, nil
+}
+
+func (plan *BuildPlan) processStagesAndAliases(
+	ctx *context.BuildContext, parsedStages dockerfile.Stages) error {
+
+	checksum := crc32.ChecksumIEEE([]byte(utils.BuildHash + fmt.Sprintf("%v", plan.opts)))
+	seedCacheID := fmt.Sprintf("%x", checksum)
+
+	existingAliases := make(map[string]struct{})
 	for i, parsedStage := range parsedStages {
+		// Record alias.
+		if parsedStage.From.Alias != "" {
+			if _, ok := existingAliases[parsedStage.From.Alias]; ok {
+				return fmt.Errorf("duplicate stage alias: %s", parsedStage.From.Alias)
+			} else if _, err := strconv.Atoi(parsedStage.From.Alias); err == nil {
+				// Note: Docker would return `name can't start with a number or
+				// contain symbols`.
+				return fmt.Errorf("stage alias cannot be a number: %s", parsedStage.From.Alias)
+			}
+		} else {
+			parsedStage.From.Alias = strconv.Itoa(i)
+		}
+		existingAliases[parsedStage.From.Alias] = struct{}{}
+
 		// Add this stage to the plan.
 		stage, err := newBuildStage(
-			ctx, parsedStage.From.Alias, parsedStage, plan.opts)
+			ctx, parsedStage.From.Alias, seedCacheID, parsedStage, plan.opts)
 		if err != nil {
-			return nil, fmt.Errorf("failed to convert parsed stage: %s", err)
+			return fmt.Errorf("failed to convert parsed stage: %s", err)
 		}
 
 		if len(stage.copyFromDirs) > 0 && !plan.opts.allowModifyFS {
 			// TODO(pourchet): Support this at some point.
 			// TODO: have a centralized place for these validations.
-			return nil, fmt.Errorf("must allow modifyfs for multi-stage dockerfiles with COPY --from")
+			return fmt.Errorf("must allow modifyfs for multi-stage dockerfiles with COPY --from")
 		}
-		plan.stages[i] = stage
-	}
-
-	if err := plan.handleCopyFromDirs(); err != nil {
-		return nil, fmt.Errorf("handle cross refs: %s", err)
-	}
 
-	return plan, nil
-}
-
-// handleCopyFromDirs goes through all of the stages in the build plan and looks
-// at the `COPY --from` steps to make sure they are valid.
-func (plan *BuildPlan) handleCopyFromDirs() error {
-	for _, stage := range plan.stages {
+		// Goes through all of the stages in the build plan and looks
+		// at the `COPY --from` steps to make sure they are valid.
 		for alias, dirs := range stage.copyFromDirs {
+			// Populate copyFromDirs.
 			plan.copyFromDirs[alias] = stringset.FromSlice(
 				append(plan.copyFromDirs[alias], dirs...),
 			).ToSlice()
-		}
-	}
-	return nil
-}
 
-// buildAliases mutates the list of stages to assign default aliases.
-// Default aliases will be integers starting from 0.
-func buildAliases(stages dockerfile.Stages) (map[string]struct{}, error) {
-	aliases := make(map[string]struct{})
-	for i, parsedStage := range stages {
-		// Check for stage alias collision if alias isn't empty.
-		if parsedStage.From.Alias != "" {
-			if _, ok := aliases[parsedStage.From.Alias]; ok {
-				return nil, fmt.Errorf("duplicate stage alias: %s", parsedStage.From.Alias)
-			} else if _, err := strconv.Atoi(parsedStage.From.Alias); err == nil {
-				// Docker would return `name can't start with a number or contain symbols`.
-				return nil, fmt.Errorf("stage alias cannot be a number: %s", parsedStage.From.Alias)
+			if _, ok := existingAliases[alias]; !ok {
+				// If the alias was an image name and not already handled,
+				// prepend a fake stage with the alias to download that image.
+				if name, err := image.ParseNameForPull(alias); err != nil || !name.IsValid() {
+					return fmt.Errorf("copy from nonexistent stage %s", alias)
+				}
+				remoteImageStage, err := newRemoteImageStage(
+					plan.baseCtx, alias, seedCacheID, plan.opts)
+				if err != nil {
+					return fmt.Errorf("new image stage: %s", err)
+				}
+
+				// Append to stage list and update cache id.
+				plan.stages = append(plan.stages, remoteImageStage)
+				// TODO: instead of chaining cache ID, it's better to calculate
+				// from scratch before executing a stage.
+				seedCacheID = remoteImageStage.nodes[len(remoteImageStage.nodes)-1].CacheID()
 			}
-		} else {
-			parsedStage.From.Alias = strconv.Itoa(i)
 		}
-		aliases[parsedStage.From.Alias] = struct{}{}
+
+		// Append to stage list and update cache id.
+		plan.stages = append(plan.stages, stage)
+		// TODO: instead of chaining cache ID, it's better to calculate from
+		// scratch before executing a stage.
+		seedCacheID = stage.nodes[len(stage.nodes)-1].CacheID()
 	}
-	return aliases, nil
+
+	return nil
 }
 
 // Execute executes all build stages in order.
@@ -137,6 +164,9 @@ func (plan *BuildPlan) Execute() (*image.DistributionManifest, error) {
 	var currStage *buildStage
 	for k := 0; k < len(plan.stages); k++ {
 		currStage = plan.stages[k]
+
+		// TODO: Implicit stages from "COPY --from=<image>" might introduce
+		// confusion here. Print stageIndexAliases instead.
 		log.Infof("* Stage %d/%d : %s", k+1, len(plan.stages), currStage.String())
 
 		// Try to pull reusable layers cached from previous builds.
@@ -179,57 +209,22 @@ func (plan *BuildPlan) Execute() (*image.DistributionManifest, error) {
 }
 
 func (plan *BuildPlan) executeStage(stage *buildStage, lastStage, copiedFrom bool) error {
-	if !plan.opts.allowModifyFS {
+	if err := stage.build(plan.cacheMgr, lastStage, copiedFrom); err != nil {
+		return fmt.Errorf("build stage %s: %s", stage.alias, err)
+	}
+
+	if plan.opts.allowModifyFS {
 		// Note: The rest of this function mostly deal with `COPY --from`
 		// related logic, and currently `COPY --from` cannot be supported with
 		// modifyfs=false. That combination was rejected in NewPlan().
-		if err := stage.build(plan.cacheMgr, lastStage, copiedFrom); err != nil {
-			return fmt.Errorf("build stage %s: %s", stage.alias, err)
+		if err := stage.checkpoint(plan.copyFromDirs[stage.alias]); err != nil {
+			return fmt.Errorf("checkpoint stage %s: %s", stage.alias, err)
 		}
-		return nil
-	}
 
-	// Handle `COPY --from=<alias>` where alias is not a stage but an image.
-	// Create and execute a fake stage with only FROM.
-	// TODO: This should be done at step level.
-	// TODO: This step prints `* Step 1/1 ...` which is misleading.
-	for alias, _ := range stage.copyFromDirs {
-		if _, ok := plan.stageAliases[alias]; !ok {
-			name, err := image.ParseNameForPull(alias)
-			if err != nil || !name.IsValid() {
-				return fmt.Errorf("copy from nonexistent stage %s", alias)
-			}
-			remoteImageStage, err := newRemoteImageStage(
-				plan.baseCtx, alias, plan.opts)
-			if err != nil {
-				return fmt.Errorf("new image stage: %s", err)
-			}
-
-			if err := remoteImageStage.build(plan.cacheMgr, false, true); err != nil {
-				return fmt.Errorf("build stage %s: %s", stage.alias, err)
-			}
-
-			if err := remoteImageStage.checkpoint(plan.copyFromDirs[alias]); err != nil {
-				return fmt.Errorf("checkpoint stage %s: %s", stage.alias, err)
-			}
-
-			if err := remoteImageStage.cleanup(); err != nil {
-				return fmt.Errorf("cleanup stage %s: %s", stage.alias, err)
-			}
+		if err := stage.cleanup(); err != nil {
+			return fmt.Errorf("cleanup stage %s: %s", stage.alias, err)
 		}
 	}
 
-	if err := stage.build(plan.cacheMgr, lastStage, copiedFrom); err != nil {
-		return fmt.Errorf("build stage %s: %s", stage.alias, err)
-	}
-
-	if err := stage.checkpoint(plan.copyFromDirs[stage.alias]); err != nil {
-		return fmt.Errorf("checkpoint stage %s: %s", stage.alias, err)
-	}
-
-	if err := stage.cleanup(); err != nil {
-		return fmt.Errorf("cleanup stage %s: %s", stage.alias, err)
-	}
-
 	return nil
 }
diff --git a/lib/builder/build_stage.go b/lib/builder/build_stage.go
index d51e014..440f9db 100644
--- a/lib/builder/build_stage.go
+++ b/lib/builder/build_stage.go
@@ -19,7 +19,6 @@ import (
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
-	"hash/crc32"
 	"io/ioutil"
 	"os"
 	"path"
@@ -32,7 +31,6 @@ import (
 	"github.com/uber/makisu/lib/log"
 	"github.com/uber/makisu/lib/parser/dockerfile"
 	"github.com/uber/makisu/lib/storage"
-	"github.com/uber/makisu/lib/utils"
 )
 
 type buildStageOptions struct {
@@ -57,7 +55,7 @@ type buildStage struct {
 
 // newBuildStage initializes a buildStage.
 func newBuildStage(
-	baseCtx *context.BuildContext, alias string, parsedStage *dockerfile.Stage,
+	baseCtx *context.BuildContext, alias, seed string, parsedStage *dockerfile.Stage,
 	planOpts *buildPlanOptions) (*buildStage, error) {
 
 	// Create a new build context for the stage.
@@ -68,7 +66,7 @@ func newBuildStage(
 	}
 
 	// Create steps from parsed stage.
-	steps, err := createDockerfileSteps(ctx, parsedStage, planOpts)
+	steps, err := createDockerfileSteps(ctx, seed, parsedStage, planOpts)
 	if err != nil {
 		return nil, fmt.Errorf("new dockerfile steps: %s", err)
 	}
@@ -76,9 +74,10 @@ func newBuildStage(
 	return newBuildStageHelper(ctx, alias, steps, planOpts)
 }
 
-// newRemoteImageStage initializes a buildStage.
+// newRemoteImageStage initializes a buildStage used for `COPY --from=<image>`.
 func newRemoteImageStage(
-	baseCtx *context.BuildContext, alias string, planOpts *buildPlanOptions) (*buildStage, error) {
+	baseCtx *context.BuildContext, alias, seed string,
+	planOpts *buildPlanOptions) (*buildStage, error) {
 
 	// Create a new build context for the stage.
 	ctx, err := context.NewBuildContext(
@@ -92,8 +91,6 @@ func newRemoteImageStage(
 	if err != nil {
 		return nil, fmt.Errorf("new from step: %s", err)
 	}
-	checksum := crc32.ChecksumIEEE([]byte(utils.BuildHash + fmt.Sprintf("%v", *planOpts)))
-	seed := fmt.Sprintf("%x", checksum)
 	if err := from.SetCacheID(ctx, seed); err != nil {
 		return nil, fmt.Errorf("set cache id: %s", err)
 	}
@@ -153,11 +150,9 @@ func newBuildStageHelper(
 
 // createDockerfileSteps returns a list of build steps given a parsed stage.
 func createDockerfileSteps(
-	ctx *context.BuildContext, stage *dockerfile.Stage,
+	ctx *context.BuildContext, seed string, stage *dockerfile.Stage,
 	planOpts *buildPlanOptions) ([]step.BuildStep, error) {
 
-	checksum := crc32.ChecksumIEEE([]byte(utils.BuildHash + fmt.Sprintf("%v", *planOpts)))
-	seed := fmt.Sprintf("%x", checksum)
 	directives := append([]dockerfile.Directive{stage.From}, stage.Directives...)
 	var steps []step.BuildStep
 	for _, directive := range directives {
@@ -305,11 +300,13 @@ func (stage *buildStage) pullCacheLayers(cacheMgr cache.Manager) {
 	// Skip the first node since it's a FROM step. We do not want to try to pull
 	// from cache because the step itself will pull the right layers when it
 	// gets executed.
-	for _, node := range stage.nodes[1:] {
-		// Stop once the cache chain is broken.
-		if node.HasCommit() || stage.opts.forceCommit {
-			if !node.pullCacheLayer(cacheMgr) {
-				return
+	if len(stage.nodes) > 1 {
+		for _, node := range stage.nodes[1:] {
+			// Stop once the cache chain is broken.
+			if node.HasCommit() || stage.opts.forceCommit {
+				if !node.pullCacheLayer(cacheMgr) {
+					return
+				}
 			}
 		}
 	}
@@ -317,13 +314,17 @@ func (stage *buildStage) pullCacheLayers(cacheMgr cache.Manager) {
 
 func (stage *buildStage) latestFetched() int {
 	latest := -1
-	for i, node := range stage.nodes[1:] {
-		// Stop once the cache chain is broken.
-		if node.HasCommit() {
-			if len(node.digestPairs) != 0 {
-				latest = i + 1
-			} else {
-				return latest
+
+	if len(stage.nodes) > 1 {
+		// Skip FROM.
+		for i, node := range stage.nodes[1:] {
+			// Stop once the cache chain is broken.
+			if node.HasCommit() {
+				if len(node.digestPairs) != 0 {
+					latest = i + 1
+				} else {
+					return latest
+				}
 			}
 		}
 	}
diff --git a/lib/builder/step/add_copy_step.go b/lib/builder/step/add_copy_step.go
index b523605..a53c94c 100644
--- a/lib/builder/step/add_copy_step.go
+++ b/lib/builder/step/add_copy_step.go
@@ -18,11 +18,9 @@ import (
 	"fmt"
 	"hash/crc32"
 	"io"
-	"math/rand"
 	"os"
 	"path/filepath"
 	"strings"
-	"time"
 
 	"github.com/uber/makisu/lib/context"
 	"github.com/uber/makisu/lib/pathutils"
@@ -91,33 +89,28 @@ func (s *addCopyStep) ContextDirs() (string, []string) {
 }
 
 // SetCacheID sets the cache ID of the step given a seed SHA256 value.
-// Calculates the ID randomly if copying from another stage, else checksums the
-// file contents.
+// Calculates the ID based on content of files. If the previous steps, current
+// step args and the contents of sources are identical, cache ID should also be
+// identical.
 func (s *addCopyStep) SetCacheID(ctx *context.BuildContext, seed string) error {
+	// Initialize the checksum with the seed, directive and args.
+	checksum := crc32.NewIEEE()
+	_, err := checksum.Write([]byte(seed + string(s.directive) + s.args))
+	if err != nil {
+		return fmt.Errorf("hash copy directive: %s", err)
+	}
 	if s.fromStage != "" {
-		// It is copying from a previous stage, return random bytes.
-		// TODO: support cache here by chaining stage sha.
-		b := make([]byte, 4)
-		r := rand.New(rand.NewSource(time.Now().UnixNano()))
-		if _, err := r.Read(b); err != nil {
-			return fmt.Errorf("read rand: %s", err)
-		}
-		s.cacheID = fmt.Sprintf("%x", b)
+		// It is copying from a previous stage, rely on the fact that cache IDs
+		// are chained between stages.
+		// TODO: Properly calculate cache ID based on content of files.
 	} else {
-		// Initialize the checksum with the seed, directive and args.
-		checksum := crc32.NewIEEE()
-		_, err := checksum.Write([]byte(seed + string(s.directive) + s.args))
-		if err != nil {
-			return fmt.Errorf("hash copy directive: %s", err)
-		}
-
-		// If the step args and the contents of sources are identical, we should
-		// be able to use the cache from the previous build.
-		if err := s.updateContextChecksum(ctx, checksum); err != nil {
+		// Update checksum based on content of files to be copied.
+		if err := s.calculateContextChecksum(ctx, checksum); err != nil {
 			return fmt.Errorf("hash context sources: %s", err)
 		}
-		s.cacheID = fmt.Sprintf("%x", checksum.Sum32())
 	}
+	s.cacheID = fmt.Sprintf("%x", checksum.Sum32())
+
 	return nil
 }
 
@@ -149,9 +142,8 @@ func (s *addCopyStep) Execute(ctx *context.BuildContext, modifyFS bool) (err err
 	return nil
 }
 
-// Updates the checksum passed in with the data stored in the context on the
-// filesystem.
-func (s *addCopyStep) updateContextChecksum(ctx *context.BuildContext, checksum io.Writer) error {
+// Updates the checksum passed in based on the content of files to be copied in.
+func (s *addCopyStep) calculateContextChecksum(ctx *context.BuildContext, checksum io.Writer) error {
 	if s.fromStage != "" {
 		return fmt.Errorf("not supported: the copy step has from stage flag")
 	}
@@ -191,6 +183,7 @@ func (s *addCopyStep) contextRootDir(ctx *context.BuildContext) string {
 	return ctx.ContextDir
 }
 
+// TODO: Consider file metadata?
 func checksumPathContents(
 	ctx *context.BuildContext, path string, fi os.FileInfo, checksum io.Writer) error {
 


Test patch:
diff --git a/lib/builder/build_stage_test.go b/lib/builder/build_stage_test.go
index bf988f7..854acab 100644
--- a/lib/builder/build_stage_test.go
+++ b/lib/builder/build_stage_test.go
@@ -123,7 +123,7 @@ func TestPullCacheLayers(t *testing.T) {
 				allowModifyFS: false,
 			}
 
-			stage, err := newBuildStage(ctx, alias, tc.stage, opts)
+			stage, err := newBuildStage(ctx, alias, tc.name, tc.stage, opts)
 			require.NoError(err)
 
 			kvStore := keyvalue.MockStore{}
diff --git a/lib/builder/step/copy_step_test.go b/lib/builder/step/copy_step_test.go
index 6ba7a93..cb609eb 100644
--- a/lib/builder/step/copy_step_test.go
+++ b/lib/builder/step/copy_step_test.go
@@ -154,16 +154,17 @@ func TestCopyStepSetCacheID(t *testing.T) {
 		defer cleanup()
 
 		step := CopyStepFixture("", "stage", []string{"."}, "tmp", false)
-		err := step.SetCacheID(context, "")
+		err := step.SetCacheID(context, "seed")
 		hash1 := step.CacheID()
 		require.NoError(err)
 
-		err = step.SetCacheID(context, hash1)
+		err = step.SetCacheID(context, "seed")
 		hash2 := step.CacheID()
 		require.NoError(err)
 
-		// hash1 and hash2 should be generated randomly.
-		require.NotEqual(hash1, hash2)
+		// It used to be the case that cache ID will be randomly generated for
+		// `COPY --from=<>`, but now it should be the same.
+		require.Equal(hash1, hash2)
 	})
 }