[moby#2112] progress.Controller should own the progress.Writer to pre…

…vent leaks

Signed-off-by: Cory Bennett <cbennett@netflix.com>

source/containerimage/pull.go

@@ -200,10 +200,7 @@ func (p *puller) CacheKey(ctx context.Context, g session.Group, index int) (cach
 		}
 
 		if len(p.manifest.Descriptors) > 0 {
-			pw, _, _ := progress.FromContext(ctx)
-			progressController := &controller.Controller{
-				Writer: pw,
-			}
+			progressController := &controller.Controller{}
 			if p.vtx != nil {
 				progressController.Digest = p.vtx.Digest()
 				progressController.Name = p.vtx.Name()

util/progress/controller/controller.go

@@ -13,57 +13,60 @@ import (
 type Controller struct {
 	count   int64
 	started *time.Time
+	writer  progress.Writer
 
 	Digest digest.Digest
 	Name   string
-	Writer progress.Writer
 }
 
 var _ progress.Controller = &Controller{}
 
 func (c *Controller) Start(ctx context.Context) (context.Context, func(error)) {
-	if c.Digest == "" {
-		return progress.WithProgress(ctx, c.Writer), func(error) {}
-	}
-
 	if atomic.AddInt64(&c.count, 1) == 1 {
 		if c.started == nil {
 			now := time.Now()
 			c.started = &now
+			c.writer, _, ctx = progress.FromContext(ctx)
+		}
+
+		if c.Digest != "" {
+			c.writer.Write(c.Digest.String(), client.Vertex{
+				Digest:  c.Digest,
+				Name:    c.Name,
+				Started: c.started,
+			})
 		}
-		c.Writer.Write(c.Digest.String(), client.Vertex{
-			Digest:  c.Digest,
-			Name:    c.Name,
-			Started: c.started,
-		})
 	}
-	return progress.WithProgress(ctx, c.Writer), func(err error) {
+	return progress.WithProgress(ctx, c.writer), func(err error) {
 		if atomic.AddInt64(&c.count, -1) == 0 {
 			now := time.Now()
 			var errString string
 			if err != nil {
 				errString = err.Error()
 			}
-			c.Writer.Write(c.Digest.String(), client.Vertex{
-				Digest:    c.Digest,
-				Name:      c.Name,
-				Started:   c.started,
-				Completed: &now,
-				Error:     errString,
-			})
+			if c.Digest != "" {
+				c.writer.Write(c.Digest.String(), client.Vertex{
+					Digest:    c.Digest,
+					Name:      c.Name,
+					Started:   c.started,
+					Completed: &now,
+					Error:     errString,
+				})
+			}
+			c.writer.Close()
 		}
 	}
 }
 
 func (c *Controller) Status(id string, action string) func() {
 	start := time.Now()
-	c.Writer.Write(id, progress.Status{
+	c.writer.Write(id, progress.Status{
 		Action:  action,
 		Started: &start,
 	})
 	return func() {
 		complete := time.Now()
-		c.Writer.Write(id, progress.Status{
+		c.writer.Write(id, progress.Status{
 			Action:    action,
 			Started:   &start,
 			Completed: &complete,

worker/base/worker.go

@@ -42,7 +42,6 @@ import (
 	"github.com/moby/buildkit/source/http"
 	"github.com/moby/buildkit/source/local"
 	"github.com/moby/buildkit/util/archutil"
-	"github.com/moby/buildkit/util/progress"
 	"github.com/moby/buildkit/util/progress/controller"
 	"github.com/moby/buildkit/worker"
 	digest "github.com/opencontainers/go-digest"
@@ -373,10 +372,9 @@ func (w *Worker) Exporter(name string, sm *session.Manager) (exporter.Exporter,
 }
 
 func (w *Worker) FromRemote(ctx context.Context, remote *solver.Remote) (ref cache.ImmutableRef, err error) {
-	pw, _, _ := progress.FromContext(ctx)
 	descHandler := &cache.DescHandler{
 		Provider: func(session.Group) content.Provider { return remote.Provider },
-		Progress: &controller.Controller{Writer: pw},
+		Progress: &controller.Controller{},
 	}
 	descHandlers := cache.DescHandlers(make(map[digest.Digest]*cache.DescHandler))
 	for _, desc := range remote.Descriptors {