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

@@ -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

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
 						}

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

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
 }