buildlet: add context argument to most of the remaining Client methods

(And update all the callers in the tree.)

Updates golang/go#35707

Change-Id: I54769bbe374f31ae1dd07776b27818db91ce8c70
Reviewed-on: https://go-review.googlesource.com/c/build/+/208157
Run-TryBot: Brad Fitzpatrick <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>
Reviewed-by: Dmitri Shuralyov <[email protected]>

Author: bradfitz

buildlet/buildletclient.go

@@ -323,7 +323,7 @@ func (c *Client) heartbeatLoop() {
 			return
 		case <-time.After(10 * time.Second):
 			t0 := time.Now()
-			if _, err := c.Status(); err != nil {
+			if _, err := c.Status(context.Background()); err != nil {
 				failInARow++
 				if failInARow == 3 {
 					log.Printf("Buildlet %v failed three heartbeats; final error: %v", c, err)
@@ -406,20 +406,20 @@ func (c *Client) doOK(req *http.Request) error {
 // If dir is empty, they're placed at the root of the buildlet's work directory.
 // The dir is created if necessary.
 // The Reader must be of a tar.gz file.
-func (c *Client) PutTar(r io.Reader, dir string) error {
+func (c *Client) PutTar(ctx context.Context, r io.Reader, dir string) error {
 	req, err := http.NewRequest("PUT", c.URL()+"/writetgz?dir="+url.QueryEscape(dir), r)
 	if err != nil {
 		return err
 	}
-	return c.doOK(req)
+	return c.doOK(req.WithContext(ctx))
 }
 
 // PutTarFromURL tells the buildlet to download the tar.gz file from tarURL
 // and write it to dir, a relative directory from the workdir.
 // If dir is empty, they're placed at the root of the buildlet's work directory.
 // The dir is created if necessary.
 // The url must be of a tar.gz file.
-func (c *Client) PutTarFromURL(tarURL, dir string) error {
+func (c *Client) PutTarFromURL(ctx context.Context, tarURL, dir string) error {
 	form := url.Values{
 		"url": {tarURL},
 	}
@@ -428,11 +428,11 @@ func (c *Client) PutTarFromURL(tarURL, dir string) error {
 		return err
 	}
 	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
-	return c.doOK(req)
+	return c.doOK(req.WithContext(ctx))
 }
 
 // Put writes the provided file to path (relative to workdir) and sets mode.
-func (c *Client) Put(r io.Reader, path string, mode os.FileMode) error {
+func (c *Client) Put(ctx context.Context, r io.Reader, path string, mode os.FileMode) error {
 	param := url.Values{
 		"path": {path},
 		"mode": {fmt.Sprint(int64(mode))},
@@ -441,7 +441,7 @@ func (c *Client) Put(r io.Reader, path string, mode os.FileMode) error {
 	if err != nil {
 		return err
 	}
-	return c.doOK(req)
+	return c.doOK(req.WithContext(ctx))
 }
 
 // GetTar returns a .tar.gz stream of the given directory, relative to the buildlet's work dir.
@@ -614,7 +614,7 @@ func (c *Client) Exec(ctx context.Context, cmd string, opts ExecOpts) (remoteErr
 }
 
 // RemoveAll deletes the provided paths, relative to the work directory.
-func (c *Client) RemoveAll(paths ...string) error {
+func (c *Client) RemoveAll(ctx context.Context, paths ...string) error {
 	if len(paths) == 0 {
 		return nil
 	}
@@ -624,7 +624,7 @@ func (c *Client) RemoveAll(paths ...string) error {
 		return err
 	}
 	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
-	return c.doOK(req)
+	return c.doOK(req.WithContext(ctx))
 }
 
 // DestroyVM shuts down the buildlet and destroys the VM instance.
@@ -681,7 +681,7 @@ type Status struct {
 }
 
 // Status returns an Status value describing this buildlet.
-func (c *Client) Status() (Status, error) {
+func (c *Client) Status(ctx context.Context) (Status, error) {
 	select {
 	case <-c.peerDead:
 		return Status{}, c.deadErr
@@ -692,6 +692,7 @@ func (c *Client) Status() (Status, error) {
 	if err != nil {
 		return Status{}, err
 	}
+	req = req.WithContext(ctx)
 	resp, err := c.doHeaderTimeout(req, 10*time.Second) // plenty of time
 	if err != nil {
 		return Status{}, err
@@ -712,11 +713,12 @@ func (c *Client) Status() (Status, error) {
 }
 
 // WorkDir returns the absolute path to the buildlet work directory.
-func (c *Client) WorkDir() (string, error) {
+func (c *Client) WorkDir(ctx context.Context) (string, error) {
 	req, err := http.NewRequest("GET", c.URL()+"/workdir", nil)
 	if err != nil {
 		return "", err
 	}
+	req = req.WithContext(ctx)
 	resp, err := c.doHeaderTimeout(req, 10*time.Second) // plenty of time
 	if err != nil {
 		return "", err
@@ -800,7 +802,7 @@ type ListDirOpts struct {
 // ListDir lists the contents of a directory.
 // The fn callback is run for each entry.
 // The directory dir itself is not included.
-func (c *Client) ListDir(dir string, opts ListDirOpts, fn func(DirEntry)) error {
+func (c *Client) ListDir(ctx context.Context, dir string, opts ListDirOpts, fn func(DirEntry)) error {
 	param := url.Values{
 		"dir":       {dir},
 		"recursive": {fmt.Sprint(opts.Recursive)},
@@ -811,7 +813,7 @@ func (c *Client) ListDir(dir string, opts ListDirOpts, fn func(DirEntry)) error
 	if err != nil {
 		return err
 	}
-	resp, err := c.do(req)
+	resp, err := c.do(req.WithContext(ctx))
 	if err != nil {
 		return err
 	}

cmd/coordinator/coordinator.go

@@ -1866,7 +1866,7 @@ func (st *buildStatus) build() error {
 
 	if st.useSnapshot() {
 		sp := st.CreateSpan("write_snapshot_tar")
-		if err := bc.PutTarFromURL(st.SnapshotURL(buildEnv), "go"); err != nil {
+		if err := bc.PutTarFromURL(st.ctx, st.SnapshotURL(buildEnv), "go"); err != nil {
 			return sp.Done(fmt.Errorf("failed to put snapshot to buildlet: %v", err))
 		}
 		sp.Done(nil)
@@ -2217,7 +2217,7 @@ func (st *buildStatus) writeGoSource() error {
 func (st *buildStatus) writeGoSourceTo(bc *buildlet.Client) error {
 	// Write the VERSION file.
 	sp := st.CreateSpan("write_version_tar")
-	if err := bc.PutTar(buildgo.VersionTgz(st.Rev), "go"); err != nil {
+	if err := bc.PutTar(st.ctx, buildgo.VersionTgz(st.Rev), "go"); err != nil {
 		return sp.Done(fmt.Errorf("writing VERSION tgz: %v", err))
 	}
 
@@ -2226,7 +2226,7 @@ func (st *buildStatus) writeGoSourceTo(bc *buildlet.Client) error {
 		return err
 	}
 	sp = st.CreateSpan("write_go_src_tar")
-	if err := bc.PutTar(srcTar, "go"); err != nil {
+	if err := bc.PutTar(st.ctx, srcTar, "go"); err != nil {
 		return sp.Done(fmt.Errorf("writing tarball from Gerrit: %v", err))
 	}
 	return sp.Done(nil)
@@ -2239,12 +2239,12 @@ func (st *buildStatus) writeBootstrapToolchain() error {
 	}
 	const bootstrapDir = "go1.4" // might be newer; name is the default
 	sp := st.CreateSpan("write_go_bootstrap_tar")
-	return sp.Done(st.bc.PutTarFromURL(u, bootstrapDir))
+	return sp.Done(st.bc.PutTarFromURL(st.ctx, u, bootstrapDir))
 }
 
 func (st *buildStatus) cleanForSnapshot(bc *buildlet.Client) error {
 	sp := st.CreateSpan("clean_for_snapshot")
-	return sp.Done(bc.RemoveAll(
+	return sp.Done(bc.RemoveAll(st.ctx,
 		"go/doc/gopher",
 		"go/pkg/bootstrap",
 	))
@@ -2257,7 +2257,7 @@ func (st *buildStatus) writeSnapshot(bc *buildlet.Client) (err error) {
 	// a couple times at 1 minute. Some buildlets might be far
 	// away on the network, so be more lenient. The timeout mostly
 	// is here to prevent infinite hangs.
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
+	ctx, cancel := context.WithTimeout(st.ctx, 5*time.Minute)
 	defer cancel()
 
 	tsp := st.CreateSpan("fetch_snapshot_reader_from_buildlet")
@@ -2295,7 +2295,7 @@ func (st *buildStatus) reportErr(err error) {
 }
 
 func (st *buildStatus) distTestList() (names []string, remoteErr, err error) {
-	workDir, err := st.bc.WorkDir()
+	workDir, err := st.bc.WorkDir(st.ctx)
 	if err != nil {
 		err = fmt.Errorf("distTestList, WorkDir: %v", err)
 		return
@@ -2441,7 +2441,7 @@ func getTestStats(sl spanlog.Logger) *buildstats.TestStats {
 func (st *buildStatus) runSubrepoTests() (remoteErr, err error) {
 	st.LogEventTime("fetching_subrepo", st.SubName)
 
-	workDir, err := st.bc.WorkDir()
+	workDir, err := st.bc.WorkDir(st.ctx)
 	if err != nil {
 		err = fmt.Errorf("error discovering workdir for helper %s: %v", st.bc.IPPort(), err)
 		return nil, err
@@ -2451,7 +2451,7 @@ func (st *buildStatus) runSubrepoTests() (remoteErr, err error) {
 
 	// Check out the provided sub-repo to the buildlet's workspace.
 	// Need to do this first, so we can run go env GOMOD in it.
-	err = buildgo.FetchSubrepo(st, st.bc, st.SubName, st.SubRev)
+	err = buildgo.FetchSubrepo(st.ctx, st, st.bc, st.SubName, st.SubRev)
 	if err != nil {
 		return nil, err
 	}
@@ -2492,7 +2492,7 @@ func (st *buildStatus) runSubrepoTests() (remoteErr, err error) {
 		// Look for inner modules, in order to test them too. See golang.org/issue/32528.
 		repoPath := importPathOfRepo(st.SubName)
 		sp := st.CreateSpan("listing_subrepo_modules", st.SubName)
-		err = st.bc.ListDir("gopath/src/"+repoPath, buildlet.ListDirOpts{Recursive: true}, func(e buildlet.DirEntry) {
+		err = st.bc.ListDir(st.ctx, "gopath/src/"+repoPath, buildlet.ListDirOpts{Recursive: true}, func(e buildlet.DirEntry) {
 			goModFile := path.Base(e.Name()) == "go.mod" && !e.IsDir()
 			if !goModFile {
 				return
@@ -2645,7 +2645,7 @@ func (st *buildStatus) fetchDependenciesToGOPATHWorkspace(goroot, gopath string)
 	// fetch checks out the provided sub-repo to the buildlet's workspace.
 	fetch := func(repo, rev string) error {
 		fetched[repo] = true
-		return buildgo.FetchSubrepo(st, st.bc, repo, rev)
+		return buildgo.FetchSubrepo(st.ctx, st, st.bc, repo, rev)
 	}
 
 	// findDeps uses 'go list' on the checked out repo to find its
@@ -2855,7 +2855,7 @@ func (st *buildStatus) runTests(helpers <-chan *buildlet.Client) (remoteErr, err
 		if rev == "" {
 			rev = "master" // should happen rarely; ok if it does.
 		}
-		b, err := st.goBuilder().EnumerateBenchmarks(st.bc, rev, st.trySet.affectedPkgs())
+		b, err := st.goBuilder().EnumerateBenchmarks(st.ctx, st.bc, rev, st.trySet.affectedPkgs())
 		sp.Done(err)
 		if err == nil {
 			benches = b
@@ -2871,7 +2871,7 @@ func (st *buildStatus) runTests(helpers <-chan *buildlet.Client) (remoteErr, err
 	st.LogEventTime("starting_tests", fmt.Sprintf("%d tests", len(set.items)))
 	startTime := time.Now()
 
-	workDir, err := st.bc.WorkDir()
+	workDir, err := st.bc.WorkDir(st.ctx)
 	if err != nil {
 		return nil, fmt.Errorf("error discovering workdir for main buildlet, %s: %v", st.bc.Name(), err)
 	}
@@ -2916,11 +2916,11 @@ func (st *buildStatus) runTests(helpers <-chan *buildlet.Client) (remoteErr, err
 					defer st.LogEventTime("DEV_HELPER_SLEEP", bc.Name())
 				}
 				st.LogEventTime("got_empty_test_helper", bc.String())
-				if err := bc.PutTarFromURL(st.SnapshotURL(buildEnv), "go"); err != nil {
+				if err := bc.PutTarFromURL(st.ctx, st.SnapshotURL(buildEnv), "go"); err != nil {
 					log.Printf("failed to extract snapshot for helper %s: %v", bc.Name(), err)
 					return
 				}
-				workDir, err := bc.WorkDir()
+				workDir, err := bc.WorkDir(st.ctx)
 				if err != nil {
 					log.Printf("error discovering workdir for helper %s: %v", bc.Name(), err)
 					return

cmd/coordinator/remote.go

@@ -628,7 +628,7 @@ func handleIncomingSSHPostAuth(s ssh.Session) {
 			err = <-errc
 		}()
 	}
-	workDir, err := rb.buildlet.WorkDir()
+	workDir, err := rb.buildlet.WorkDir(ctx)
 	if err != nil {
 		fmt.Fprintf(s, "Error getting WorkDir: %v\n", err)
 		return

cmd/coordinator/reverse.go

@@ -250,7 +250,7 @@ func (p *reverseBuildletPool) healthCheckBuildlet(b *reverseBuildlet) bool {
 	b.inUseTime = time.Now()
 	res := make(chan error, 1)
 	go func() {
-		_, err := b.client.Status()
+		_, err := b.client.Status(context.Background())
 		res <- err
 	}()
 	p.mu.Unlock()
@@ -328,7 +328,7 @@ func (p *reverseBuildletPool) GetBuildlet(ctx context.Context, hostType string,
 func (p *reverseBuildletPool) cleanedBuildlet(b *buildlet.Client, lg logger) (*buildlet.Client, error) {
 	// Clean up any files from previous builds.
 	sp := lg.CreateSpan("clean_buildlet", b.String())
-	err := b.RemoveAll(".")
+	err := b.RemoveAll(context.Background(), ".")
 	sp.Done(err)
 	if err != nil {
 		b.Close()
@@ -598,7 +598,7 @@ func handleReverse(w http.ResponseWriter, r *http.Request) {
 		}
 	}()
 	tstatus := time.Now()
-	status, err := client.Status()
+	status, err := client.Status(context.Background())
 	if err != nil {
 		log.Printf("Reverse connection %s/%s for %s did not answer status after %v: %v",
 			hostname, r.RemoteAddr, hostType, time.Since(tstatus), err)

cmd/debugnewvm/debugnewvm.go

@@ -124,7 +124,7 @@ func main() {
 	if err != nil {
 		log.Fatalf("StartNewVM: %v", err)
 	}
-	dir, err := bc.WorkDir()
+	dir, err := bc.WorkDir(ctx)
 	log.Printf("WorkDir: %v, %v", dir, err)
 
 	if *sleepSec > 0 {
@@ -146,7 +146,7 @@ func main() {
 		if u := bconf.GoBootstrapURL(env); u != "" {
 			log.Printf("Pushing 'go1.4' Go bootstrap dir ...")
 			const bootstrapDir = "go1.4" // might be newer; name is the default
-			if err := bc.PutTarFromURL(u, bootstrapDir); err != nil {
+			if err := bc.PutTarFromURL(ctx, u, bootstrapDir); err != nil {
 				bc.Close()
 				log.Fatalf("Putting Go bootstrap: %v", err)
 			}
@@ -155,13 +155,13 @@ func main() {
 		// Push Go code
 		log.Printf("Pushing 'go' dir...")
 		goTarGz := "https://go.googlesource.com/go/+archive/" + *buildRev + ".tar.gz"
-		if err := bc.PutTarFromURL(goTarGz, "go"); err != nil {
+		if err := bc.PutTarFromURL(ctx, goTarGz, "go"); err != nil {
 			bc.Close()
 			log.Fatalf("Putting go code: %v", err)
 		}
 
 		// Push a synthetic VERSION file to prevent git usage:
-		if err := bc.PutTar(buildgo.VersionTgz(*buildRev), "go"); err != nil {
+		if err := bc.PutTar(ctx, buildgo.VersionTgz(*buildRev), "go"); err != nil {
 			bc.Close()
 			log.Fatalf("Putting VERSION file: %v", err)
 		}

cmd/gomote/ls.go

@@ -5,6 +5,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"os"
@@ -44,7 +45,7 @@ func ls(args []string) error {
 		Digest:    digest,
 		Skip:      strings.Split(skip, ","),
 	}
-	return bc.ListDir(dir, opts, func(bi buildlet.DirEntry) {
+	return bc.ListDir(context.Background(), dir, opts, func(bi buildlet.DirEntry) {
 		fmt.Fprintf(os.Stdout, "%s\n", bi)
 	})
 }

cmd/gomote/ping.go

@@ -5,6 +5,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"os"
@@ -29,13 +30,14 @@ func ping(args []string) error {
 	if err != nil {
 		return err
 	}
-	wd, err := bc.WorkDir()
+	ctx := context.Background()
+	wd, err := bc.WorkDir(ctx)
 	if err != nil {
 		return err
 	}
 	if status {
 		fmt.Printf("workdir: %v\n", wd)
-		s, err := bc.Status()
+		s, err := bc.Status(ctx)
 		if err != nil {
 			return err
 		}