Move in-mem step cache logic from BuildStage to CacheManager (#260)

* comments

* comments

* Add in-mem kv store

* Add test

Makefile

@@ -58,8 +58,8 @@ cbins:
 
 $(ALL_SRC): ;
 
-# TODO(pourchet): Remove this hack to make dep more reliable. For some reason `go mod download` fails
-# sometimes on TravisCI, so run it a few times if it fails.
+# TODO(pourchet): Remove this hack to make dep more reliable. For some reason `go mod download`
+# fails sometimes on TravisCI, so run it a few times if it fails.
 vendor: go.mod go.sum
 	go mod download || go mod download || go mod download
 	go mod vendor
@@ -75,6 +75,8 @@ mocks: ext-tools
 	@echo "Generating mocks"
 	mkdir -p mocks/net/http
 	$(EXT_TOOLS_DIR)/mockgen -destination=mocks/net/http/mockhttp.go -package=mockhttp net/http RoundTripper
+	mkdir -p mocks/lib/registry
+	$(EXT_TOOLS_DIR)/mockgen -destination=mocks/lib/registry/mockclient.go -package=mockregistry github.com/uber/makisu/lib/registry Client
 
 env: test/python/requirements.txt
 	[ -d env ] || virtualenv --setuptools env

bin/makisu/cmd/build.go

@@ -72,9 +72,9 @@ type buildCmd struct {
 func getBuildCmd() *buildCmd {
 	buildCmd := &buildCmd{
 		Command: &cobra.Command{
-			Use:                   "build -t=<image_tag> [flags] <context_path>",
+			Use: "build -t=<image_tag> [flags] <context_path>",
 			DisableFlagsInUseLine: true,
-			Short:                 "Build docker image, optionally push to registries and/or load into docker daemon",
+			Short: "Build docker image, optionally push to registries and/or load into docker daemon",
 		},
 	}
 	buildCmd.Args = func(cmd *cobra.Command, args []string) error {
@@ -245,7 +245,6 @@ func (cmd *buildCmd) Build(contextDir string) error {
 			}
 			defer rootPreserver.RestoreRoot()
 		}
-		log.Debugf("build.Cmd.Build() first call")
 		buildContext.MemFS.Remove()
 		defer buildContext.MemFS.Remove()
 	}

lib/builder/build_node.go

@@ -166,6 +166,7 @@ func (n *buildNode) pushCacheLayer(cacheMgr cache.Manager) error {
 func (n *buildNode) pullCacheLayer(cacheMgr cache.Manager) bool {
 	digestPair, err := cacheMgr.PullCache(n.CacheID())
 	if err != nil {
+		// TODO: distinguish cache not found and pull failure.
 		log.Errorf("Failed to fetch intermediate layer with cache ID %s: %s", n.CacheID(), err)
 		return false
 	} else if digestPair == nil {

lib/builder/build_plan.go

@@ -70,11 +70,10 @@ func NewBuildPlan(
 		return nil, fmt.Errorf("build alias list: %s", err)
 	}
 
-	digestPairs := make(image.DigestPairMap)
 	for i, parsedStage := range parsedStages {
 		// Add this stage to the plan.
 		stage, err := newBuildStage(
-			ctx, parsedStage.From.Alias, parsedStage, digestPairs, plan.opts)
+			ctx, parsedStage.From.Alias, parsedStage, plan.opts)
 		if err != nil {
 			return nil, fmt.Errorf("failed to convert parsed stage: %s", err)
 		}
@@ -86,7 +85,7 @@ func NewBuildPlan(
 		plan.stages[i] = stage
 	}
 
-	if err := plan.handleCopyFromDirs(aliases, digestPairs); err != nil {
+	if err := plan.handleCopyFromDirs(aliases); err != nil {
 		return nil, fmt.Errorf("handle cross refs: %s", err)
 	}
 	return plan, nil
@@ -95,9 +94,7 @@ func NewBuildPlan(
 // handleCopyFromDirs goes through all of the stages in the build plan and looks
 // at the `COPY --from` steps to make sure they are valid. If the --from source
 // is another image, we create a new image stage in the build plan.
-func (plan *BuildPlan) handleCopyFromDirs(
-	aliases map[string]bool, digestPairs image.DigestPairMap) error {
-
+func (plan *BuildPlan) handleCopyFromDirs(aliases map[string]bool) error {
 	for _, stage := range plan.stages {
 		for alias, dirs := range stage.copyFromDirs {
 			if _, ok := aliases[alias]; !ok {
@@ -110,7 +107,7 @@ func (plan *BuildPlan) handleCopyFromDirs(
 					return fmt.Errorf("copy from nonexistent stage %s", alias)
 				}
 				remoteImageStage, err := newRemoteImageStage(
-					plan.baseCtx, alias, digestPairs, plan.opts)
+					plan.baseCtx, alias, plan.opts)
 				if err != nil {
 					return fmt.Errorf("new image stage: %s", err)
 				}
@@ -135,7 +132,7 @@ func buildAliases(stages dockerfile.Stages) (map[string]bool, error) {
 			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`
+				// 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)
 			}
 		} else {
@@ -148,13 +145,6 @@ func buildAliases(stages dockerfile.Stages) (map[string]bool, error) {
 
 // Execute executes all build stages in order.
 func (plan *BuildPlan) Execute() (*image.DistributionManifest, error) {
-	// Execute pre-build procedures. Try to pull some reusable layers from the
-	// registry.
-	// TODO: Pull in parallel
-	for _, stage := range plan.stages {
-		stage.pullCacheLayers(plan.cacheMgr)
-	}
-
 	for alias, stage := range plan.remoteImageStages {
 		// Building that pseudo stage will unpack the image directly into the
 		// stage's cross stage directory.
@@ -174,6 +164,9 @@ func (plan *BuildPlan) Execute() (*image.DistributionManifest, error) {
 		currStage = plan.stages[k]
 		log.Infof("* Stage %d/%d : %s", k+1, len(plan.stages), currStage.String())
 
+		// Try to pull reusable layers cached from previous builds.
+		currStage.pullCacheLayers(plan.cacheMgr)
+
 		lastStage := k == len(plan.stages)-1
 		_, copiedFrom := plan.copyFromDirs[currStage.alias]
 

lib/builder/build_stage.go

@@ -46,20 +46,19 @@ type buildStageOptions struct {
 
 // buildStage represents a sequence of steps to build intermediate layers or a final image.
 type buildStage struct {
-	ctx               *context.BuildContext
-	alias             string
-	copyFromDirs      map[string][]string
-	nodes             []*buildNode
-	lastImageConfig   *image.Config
-	sharedDigestPairs image.DigestPairMap
+	ctx             *context.BuildContext
+	alias           string
+	copyFromDirs    map[string][]string
+	nodes           []*buildNode
+	lastImageConfig *image.Config
 
 	opts *buildStageOptions
 }
 
 // newBuildStage initializes a buildStage.
 func newBuildStage(
 	baseCtx *context.BuildContext, alias string, parsedStage *dockerfile.Stage,
-	digestPairs image.DigestPairMap, planOpts *buildPlanOptions) (*buildStage, error) {
+	planOpts *buildPlanOptions) (*buildStage, error) {
 
 	// Create a new build context for the stage.
 	ctx, err := context.NewBuildContext(
@@ -74,13 +73,12 @@ func newBuildStage(
 		return nil, fmt.Errorf("new dockerfile steps: %s", err)
 	}
 
-	return newBuildStageHelper(ctx, alias, steps, digestPairs, planOpts)
+	return newBuildStageHelper(ctx, alias, steps, planOpts)
 }
 
 // newRemoteImageStage initializes a buildStage.
 func newRemoteImageStage(
-	baseCtx *context.BuildContext, alias string, digestPairs image.DigestPairMap,
-	planOpts *buildPlanOptions) (*buildStage, error) {
+	baseCtx *context.BuildContext, alias string, planOpts *buildPlanOptions) (*buildStage, error) {
 
 	// Create a new build context for the stage.
 	ctx, err := context.NewBuildContext(
@@ -105,12 +103,12 @@ func newRemoteImageStage(
 		allowModifyFS: planOpts.allowModifyFS,
 	}
 
-	return newBuildStageHelper(ctx, alias, steps, digestPairs, opts)
+	return newBuildStageHelper(ctx, alias, steps, opts)
 }
 
 func newBuildStageHelper(
 	ctx *context.BuildContext, alias string, steps []step.BuildStep,
-	digestPairs image.DigestPairMap, planOpts *buildPlanOptions) (*buildStage, error) {
+	planOpts *buildPlanOptions) (*buildStage, error) {
 
 	// Convert each step to a build node.
 	var requireOnDisk bool
@@ -134,11 +132,10 @@ func newBuildStageHelper(
 	}
 
 	stage := &buildStage{
-		ctx:               ctx,
-		copyFromDirs:      copyFromDirs,
-		alias:             alias,
-		nodes:             nodes,
-		sharedDigestPairs: digestPairs,
+		ctx:          ctx,
+		copyFromDirs: copyFromDirs,
+		alias:        alias,
+		nodes:        nodes,
 		opts: &buildStageOptions{
 			allowModifyFS: planOpts.allowModifyFS,
 			forceCommit:   planOpts.forceCommit,
@@ -172,14 +169,6 @@ func createDockerfileSteps(
 // build performs the build for that stage. There are side effects that should
 // be expected on each node within the stage.
 func (stage *buildStage) build(cacheMgr cache.Manager, lastStage, copiedFrom bool) error {
-	// Reuse the digestpairs that other stages have populated.
-	for _, node := range stage.nodes {
-		if pairs, ok := stage.sharedDigestPairs[node.CacheID()]; ok {
-			log.Infof("* Reusing digest pairs computed from earlier step %s", node.CacheID())
-			node.digestPairs = pairs
-		}
-	}
-
 	var err error
 	diffIDs := make([]image.Digest, 0)
 	histories := make([]image.History, 0)
@@ -214,11 +203,6 @@ func (stage *buildStage) build(cacheMgr cache.Manager, lastStage, copiedFrom boo
 				Author:    "makisu",
 			})
 		}
-
-		// Update the shared map of cacheID to digest pair.
-		if len(node.digestPairs) != 0 {
-			stage.sharedDigestPairs[node.CacheID()] = node.digestPairs
-		}
 	}
 	stage.lastImageConfig.Created = time.Now()
 	stage.lastImageConfig.History = histories
@@ -310,12 +294,12 @@ func (stage *buildStage) saveManifest(
 	return manifest, nil
 }
 
-// pullCacheLayers attempts to pull reusable layers from the distributed cache. Terminates once
-// a node that can be cached fails to pull its layer.
+// pullCacheLayers attempts to pull reusable layers from the distributed cache.
+// Terminates once a node that can be cached fails to pull its layer.
 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.
+	// 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 {

lib/builder/build_stage_test.go

@@ -123,10 +123,10 @@ func TestPullCacheLayers(t *testing.T) {
 				allowModifyFS: false,
 			}
 
-			stage, err := newBuildStage(ctx, alias, tc.stage, image.DigestPairMap{}, opts)
+			stage, err := newBuildStage(ctx, alias, tc.stage, opts)
 			require.NoError(err)
 
-			kvStore := keyvalue.MemStore{}
+			kvStore := keyvalue.MockStore{}
 			cacheMgr := cache.New(ctx.ImageStore, kvStore, registry.NoopClientFixture())
 
 			for i, node := range stage.nodes {

lib/builder/step/add_copy_step.go

@@ -90,8 +90,9 @@ func (s *addCopyStep) ContextDirs() (string, []string) {
 	return s.fromStage, s.fromPaths
 }
 
-// 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.
+// 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.
 func (s *addCopyStep) SetCacheID(ctx *context.BuildContext, seed string) error {
 	if s.fromStage != "" {
 		// It is copying from a previous stage, return random bytes.
@@ -120,7 +121,8 @@ func (s *addCopyStep) SetCacheID(ctx *context.BuildContext, seed string) error {
 	return nil
 }
 
-// Execute executes the add/copy step. If modifyFS is true, actually performs the on-disk copy.
+// Execute executes the add/copy step. If modifyFS is true, actually performs
+// the on-disk copy.
 func (s *addCopyStep) Execute(ctx *context.BuildContext, modifyFS bool) (err error) {
 	sourceRoot := s.contextRootDir(ctx)
 	sources := s.resolveFromPaths(ctx)
@@ -188,7 +190,9 @@ func (s *addCopyStep) contextRootDir(ctx *context.BuildContext) string {
 	return ctx.ContextDir
 }
 
-func checksumPathContents(ctx *context.BuildContext, path string, fi os.FileInfo, checksum io.Writer) error {
+func checksumPathContents(
+	ctx *context.BuildContext, path string, fi os.FileInfo, checksum io.Writer) error {
+
 	// Skip special files.
 	if utils.IsSpecialFile(fi) {
 		if fi.IsDir() {

lib/builder/step/base_step.go

@@ -66,13 +66,13 @@ func (s *baseStep) SetCacheID(ctx *context.BuildContext, seed string) error {
 	return nil
 }
 
-// SetWorkingDir set the working dir of the current step
-// Exporting the logic to this method allows for an easier `ApplyCtxAndConfig` overwriting
+// SetWorkingDir set the working dir of the current step.
+// Exporting the logic to this method allows overwriting `ApplyCtxAndConfig`.
 func (s *baseStep) SetWorkingDir(
 	ctx *context.BuildContext, imageConfig *image.Config) error {
 	s.workingDir = ctx.RootDir // Default workingDir to root.
 
-	// Set working dir from imageConfig
+	// Set working dir from imageConfig.
 	if imageConfig != nil && imageConfig.Config.WorkingDir != "" {
 		s.workingDir = os.ExpandEnv(imageConfig.Config.WorkingDir)
 	}

lib/builder/step/step.go

@@ -116,7 +116,8 @@ func NewDockerfileStep(
 		step, err = NewFromStep(s.Args, s.Image, s.Alias)
 	case *dockerfile.HealthcheckDirective:
 		s, _ := d.(*dockerfile.HealthcheckDirective)
-		step, err = NewHealthcheckStep(s.Args, s.Interval, s.Timeout, s.StartPeriod, s.Retries, s.Test, s.Commit)
+		step, err = NewHealthcheckStep(
+			s.Args, s.Interval, s.Timeout, s.StartPeriod, s.Retries, s.Test, s.Commit)
 	case *dockerfile.LabelDirective:
 		s, _ := d.(*dockerfile.LabelDirective)
 		step = NewLabelStep(s.Args, s.Labels, s.Commit)