cmd/gomote: add -collect flag to run and refactor output writing

For golang/go#53956.

Change-Id: I13e0c2a3a0d75a7e7bdd072963e158ccb684a3bf
Reviewed-on: https://go-review.googlesource.com/c/build/+/418781
Run-TryBot: Michael Knyszek <[email protected]>
Reviewed-by: Carlos Amedee <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>
Auto-Submit: Michael Knyszek <[email protected]>

cmd/gomote/create.go

@@ -13,6 +13,7 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"path/filepath"
 	"sort"
 	"strings"
 	"sync"
@@ -173,24 +174,18 @@ func create(args []string) error {
 	}
 	builderType := fs.Arg(0)
 
-	var tmpOutDir string
-	var err error
-	if setup {
-		tmpOutDir, err = os.MkdirTemp("", "gomote")
-		if err != nil {
-			return fmt.Errorf("failed to create a temporary directory for setup output: %v", err)
-		}
-	}
-
 	var groupMu sync.Mutex
 	group := activeGroup
+	var err error
 	if newGroup != "" {
 		group, err = doCreateGroup(newGroup)
 		if err != nil {
 			return err
 		}
 	}
 
+	var tmpOutDir string
+	var tmpOutDirOnce sync.Once
 	eg, ctx := errgroup.WithContext(context.Background())
 	client := gomoteServerClient(ctx)
 	for i := 0; i < count; i++ {
@@ -225,6 +220,17 @@ func create(args []string) error {
 			if !setup {
 				return nil
 			}
+
+			// -setup is set, so push GOROOT and run make.bash.
+
+			tmpOutDirOnce.Do(func() {
+				tmpOutDir, err = os.MkdirTemp("", "gomote")
+			})
+			if err != nil {
+				return fmt.Errorf("failed to create a temporary directory for setup output: %v", err)
+			}
+
+			// Push GOROOT.
 			detailedProgress := count == 1
 			goroot, err := getGOROOT()
 			if err != nil {
@@ -236,14 +242,33 @@ func create(args []string) error {
 			if err := doPush(ctx, inst, goroot, false, detailedProgress); err != nil {
 				return err
 			}
+
+			// Run make.bash or make.bat.
 			cmd := "go/src/make.bash"
 			if strings.Contains(builderType, "windows") {
 				cmd = "go/src/make.bat"
 			}
-			if !detailedProgress {
+
+			// Create a file to write output to so it doesn't get lost.
+			outf, err := os.Create(filepath.Join(tmpOutDir, fmt.Sprintf("%s.stdout", inst)))
+			if err != nil {
+				return err
+			}
+			defer func() {
+				outf.Close()
+				fmt.Fprintf(os.Stderr, "# Wrote results from %q to %q.\n", inst, outf.Name())
+			}()
+			fmt.Fprintf(os.Stderr, "# Streaming results from %q to %q...\n", inst, outf.Name())
+
+			// If this is the only command running, print to stdout too, for convenience and
+			// backwards compatibility.
+			outputs := []io.Writer{outf}
+			if detailedProgress {
+				outputs = append(outputs, os.Stdout)
+			} else {
 				fmt.Fprintf(os.Stderr, "# Running %q on %q...\n", cmd, inst)
 			}
-			return doRun(ctx, inst, tmpOutDir, cmd, []string{}, count == 1)
+			return doRun(ctx, inst, cmd, []string{}, runWriters(outputs...))
 		})
 	}
 	if err := eg.Wait(); err != nil {

cmd/gomote/get.go

@@ -71,6 +71,10 @@ func getTar(args []string) error {
 	}
 
 	name := fs.Arg(0)
+	return doGetTar(name, dir, os.Stdout)
+}
+
+func doGetTar(name, dir string, out io.Writer) error {
 	ctx := context.Background()
 	client := gomoteServerClient(ctx)
 	resp, err := client.ReadTGZToURL(ctx, &protos.ReadTGZToURLRequest{
@@ -95,7 +99,7 @@ func getTar(args []string) error {
 		return fmt.Errorf("unable to download tgz: %s", err)
 	}
 	defer r.Body.Close()
-	_, err = io.Copy(os.Stdout, r.Body)
+	_, err = io.Copy(out, r.Body)
 	if err != nil {
 		return fmt.Errorf("unable to copy tgz to stdout: %s", err)
 	}

cmd/gomote/run.go

@@ -6,12 +6,14 @@ package main
 
 import (
 	"context"
+	"errors"
 	"flag"
 	"fmt"
 	"io"
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 
 	"golang.org/x/build/buildlet"
 	"golang.org/x/build/dashboard"
@@ -136,6 +138,9 @@ func run(args []string) error {
 	var builderEnv string
 	fs.StringVar(&builderEnv, "builderenv", "", "Optional alternate builder to act like. Must share the same underlying buildlet host type, or it's an error. For instance, linux-amd64-race or linux-386-387 are compatible with linux-amd64, but openbsd-amd64 and openbsd-386 are different hosts.")
 
+	var collect bool
+	fs.BoolVar(&collect, "collect", false, "Collect artifacts (stdout, work dir .tar.gz) into $PWD once complete.")
+
 	fs.Parse(args)
 	if fs.NArg() == 0 {
 		fs.Usage()
@@ -172,7 +177,6 @@ func run(args []string) error {
 		cmdArgs = fs.Args()[2:]
 	}
 
-	detailedProgress := len(runSet) == 1
 	var pathOpt []string
 	if path == "EMPTY" {
 		pathOpt = []string{} // non-nil
@@ -183,71 +187,128 @@ func run(args []string) error {
 	// Create temporary directory for output.
 	// This is useful even if we don't have multiple gomotes running, since
 	// it's easy to accidentally lose the output.
-	tmpOutDir, err := os.MkdirTemp("", "gomote")
-	if err != nil {
-		return err
+	var outDir string
+	if collect {
+		outDir, err = os.Getwd()
+		if err != nil {
+			return err
+		}
+	} else {
+		outDir, err = os.MkdirTemp("", "gomote")
+		if err != nil {
+			return err
+		}
 	}
 
+	var cmdsFailedMu sync.Mutex
+	var cmdsFailed []*cmdFailedError
 	eg, ctx := errgroup.WithContext(context.Background())
 	for _, inst := range runSet {
 		inst := inst
-		if !detailedProgress {
+		if len(runSet) > 1 {
+			// There's more than one instance running the command, so let's
+			// be explicit about that.
 			fmt.Fprintf(os.Stderr, "# Running command on %q...\n", inst)
 		}
 		eg.Go(func() error {
-			return doRun(
+			// Create a file to write output to so it doesn't get lost.
+			outf, err := os.Create(filepath.Join(outDir, fmt.Sprintf("%s.stdout", inst)))
+			if err != nil {
+				return err
+			}
+			defer func() {
+				outf.Close()
+				fmt.Fprintf(os.Stderr, "# Wrote results from %q to %q.\n", inst, outf.Name())
+			}()
+			fmt.Fprintf(os.Stderr, "# Streaming results from %q to %q...\n", inst, outf.Name())
+
+			outputs := []io.Writer{outf}
+			// If this is the only command running, print to stdout too, for convenience and
+			// backwards compatibility.
+			if len(runSet) == 1 {
+				outputs = append(outputs, os.Stdout)
+			}
+			err = doRun(
 				ctx,
 				inst,
-				tmpOutDir,
 				cmd,
 				cmdArgs,
-				detailedProgress,
 				runDir(dir),
 				runBuilderEnv(builderEnv),
 				runEnv(env),
 				runPath(pathOpt),
 				runSystem(sys),
 				runDebug(debug),
 				runFirewall(firewall),
+				runWriters(outputs...),
 			)
+			// If it's just that the command failed, don't exit just yet, and don't return
+			// an error to the errgroup because we want the other commands to keep going.
+			if err != nil {
+				ce, ok := err.(*cmdFailedError)
+				if !ok {
+					return err
+				}
+				cmdsFailedMu.Lock()
+				cmdsFailed = append(cmdsFailed, ce)
+				cmdsFailedMu.Unlock()
+				// Write out the error.
+				_, err := io.MultiWriter(outputs...).Write([]byte(err.Error() + "\n"))
+				if err != nil {
+					fmt.Fprintf(os.Stderr, "failed to write error to output: %v", err)
+				}
+			}
+			if collect {
+				f, err := os.Create(fmt.Sprintf("%s.tar.gz", inst))
+				if err != nil {
+					fmt.Fprintf(os.Stderr, "failed to create file to write instance tarball: %v", err)
+					return nil
+				}
+				defer f.Close()
+				fmt.Fprintf(os.Stderr, "# Downloading work dir tarball for %q to %q...\n", inst, f.Name())
+				if err := doGetTar(inst, ".", f); err != nil {
+					fmt.Fprintf(os.Stderr, "failed to retrieve instance tarball: %v", err)
+					return nil
+				}
+			}
+			return nil
 		})
 	}
-	return eg.Wait()
-}
-
-func doRun(ctx context.Context, inst, tmpOutDir, cmd string, cmdArgs []string, detailedProgress bool, opts ...runOpt) error {
-	req := &protos.ExecuteCommandRequest{
-		AppendEnvironment: []string{},
-		Args:              cmdArgs,
-		Command:           cmd,
-		Path:              []string{},
-		GomoteId:          inst,
+	if err := eg.Wait(); err != nil {
+		return err
 	}
-	for _, opt := range opts {
-		opt(req)
+	// Handle failed commands separately so that we can let all the instances finish
+	// running. We still want to handle them, though, because we want to make sure
+	// we exit with a non-zero exit code to reflect the command failure.
+	for _, ce := range cmdsFailed {
+		fmt.Fprintf(os.Stderr, "# Command %q failed on %q: %v\n", ce.cmd, ce.inst, err)
 	}
-	if !req.SystemLevel {
-		req.SystemLevel = strings.HasPrefix(cmd, "/")
+	if len(cmdsFailed) > 0 {
+		return errors.New("one or more commands failed")
 	}
+	return nil
+}
 
-	outf, err := os.Create(filepath.Join(tmpOutDir, fmt.Sprintf("%s.stdout", inst)))
-	if err != nil {
-		return err
+func doRun(ctx context.Context, inst, cmd string, cmdArgs []string, opts ...runOpt) error {
+	cfg := &runCfg{
+		req: protos.ExecuteCommandRequest{
+			AppendEnvironment: []string{},
+			Args:              cmdArgs,
+			Command:           cmd,
+			Path:              []string{},
+			GomoteId:          inst,
+		},
 	}
-	defer func() {
-		outf.Close()
-		fmt.Fprintf(os.Stderr, "# Wrote results from %q to %q.\n", inst, outf.Name())
-	}()
-	fmt.Fprintf(os.Stderr, "# Streaming results from %q to %q...\n", inst, outf.Name())
-	var outWriter io.Writer
-	if detailedProgress {
-		outWriter = io.MultiWriter(os.Stdout, outf)
-	} else {
-		outWriter = outf
+	for _, opt := range opts {
+		opt(cfg)
+	}
+	if !cfg.req.SystemLevel {
+		cfg.req.SystemLevel = strings.HasPrefix(cmd, "/")
 	}
 
+	outWriter := io.MultiWriter(cfg.outputs...)
 	client := gomoteServerClient(ctx)
-	stream, err := client.ExecuteCommand(ctx, req)
+	stream, err := client.ExecuteCommand(ctx, &cfg.req)
 	if err != nil {
 		return fmt.Errorf("unable to execute %s: %s", cmd, statusFromError(err))
 	}
@@ -259,7 +320,7 @@ func doRun(ctx context.Context, inst, tmpOutDir, cmd string, cmdArgs []string, d
 		if err != nil {
 			// execution error
 			if status.Code(err) == codes.Aborted {
-				return fmt.Errorf("Error trying to execute %s: %v", cmd, statusFromError(err))
+				return &cmdFailedError{inst: inst, cmd: cmd, err: err}
 			}
 			// remote error
 			return fmt.Errorf("unable to execute %s: %s", cmd, statusFromError(err))
@@ -268,46 +329,66 @@ func doRun(ctx context.Context, inst, tmpOutDir, cmd string, cmdArgs []string, d
 	}
 }
 
-type runOpt func(*protos.ExecuteCommandRequest)
+type cmdFailedError struct {
+	inst, cmd string
+	err       error
+}
+
+func (e *cmdFailedError) Error() string {
+	return fmt.Sprintf("Error trying to execute %s: %v", e.cmd, statusFromError(e.err))
+}
+
+type runCfg struct {
+	outputs []io.Writer
+	req     protos.ExecuteCommandRequest
+}
+
+type runOpt func(*runCfg)
 
 func runBuilderEnv(builderEnv string) runOpt {
-	return func(r *protos.ExecuteCommandRequest) {
-		r.ImitateHostType = builderEnv
+	return func(r *runCfg) {
+		r.req.ImitateHostType = builderEnv
 	}
 }
 
 func runDir(dir string) runOpt {
-	return func(r *protos.ExecuteCommandRequest) {
-		r.Directory = dir
+	return func(r *runCfg) {
+		r.req.Directory = dir
 	}
 }
 
 func runEnv(env []string) runOpt {
-	return func(r *protos.ExecuteCommandRequest) {
-		r.AppendEnvironment = append(r.AppendEnvironment, env...)
+	return func(r *runCfg) {
+		r.req.AppendEnvironment = append(r.req.AppendEnvironment, env...)
 	}
 }
 
 func runPath(path []string) runOpt {
-	return func(r *protos.ExecuteCommandRequest) {
-		r.Path = append(r.Path, path...)
+	return func(r *runCfg) {
+		r.req.Path = append(r.req.Path, path...)
 	}
 }
 
 func runDebug(debug bool) runOpt {
-	return func(r *protos.ExecuteCommandRequest) {
-		r.Debug = debug
+	return func(r *runCfg) {
+		r.req.Debug = debug
 	}
 }
 
 func runSystem(sys bool) runOpt {
-	return func(r *protos.ExecuteCommandRequest) {
-		r.SystemLevel = sys
+	return func(r *runCfg) {
+		r.req.SystemLevel = sys
 	}
 }
 
 func runFirewall(firewall bool) runOpt {
-	return func(r *protos.ExecuteCommandRequest) {
-		r.AppendEnvironment = append(r.AppendEnvironment, "GO_DISABLE_OUTBOUND_NETWORK="+fmt.Sprint(firewall))
+	return func(r *runCfg) {
+		r.req.AppendEnvironment = append(r.req.AppendEnvironment, "GO_DISABLE_OUTBOUND_NETWORK="+fmt.Sprint(firewall))
+	}
+}
+
+func runWriters(writers ...io.Writer) runOpt {
+	return func(r *runCfg) {
+		r.outputs = writers
 	}
 }