From 40fad4bbb57808258aded2f40074350c6dec10f8 Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Fri, 16 Oct 2020 22:49:29 -0700
Subject: [PATCH] progress: make sure all channels have written before
 returning

Possible write on closed channel on cancellation before.

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 bake/remote.go          | 4 +++-
 build/build.go          | 4 +++-
 build/url.go            | 4 +++-
 util/progress/writer.go | 6 ++++--
 4 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/bake/remote.go b/bake/remote.go
index 4e28946b440..ff50a7c26ba 100644
--- a/bake/remote.go
+++ b/bake/remote.go
@@ -48,6 +48,8 @@ func ReadRemoteFiles(ctx context.Context, dis []build.DriverInfo, url string, na
 		return nil, nil, err
 	}
 
+	ch, done := progress.NewChannel(pw)
+	defer func() { <-done }()
 	_, err = c.Build(ctx, client.SolveOpt{}, "buildx", func(ctx context.Context, c gwclient.Client) (*gwclient.Result, error) {
 		def, err := st.Marshal(ctx)
 		if err != nil {
@@ -71,7 +73,7 @@ func ReadRemoteFiles(ctx context.Context, dis []build.DriverInfo, url string, na
 			files, err = filesFromRef(ctx, ref, names)
 		}
 		return nil, err
-	}, progress.NewChannel(pw))
+	}, ch)
 
 	if err != nil {
 		return nil, nil, err
diff --git a/build/build.go b/build/build.go
index f654d62162f..7f0256fbd4b 100644
--- a/build/build.go
+++ b/build/build.go
@@ -658,7 +658,9 @@ func Build(ctx context.Context, drivers []DriverInfo, opt map[string]Options, do
 
 					eg.Go(func() error {
 						defer wg.Done()
-						rr, err := c.Solve(ctx, nil, so, progress.NewChannel(pw))
+						ch, done := progress.NewChannel(pw)
+						defer func() { <-done }()
+						rr, err := c.Solve(ctx, nil, so, ch)
 						if err != nil {
 							return err
 						}
diff --git a/build/url.go b/build/url.go
index bbde6b478ab..05a3da94cfe 100644
--- a/build/url.go
+++ b/build/url.go
@@ -19,6 +19,8 @@ func createTempDockerfileFromURL(ctx context.Context, d driver.Driver, url strin
 		return "", err
 	}
 	var out string
+	ch, done := progress.NewChannel(pw)
+	defer func() { <-done }()
 	_, err = c.Build(ctx, client.SolveOpt{}, "buildx", func(ctx context.Context, c gwclient.Client) (*gwclient.Result, error) {
 		def, err := llb.HTTP(url, llb.Filename("Dockerfile"), llb.WithCustomNamef("[internal] load %s", url)).Marshal(ctx)
 		if err != nil {
@@ -60,7 +62,7 @@ func createTempDockerfileFromURL(ctx context.Context, d driver.Driver, url strin
 		}
 		out = dir
 		return nil, nil
-	}, progress.NewChannel(pw))
+	}, ch)
 
 	if err != nil {
 		return "", err
diff --git a/util/progress/writer.go b/util/progress/writer.go
index d2a77e1711b..44194b5938f 100644
--- a/util/progress/writer.go
+++ b/util/progress/writer.go
@@ -39,16 +39,18 @@ func Write(w Writer, name string, f func() error) {
 	})
 }
 
-func NewChannel(w Writer) chan *client.SolveStatus {
+func NewChannel(w Writer) (chan *client.SolveStatus, chan struct{}) {
 	ch := make(chan *client.SolveStatus)
+	done := make(chan struct{})
 	go func() {
 		for {
 			v, ok := <-ch
 			if !ok {
+				close(done)
 				return
 			}
 			w.Write(v)
 		}
 	}()
-	return ch
+	return ch, done
 }