overlay differ: cancel diff calculation cleanly

cache/blobs_linux.go

@@ -11,6 +11,7 @@ import (
 	"github.com/containerd/containerd/content"
 	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/mount"
+	"github.com/moby/buildkit/util/bklog"
 	"github.com/moby/buildkit/util/overlay"
 	digest "github.com/opencontainers/go-digest"
 	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
@@ -41,8 +42,10 @@ func (sr *immutableRef) tryComputeOverlayBlob(ctx context.Context, lower, upper
 		return emptyDesc, false, errors.Wrap(err, "failed to open writer")
 	}
 	defer func() {
-		if err != nil {
-			cw.Close()
+		if cw != nil {
+			if cerr := cw.Close(); cerr != nil {
+				bklog.G(ctx).WithError(cerr).Warnf("failed to close writer %q", ref)
+			}
 		}
 	}()
 
@@ -82,6 +85,10 @@ func (sr *immutableRef) tryComputeOverlayBlob(ctx context.Context, lower, upper
 			return emptyDesc, false, errors.Wrap(err, "failed to commit")
 		}
 	}
+	if err := cw.Close(); err != nil {
+		return emptyDesc, false, err
+	}
+	cw = nil
 	cinfo, err := sr.cm.ContentStore.Info(ctx, dgst)
 	if err != nil {
 		return emptyDesc, false, errors.Wrap(err, "failed to get info from content store")

util/overlay/overlay_linux.go

@@ -128,18 +128,30 @@ func WriteUpperdir(ctx context.Context, w io.Writer, upperdir string, lower []mo
 	}
 	return mount.WithTempMount(ctx, lower, func(lowerRoot string) error {
 		return mount.WithTempMount(ctx, upperView, func(upperViewRoot string) error {
-			cw := archive.NewChangeWriter(w, upperViewRoot)
+			cw := archive.NewChangeWriter(&cancellableWriter{ctx, w}, upperViewRoot)
 			if err := Changes(ctx, cw.HandleChange, upperdir, upperViewRoot, lowerRoot); err != nil {
 				if err2 := cw.Close(); err2 != nil {
-					return errors.Wrapf(err, "failed torecord upperdir changes (close error: %v)", err2)
+					return errors.Wrapf(err, "failed to record upperdir changes (close error: %v)", err2)
 				}
-				return errors.Wrapf(err, "failed torecord upperdir changes")
+				return errors.Wrapf(err, "failed to record upperdir changes")
 			}
 			return cw.Close()
 		})
 	})
 }
 
+type cancellableWriter struct {
+	ctx context.Context
+	w   io.Writer
+}
+
+func (w *cancellableWriter) Write(p []byte) (int, error) {
+	if err := w.ctx.Err(); err != nil {
+		return 0, err
+	}
+	return w.w.Write(p)
+}
+
 // Changes is continuty's `fs.Change`-like method but leverages overlayfs's
 // "upperdir" for computing the diff. "upperdirView" is overlayfs mounted view of
 // the upperdir that doesn't contain whiteouts. This is used for computing
@@ -149,6 +161,9 @@ func Changes(ctx context.Context, changeFn fs.ChangeFunc, upperdir, upperdirView
 		if err != nil {
 			return err
 		}
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
 
 		// Rebase path
 		path, err = filepath.Rel(upperdir, path)