client: keep resource initialization and cleanup together

Signed-off-by: Andrea Luzzardi <aluzzardi@gmail.com>

dagger/client.go

@@ -7,7 +7,6 @@ import (
 	"io"
 	"os"
 	"strings"
-	"time"
 
 	"github.com/pkg/errors"
 	"golang.org/x/sync/errgroup"
@@ -112,38 +111,58 @@ func (c *Client) Compute(ctx context.Context) (*Value, error) {
 	eg, ctx := errgroup.WithContext(ctx)
 	events := make(chan *bk.SolveStatus)
 	outr, outw := io.Pipe()
+
 	// Spawn build function
-	eg.Go(c.buildfn(ctx, events, outw))
+	eg.Go(func() error {
+		defer outw.Close()
+		return c.buildfn(ctx, events, outw)
+	})
+
 	// Spawn print function(s)
 	dispCtx := context.TODO()
-	var eventsdup chan *bk.SolveStatus
 	if os.Getenv("DOCKER_OUTPUT") != "" {
-		eventsdup = make(chan *bk.SolveStatus)
-		eg.Go(c.dockerprintfn(dispCtx, eventsdup, os.Stderr))
+		// Multiplex events
+		eventsPrint := make(chan *bk.SolveStatus)
+		eventsDockerPrint := make(chan *bk.SolveStatus)
+		eg.Go(func() error {
+			defer close(eventsPrint)
+			defer close(eventsDockerPrint)
+
+			for e := range events {
+				eventsPrint <- e
+				eventsDockerPrint <- e
 			}
-	eg.Go(c.printfn(dispCtx, events, eventsdup))
+			return nil
+		})
+
+		eg.Go(func() error {
+			return c.printfn(dispCtx, eventsPrint)
+		})
+
+		eg.Go(func() error {
+			return c.dockerprintfn(dispCtx, eventsDockerPrint, os.Stderr)
+		})
+	} else {
+		eg.Go(func() error {
+			return c.printfn(dispCtx, events)
+		})
+	}
+
 	// Retrieve output
-	eg.Go(c.outputfn(ctx, outr, out))
+	eg.Go(func() error {
+		defer outr.Close()
+		return c.outputfn(ctx, outr, out)
+	})
 	return out, eg.Wait()
 }
 
-func (c *Client) buildfn(ctx context.Context, ch chan *bk.SolveStatus, w io.WriteCloser) func() error {
-	return func() (err error) {
-		defer func() {
-			debugf("buildfn complete, err=%q", err)
-			if err != nil {
-				// Close exporter pipe so that export processor can return
-				w.Close()
-			}
-		}()
+func (c *Client) buildfn(ctx context.Context, ch chan *bk.SolveStatus, w io.WriteCloser) error {
 	boot, err := c.BootScript()
 	if err != nil {
-			close(ch)
 		return errors.Wrap(err, "assemble boot script")
 	}
 	bootSource, err := boot.Value().Source()
 	if err != nil {
-			close(ch)
 		return errors.Wrap(err, "serialize boot script")
 	}
 	debugf("client: assembled boot script: %s\n", bootSource)
@@ -167,7 +186,6 @@ func (c *Client) buildfn(ctx context.Context, ch chan *bk.SolveStatus, w io.Writ
 	// Connect local dirs
 	localdirs, err := c.LocalDirs()
 	if err != nil {
-			close(ch)
 		return errors.Wrap(err, "connect local dirs")
 	}
 	for _, dir := range localdirs {
@@ -184,12 +202,10 @@ func (c *Client) buildfn(ctx context.Context, ch chan *bk.SolveStatus, w io.Writ
 		fmt.Printf("exporter response: %s=%s\n", k, v)
 	}
 	return nil
-	}
 }
 
 // Read tar export stream from buildkit Build(), and extract cue output
-func (c *Client) outputfn(_ context.Context, r io.Reader, out *Value) func() error {
-	return func() error {
+func (c *Client) outputfn(_ context.Context, r io.Reader, out *Value) error {
 	defer debugf("outputfn complete")
 	tr := tar.NewReader(r)
 	for {
@@ -217,7 +233,6 @@ func (c *Client) outputfn(_ context.Context, r io.Reader, out *Value) func() err
 		debugf("outputfn: DONE: compiling & merging %q", h.Name)
 	}
 	return nil
-	}
 }
 
 // Status of a node in the config tree being computed
@@ -269,31 +284,21 @@ func (n Node) LogError(errmsg string) {
 	n.Logf("ERROR: %s", bkCleanError(errmsg))
 }
 
-func (c *Client) printfn(ctx context.Context, ch, ch2 chan *bk.SolveStatus) func() error {
-	return func() error {
+func (c *Client) printfn(ctx context.Context, ch chan *bk.SolveStatus) error {
 	// Node status mapped to buildkit vertex digest
 	nodesByDigest := map[string]*Node{}
 	// Node status mapped to cue path
 	nodesByPath := map[string]*Node{}
 
 	defer debugf("printfn complete")
-		if ch2 != nil {
-			defer close(ch2)
-		}
-		ticker := time.NewTicker(150 * time.Millisecond)
-		defer ticker.Stop()
 	for {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
-			case <-ticker.C:
 		case status, ok := <-ch:
 			if !ok {
 				return nil
 			}
-				if ch2 != nil {
-					ch2 <- status
-				}
 			debugf("status event: vertexes:%d statuses:%d logs:%d\n",
 				len(status.Vertexes),
 				len(status.Statuses),
@@ -328,7 +333,6 @@ func (c *Client) printfn(ctx context.Context, ch, ch2 chan *bk.SolveStatus) func
 			// see proto 67
 		}
 	}
-	}
 }
 
 // A helper to remove noise from buildkit error messages.
@@ -346,11 +350,9 @@ func bkCleanError(msg string) string {
 	return msg
 }
 
-func (c *Client) dockerprintfn(ctx context.Context, ch chan *bk.SolveStatus, out io.Writer) func() error {
-	return func() error {
+func (c *Client) dockerprintfn(ctx context.Context, ch chan *bk.SolveStatus, out io.Writer) error {
 	defer debugf("dockerprintfn complete")
 	var cons console.Console
 	// FIXME: use smarter writer from blr
 	return progressui.DisplaySolveStatus(ctx, "", cons, out, ch)
-	}
 }