buildlet: add context argument to Client.Exec

(And update all the callers in the tree.)

Updates golang/go#35707

Change-Id: I504ef73ea4ba7f8028f47a658c1cd519c7b5d986
Reviewed-on: https://go-review.googlesource.com/c/build/+/207908
Reviewed-by: Dmitri Shuralyov <[email protected]>

Author: bradfitz

buildlet/buildletclient.go

@@ -507,9 +507,6 @@ type ExecOpts struct {
 	// response from the buildlet, but before the output begins
 	// writing to Output.
 	OnStartExec func()
-
-	// Timeout is an optional duration before ErrTimeout is returned.
-	Timeout time.Duration
 }
 
 var ErrTimeout = errors.New("buildlet: timeout waiting for command to complete")
@@ -521,7 +518,10 @@ var ErrTimeout = errors.New("buildlet: timeout waiting for command to complete")
 // were system errors preventing the command from being started or
 // seen to completition. If execErr is non-nil, the remoteErr is
 // meaningless.
-func (c *Client) Exec(cmd string, opts ExecOpts) (remoteErr, execErr error) {
+//
+// If the context's deadline is exceeded, the returned execErr is
+// ErrTimeout.
+func (c *Client) Exec(ctx context.Context, cmd string, opts ExecOpts) (remoteErr, execErr error) {
 	var mode string
 	if opts.SystemLevel {
 		mode = "sys"
@@ -545,6 +545,7 @@ func (c *Client) Exec(cmd string, opts ExecOpts) (remoteErr, execErr error) {
 	if err != nil {
 		return nil, err
 	}
+	req = req.WithContext(ctx)
 	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
 
 	// The first thing the buildlet's exec handler does is flush the headers, so
@@ -596,17 +597,16 @@ func (c *Client) Exec(cmd string, opts ExecOpts) (remoteErr, execErr error) {
 			resc <- errs{} // success
 		}
 	}()
-	var timer <-chan time.Time
-	if opts.Timeout > 0 {
-		t := time.NewTimer(opts.Timeout)
-		defer t.Stop()
-		timer = t.C
-	}
 	select {
-	case <-timer:
-		c.MarkBroken()
-		return nil, ErrTimeout
 	case res := <-resc:
+		if res.execErr != nil {
+			c.MarkBroken()
+			if res.execErr == context.DeadlineExceeded {
+				// Historical pre-context value.
+				// TODO: update docs & callers to just use the context value.
+				res.execErr = ErrTimeout
+			}
+		}
 		return res.remoteErr, res.execErr
 	case <-c.peerDead:
 		return nil, c.deadErr

cmd/coordinator/coordinator.go

@@ -2061,7 +2061,7 @@ func (st *buildStatus) runAllSharded() (remoteErr, err error) {
 	st.getHelpersReadySoon()
 
 	if !st.useSnapshot() {
-		remoteErr, err = st.goBuilder().RunMake(st.bc, st)
+		remoteErr, err = st.goBuilder().RunMake(st.ctx, st.bc, st)
 		if err != nil {
 			return nil, err
 		}
@@ -2128,7 +2128,7 @@ func (st *buildStatus) crossCompileMakeAndSnapshot(config *dashboard.CrossCompil
 
 	goos, goarch := st.conf.GOOS(), st.conf.GOARCH()
 
-	remoteErr, err := kubeBC.Exec("/bin/bash", buildlet.ExecOpts{
+	remoteErr, err := kubeBC.Exec(st.ctx, "/bin/bash", buildlet.ExecOpts{
 		SystemLevel: true,
 		Args: []string{
 			"-c",
@@ -2175,7 +2175,7 @@ func (st *buildStatus) crossCompileMakeAndSnapshot(config *dashboard.CrossCompil
 func (st *buildStatus) runAllLegacy() (remoteErr, err error) {
 	allScript := st.conf.AllScript()
 	sp := st.CreateSpan("legacy_all_path", allScript)
-	remoteErr, err = st.bc.Exec(path.Join("go", allScript), buildlet.ExecOpts{
+	remoteErr, err = st.bc.Exec(st.ctx, path.Join("go", allScript), buildlet.ExecOpts{
 		Output:   st,
 		ExtraEnv: st.conf.Env(),
 		Debug:    true,
@@ -2310,7 +2310,7 @@ func (st *buildStatus) distTestList() (names []string, remoteErr, err error) {
 		args = append(args, "--compile-only")
 	}
 	var buf bytes.Buffer
-	remoteErr, err = st.bc.Exec("go/bin/go", buildlet.ExecOpts{
+	remoteErr, err = st.bc.Exec(st.ctx, "go/bin/go", buildlet.ExecOpts{
 		Output:      &buf,
 		ExtraEnv:    append(st.conf.Env(), "GOROOT="+goroot),
 		OnStartExec: func() { st.LogEventTime("discovering_tests") },
@@ -2533,7 +2533,7 @@ func (st *buildStatus) runSubrepoTests() (remoteErr, err error) {
 		repoPath := importPathOfRepo(st.SubName)
 		var buf bytes.Buffer
 		sp := st.CreateSpan("listing_subrepo_packages", st.SubName)
-		rErr, err := st.bc.Exec("go/bin/go", buildlet.ExecOpts{
+		rErr, err := st.bc.Exec(st.ctx, "go/bin/go", buildlet.ExecOpts{
 			Output:   &buf,
 			Dir:      "gopath/src/" + repoPath,
 			ExtraEnv: append(st.conf.Env(), "GOROOT="+goroot, "GOPATH="+gopath),
@@ -2580,7 +2580,7 @@ func (st *buildStatus) runSubrepoTests() (remoteErr, err error) {
 
 	var remoteErrors []error
 	for _, tr := range testRuns {
-		rErr, err := st.bc.Exec("go/bin/go", buildlet.ExecOpts{
+		rErr, err := st.bc.Exec(st.ctx, "go/bin/go", buildlet.ExecOpts{
 			Debug:    true, // make buildlet print extra debug in output for failures
 			Output:   st,
 			Dir:      tr.Dir,
@@ -2610,7 +2610,7 @@ func (st *buildStatus) runSubrepoTests() (remoteErr, err error) {
 // It uses module-specific environment variables from st.conf.ModulesEnv.
 func (st *buildStatus) goMod(importPath, goroot, gopath string) (string, error) {
 	var buf bytes.Buffer
-	rErr, err := st.bc.Exec("go/bin/go", buildlet.ExecOpts{
+	rErr, err := st.bc.Exec(st.ctx, "go/bin/go", buildlet.ExecOpts{
 		Output:   &buf,
 		Dir:      "gopath/src/" + importPath,
 		ExtraEnv: append(append(st.conf.Env(), "GOROOT="+goroot, "GOPATH="+gopath), st.conf.ModulesEnv(st.SubName)...),
@@ -2653,7 +2653,7 @@ func (st *buildStatus) fetchDependenciesToGOPATHWorkspace(goroot, gopath string)
 	findDeps := func(repo string) (rErr, err error) {
 		repoPath := importPathOfRepo(repo)
 		var buf bytes.Buffer
-		rErr, err = st.bc.Exec("go/bin/go", buildlet.ExecOpts{
+		rErr, err = st.bc.Exec(st.ctx, "go/bin/go", buildlet.ExecOpts{
 			Output:   &buf,
 			Dir:      "gopath/src/" + repoPath,
 			ExtraEnv: append(st.conf.Env(), "GOROOT="+goroot, "GOPATH="+gopath),
@@ -3138,12 +3138,16 @@ func (st *buildStatus) runTestsOnBuildlet(bc *buildlet.Client, tis []*testItem,
 	var buf bytes.Buffer
 	t0 := time.Now()
 	timeout := st.conf.DistTestsExecTimeout(names)
+
+	ctx, cancel := context.WithTimeout(st.ctx, timeout)
+	defer cancel()
+
 	var remoteErr, err error
 	if ti := tis[0]; ti.bench != nil {
 		pbr, perr := st.parentRev()
 		// TODO(quentin): Error if parent commit could not be determined?
 		if perr == nil {
-			remoteErr, err = ti.bench.Run(buildEnv, st, st.conf, bc, &buf, []buildgo.BuilderRev{st.BuilderRev, pbr})
+			remoteErr, err = ti.bench.Run(st.ctx, buildEnv, st, st.conf, bc, &buf, []buildgo.BuilderRev{st.BuilderRev, pbr})
 		}
 	} else {
 		env := append(st.conf.Env(),
@@ -3153,7 +3157,7 @@ func (st *buildStatus) runTestsOnBuildlet(bc *buildlet.Client, tis []*testItem,
 		)
 		env = append(env, st.conf.ModulesEnv("go")...)
 
-		remoteErr, err = bc.Exec("go/bin/go", buildlet.ExecOpts{
+		remoteErr, err = bc.Exec(ctx, "go/bin/go", buildlet.ExecOpts{
 			// We set Dir to "." instead of the default ("go/bin") so when the dist tests
 			// try to run os/exec.Command("go", "test", ...), the LookPath of "go" doesn't
 			// return "./go.exe" (which exists in the current directory: "go/bin") and then
@@ -3163,7 +3167,6 @@ func (st *buildStatus) runTestsOnBuildlet(bc *buildlet.Client, tis []*testItem,
 			Dir:      ".",
 			Output:   &buf, // see "maybe stream lines" TODO below
 			ExtraEnv: env,
-			Timeout:  timeout,
 			Path:     []string{"$WORKDIR/go/bin", "$PATH"},
 			Args:     args,
 		})

cmd/debugnewvm/debugnewvm.go

@@ -128,12 +128,12 @@ func main() {
 	log.Printf("WorkDir: %v, %v", dir, err)
 
 	if *sleepSec > 0 {
-		bc.Exec("sysctl", buildlet.ExecOpts{
+		bc.Exec(ctx, "sysctl", buildlet.ExecOpts{
 			Output:      os.Stdout,
 			SystemLevel: true,
 			Args:        []string{"kern.timecounter.hardware"},
 		})
-		bc.Exec("bash", buildlet.ExecOpts{
+		bc.Exec(ctx, "bash", buildlet.ExecOpts{
 			Output:      os.Stdout,
 			SystemLevel: true,
 			Args:        []string{"-c", "rdate -p -v time.nist.gov; sleep " + fmt.Sprint(*sleepSec) + "; rdate -p -v time.nist.gov"},
@@ -172,7 +172,7 @@ func main() {
 		}
 		t0 := time.Now()
 		log.Printf("Running %s ...", script)
-		remoteErr, err := bc.Exec(path.Join("go", script), buildlet.ExecOpts{
+		remoteErr, err := bc.Exec(ctx, path.Join("go", script), buildlet.ExecOpts{
 			Output:   os.Stdout,
 			ExtraEnv: bconf.Env(),
 			Debug:    true,

cmd/gomote/run.go

@@ -5,6 +5,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"os"
@@ -71,7 +72,7 @@ func run(args []string) error {
 	}
 	env = append(env, "GO_DISABLE_OUTBOUND_NETWORK="+fmt.Sprint(firewall))
 
-	remoteErr, execErr := bc.Exec(cmd, buildlet.ExecOpts{
+	remoteErr, execErr := bc.Exec(context.Background(), cmd, buildlet.ExecOpts{
 		Dir:         dir,
 		SystemLevel: sys || strings.HasPrefix(cmd, "/"),
 		Output:      os.Stdout,

cmd/perfrun/perfrun.go

@@ -9,6 +9,7 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"flag"
 	"fmt"
@@ -58,7 +59,7 @@ func runBench(out io.Writer, bench, src string, commits []string) error {
 		// Build binaries
 		log.Printf("Building bench binary for rev %s", rev)
 		var buf bytes.Buffer
-		remoteErr, err := bc.Exec(path.Join(dir, "bin", "go"), buildlet.ExecOpts{
+		remoteErr, err := bc.Exec(context.Background(), path.Join(dir, "bin", "go"), buildlet.ExecOpts{
 			Output:   &buf,
 			ExtraEnv: []string{"GOROOT=" + path.Join(workDir, dir)},
 			Args:     []string{"test", "-c"},
@@ -81,7 +82,7 @@ func runBench(out io.Writer, bench, src string, commits []string) error {
 		log.Printf("Starting bench run %d", i)
 		for _, rev := range commits {
 			var buf bytes.Buffer
-			remoteErr, err := bc.Exec(path.Join("go-"+rev, "test/bench/go1/go1.test"), buildlet.ExecOpts{
+			remoteErr, err := bc.Exec(context.Background(), path.Join("go-"+rev, "test/bench/go1/go1.test"), buildlet.ExecOpts{
 				Output: &buf,
 				Args:   []string{"-test.bench", ".", "-test.benchmem"},
 			})

cmd/release/release.go

@@ -366,7 +366,7 @@ func (b *Build) make() error {
 	if *watch && *target != "" {
 		execOut = io.MultiWriter(out, os.Stdout)
 	}
-	remoteErr, err := client.Exec(filepath.Join(goDir, bc.MakeScript()), buildlet.ExecOpts{
+	remoteErr, err := client.Exec(context.Background(), filepath.Join(goDir, bc.MakeScript()), buildlet.ExecOpts{
 		Output:   execOut,
 		ExtraEnv: env,
 		Args:     bc.MakeScriptArgs(),
@@ -396,7 +396,7 @@ func (b *Build) make() error {
 		if len(args) > 0 && args[0] == "run" && hostArch != "" {
 			cmdEnv = setGOARCH(cmdEnv, hostArch)
 		}
-		remoteErr, err := client.Exec(goCmd, buildlet.ExecOpts{
+		remoteErr, err := client.Exec(context.Background(), goCmd, buildlet.ExecOpts{
 			Output:   execOut,
 			Dir:      ".", // root of buildlet work directory
 			Args:     args,
@@ -595,7 +595,7 @@ func (b *Build) make() error {
 		if *watch && *target != "" {
 			execOut = io.MultiWriter(out, os.Stdout)
 		}
-		remoteErr, err := client.Exec(filepath.Join(goDir, bc.AllScript()), buildlet.ExecOpts{
+		remoteErr, err := client.Exec(context.Background(), filepath.Join(goDir, bc.AllScript()), buildlet.ExecOpts{
 			Output:   execOut,
 			ExtraEnv: env,
 			Args:     bc.AllScriptArgs(),
@@ -814,7 +814,7 @@ func (b *Build) checkRelocations(client *buildlet.Client) error {
 	}
 	var out bytes.Buffer
 	file := fmt.Sprintf("go/pkg/linux_%s/runtime/cgo.a", b.Arch)
-	remoteErr, err := client.Exec("readelf", buildlet.ExecOpts{
+	remoteErr, err := client.Exec(context.Background(), "readelf", buildlet.ExecOpts{
 		Output:      &out,
 		Args:        []string{"-r", "--wide", file},
 		SystemLevel: true, // look for readelf in system's PATH

internal/buildgo/benchmarks.go

@@ -58,7 +58,7 @@ func buildGo1(conf *dashboard.BuildConfig, bc *buildlet.Client, goroot string, w
 	if found {
 		return nil, nil
 	}
-	return bc.Exec(path.Join(goroot, "bin", "go"), buildlet.ExecOpts{
+	return bc.Exec(context.TODO(), path.Join(goroot, "bin", "go"), buildlet.ExecOpts{
 		Output:   w,
 		ExtraEnv: []string{"GOROOT=" + conf.FilePathJoin(workDir, goroot)},
 		Args:     []string{"test", "-c"},
@@ -72,7 +72,7 @@ func buildPkg(conf *dashboard.BuildConfig, bc *buildlet.Client, goroot string, w
 	if err != nil {
 		return nil, err
 	}
-	return bc.Exec(path.Join(goroot, "bin", "go"), buildlet.ExecOpts{
+	return bc.Exec(context.TODO(), path.Join(goroot, "bin", "go"), buildlet.ExecOpts{
 		Output:   w,
 		ExtraEnv: []string{"GOROOT=" + conf.FilePathJoin(workDir, goroot)},
 		Args:     []string{"test", "-c", "-o", conf.FilePathJoin(workDir, goroot, name), pkg},
@@ -90,7 +90,7 @@ func buildXBenchmark(sl spanlog.Logger, conf *dashboard.BuildConfig, bc *buildle
 			return nil, err
 		}
 	}
-	return bc.Exec(path.Join(goroot, "bin/go"), buildlet.ExecOpts{
+	return bc.Exec(context.TODO(), path.Join(goroot, "bin/go"), buildlet.ExecOpts{
 		Output: w,
 		ExtraEnv: []string{
 			"GOROOT=" + conf.FilePathJoin(workDir, goroot),
@@ -134,7 +134,7 @@ func (gb GoBuilder) EnumerateBenchmarks(bc *buildlet.Client, benchmarksRev strin
 	// Enumerate x/benchmarks
 	if benchmarksRev != "" {
 		var buf bytes.Buffer
-		remoteErr, err := bc.Exec(path.Join(gb.Goroot, "bin/go"), buildlet.ExecOpts{
+		remoteErr, err := bc.Exec(context.TODO(), path.Join(gb.Goroot, "bin/go"), buildlet.ExecOpts{
 			Output: &buf,
 			ExtraEnv: []string{
 				"GOROOT=" + gb.Conf.FilePathJoin(workDir, gb.Goroot),
@@ -161,7 +161,7 @@ func (gb GoBuilder) EnumerateBenchmarks(bc *buildlet.Client, benchmarksRev strin
 	if len(pkgs) > 0 {
 		// Find packages that actually have benchmarks or tests.
 		var buf bytes.Buffer
-		remoteErr, err := bc.Exec(path.Join(gb.Goroot, "bin/go"), buildlet.ExecOpts{
+		remoteErr, err := bc.Exec(context.TODO(), path.Join(gb.Goroot, "bin/go"), buildlet.ExecOpts{
 			Output: &buf,
 			ExtraEnv: []string{
 				"GOROOT=" + gb.Conf.FilePathJoin(workDir, gb.Goroot),
@@ -205,7 +205,7 @@ func runOneBenchBinary(conf *dashboard.BuildConfig, bc *buildlet.Client, w io.Wr
 		return nil, fmt.Errorf("runOneBenchBinary, WorkDir: %v", err)
 	}
 	// Some benchmarks need GOROOT so they can invoke cmd/go.
-	return bc.Exec(binaryPath, buildlet.ExecOpts{
+	return bc.Exec(context.TODO(), binaryPath, buildlet.ExecOpts{
 		Output: w,
 		Dir:    dir,
 		Args:   args,
@@ -226,7 +226,7 @@ func runOneBenchBinary(conf *dashboard.BuildConfig, bc *buildlet.Client, w io.Wr
 	})
 }
 
-func buildRev(buildEnv *buildenv.Environment, sl spanlog.Logger, conf *dashboard.BuildConfig, bc *buildlet.Client, w io.Writer, goroot string, br BuilderRev) error {
+func buildRev(ctx context.Context, buildEnv *buildenv.Environment, sl spanlog.Logger, conf *dashboard.BuildConfig, bc *buildlet.Client, w io.Writer, goroot string, br BuilderRev) error {
 	if br.SnapshotExists(context.TODO(), buildEnv) {
 		return bc.PutTarFromURL(br.SnapshotURL(buildEnv), goroot)
 	}
@@ -246,7 +246,7 @@ func buildRev(buildEnv *buildenv.Environment, sl spanlog.Logger, conf *dashboard
 		Conf:       conf,
 		Goroot:     goroot,
 	}
-	remoteErr, err := builder.RunMake(bc, w)
+	remoteErr, err := builder.RunMake(ctx, bc, w)
 	if err != nil {
 		return err
 	}
@@ -257,12 +257,12 @@ func buildRev(buildEnv *buildenv.Environment, sl spanlog.Logger, conf *dashboard
 // Build output is sent to w. Benchmark output is stored in b.output.
 // revs must contain exactly two revs. The first rev is assumed to be present in "go", and the second will be placed into "go-parent".
 // TODO(quentin): Support len(revs) != 2.
-func (b *BenchmarkItem) Run(buildEnv *buildenv.Environment, sl spanlog.Logger, conf *dashboard.BuildConfig, bc *buildlet.Client, w io.Writer, revs []BuilderRev) (remoteErr, err error) {
+func (b *BenchmarkItem) Run(ctx context.Context, buildEnv *buildenv.Environment, sl spanlog.Logger, conf *dashboard.BuildConfig, bc *buildlet.Client, w io.Writer, revs []BuilderRev) (remoteErr, err error) {
 	// Ensure we have a built parent repo.
 	if err := bc.ListDir("go-parent", buildlet.ListDirOpts{}, func(buildlet.DirEntry) {}); err != nil {
 		pbr := revs[1]
 		sp := sl.CreateSpan("bench_build_parent", bc.Name())
-		err = buildRev(buildEnv, sl, conf, bc, w, "go-parent", pbr)
+		err = buildRev(ctx, buildEnv, sl, conf, bc, w, "go-parent", pbr)
 		sp.Done(err)
 		if err != nil {
 			return nil, err

internal/buildgo/buildgo.go

@@ -103,10 +103,10 @@ type GoBuilder struct {
 // goroot is relative to the workdir with forward slashes.
 // w is the Writer to send build output to.
 // remoteErr and err are as described at the top of this file.
-func (gb GoBuilder) RunMake(bc *buildlet.Client, w io.Writer) (remoteErr, err error) {
+func (gb GoBuilder) RunMake(ctx context.Context, bc *buildlet.Client, w io.Writer) (remoteErr, err error) {
 	// Build the source code.
 	makeSpan := gb.CreateSpan("make", gb.Conf.MakeScript())
-	remoteErr, err = bc.Exec(path.Join(gb.Goroot, gb.Conf.MakeScript()), buildlet.ExecOpts{
+	remoteErr, err = bc.Exec(ctx, path.Join(gb.Goroot, gb.Conf.MakeScript()), buildlet.ExecOpts{
 		Output:   w,
 		ExtraEnv: append(gb.Conf.Env(), "GOBIN="),
 		Debug:    true,
@@ -125,7 +125,7 @@ func (gb GoBuilder) RunMake(bc *buildlet.Client, w io.Writer) (remoteErr, err er
 	// Need to run "go install -race std" before the snapshot + tests.
 	if pkgs := gb.Conf.GoInstallRacePackages(); len(pkgs) > 0 {
 		sp := gb.CreateSpan("install_race_std")
-		remoteErr, err = bc.Exec(path.Join(gb.Goroot, "bin/go"), buildlet.ExecOpts{
+		remoteErr, err = bc.Exec(ctx, path.Join(gb.Goroot, "bin/go"), buildlet.ExecOpts{
 			Output:   w,
 			ExtraEnv: append(gb.Conf.Env(), "GOBIN="),
 			Debug:    true,
@@ -143,7 +143,7 @@ func (gb GoBuilder) RunMake(bc *buildlet.Client, w io.Writer) (remoteErr, err er
 	}
 
 	if gb.Name == "linux-amd64-racecompile" {
-		return gb.runConcurrentGoBuildStdCmd(bc, w)
+		return gb.runConcurrentGoBuildStdCmd(ctx, bc, w)
 	}
 
 	return nil, nil
@@ -155,9 +155,9 @@ func (gb GoBuilder) RunMake(bc *buildlet.Client, w io.Writer) (remoteErr, err er
 // with -gcflags=-c=8 using a race-enabled cmd/compile (built by
 // caller, runMake, per builder config).
 // The idea is that this might find data races in cmd/compile.
-func (gb GoBuilder) runConcurrentGoBuildStdCmd(bc *buildlet.Client, w io.Writer) (remoteErr, err error) {
+func (gb GoBuilder) runConcurrentGoBuildStdCmd(ctx context.Context, bc *buildlet.Client, w io.Writer) (remoteErr, err error) {
 	span := gb.CreateSpan("go_build_c128_std_cmd")
-	remoteErr, err = bc.Exec(path.Join(gb.Goroot, "bin/go"), buildlet.ExecOpts{
+	remoteErr, err = bc.Exec(ctx, path.Join(gb.Goroot, "bin/go"), buildlet.ExecOpts{
 		Output:   w,
 		ExtraEnv: append(gb.Conf.Env(), "GOBIN="),
 		Debug:    true,