Only cleanup fs when necessary

bin/makisu/cmd/build.go

@@ -236,19 +236,14 @@ func (cmd *buildCmd) Build(contextDir string) error {
 	}
 	defer buildContext.Cleanup()
 
-	// Make sure sandbox is cleaned after build.
-	// Optionally remove everything before and after build.
+	// Optionally stash everything before build and restore after build.
 	defer storage.CleanupSandbox(cmd.storageDir)
-	if cmd.allowModifyFS {
-		if cmd.preserveRoot {
-			rootPreserver, err := storage.NewRootPreserver("/", cmd.storageDir, pathutils.DefaultBlacklist)
-			if err != nil {
-				return fmt.Errorf("failed to preserve root: %s", err)
-			}
-			defer rootPreserver.RestoreRoot()
+	if cmd.allowModifyFS && cmd.preserveRoot {
+		rootPreserver, err := storage.NewRootPreserver("/", cmd.storageDir, pathutils.DefaultBlacklist)
+		if err != nil {
+			return fmt.Errorf("failed to preserve root: %s", err)
 		}
-		buildContext.MemFS.Remove()
-		defer buildContext.MemFS.Remove()
+		defer rootPreserver.RestoreRoot()
 	}
 
 	// Create and execute build plan.

lib/builder/build_plan.go

@@ -234,6 +234,17 @@ func (plan *BuildPlan) Execute() (*image.DistributionManifest, error) {
 }
 
 func (plan *BuildPlan) executeStage(stage *buildStage, lastStage, copiedFrom bool) error {
+	if stage.opts.requireOnDisk {
+		if !stage.opts.allowModifyFS {
+			return fmt.Errorf("fs not allowed to be modified")
+		}
+		// Clean up local fs before build if current stage relies on local fs
+		// (i.e. contains RUN, or is source of COPY --from).
+		if err := stage.cleanup(); err != nil {
+			return fmt.Errorf("cleanup stage %s before build: %s", stage.alias, err)
+		}
+	}
+
 	if err := stage.build(plan.cacheMgr, lastStage, copiedFrom); err != nil {
 		return fmt.Errorf("build stage %s: %s", stage.alias, err)
 	}
@@ -245,10 +256,6 @@ func (plan *BuildPlan) executeStage(stage *buildStage, lastStage, copiedFrom boo
 		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

lib/builder/build_stage.go

@@ -175,9 +175,6 @@ func (stage *buildStage) build(cacheMgr cache.Manager, lastStage, copiedFrom boo
 	for i, node := range stage.nodes {
 		// Build current step from the previous image config (possibly cached).
 		modifyFS := stage.opts.requireOnDisk || copiedFrom
-		if modifyFS && !stage.opts.allowModifyFS {
-			return fmt.Errorf("fs not allowed to be modified")
-		}
 		skipBuild := i < stage.latestFetched() && i > 0
 		lastStep := i == len(stage.nodes)-1
 		forceCommit := i == 0 || (lastStage && lastStep) || stage.opts.forceCommit