[exec] Respect context.Done()

- Use context.WithTimeout() to implement timeouts.
- Kill the subprocess if the context was canceled, either due to the
  configured timeout or parent context cancellation.
- Add util.NonCancelableContext() which wraps a parent context but
  ignores parent.Done() et al.

Change-Id: I6102961c3e0446cc123a4d1b61a87441f51b3667
Reviewed-on: https://skia-review.googlesource.com/c/buildbot/+/234036
Commit-Queue: Eric Boren <borenet@google.com>
Reviewed-by: Ben Wagner aka dogben <benjaminwagner@google.com>
diff --git a/autoroll/go/repo_manager/android_repo_manager_test.go b/autoroll/go/repo_manager/android_repo_manager_test.go
index d74c02a..0dd4dec 100644
--- a/autoroll/go/repo_manager/android_repo_manager_test.go
+++ b/autoroll/go/repo_manager/android_repo_manager_test.go
@@ -54,7 +54,7 @@
 	wd, err := ioutil.TempDir("", "")
 	assert.NoError(t, err)
 	mockRun := exec.CommandCollector{}
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		if strings.Contains(cmd.Name, "repo") {
 			return nil
 		}
diff --git a/autoroll/go/repo_manager/copy_repo_manager_test.go b/autoroll/go/repo_manager/copy_repo_manager_test.go
index 702d86a..726d0f0 100644
--- a/autoroll/go/repo_manager/copy_repo_manager_test.go
+++ b/autoroll/go/repo_manager/copy_repo_manager_test.go
@@ -53,7 +53,7 @@
 	parent.Commit(ctx)
 
 	mockRun := &exec.CommandCollector{}
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		if cmd.Name == "git" && cmd.Args[0] == "cl" {
 			if cmd.Args[1] == "upload" {
 				return nil
@@ -67,7 +67,7 @@
 				return nil
 			}
 		}
-		return exec.DefaultRun(cmd)
+		return exec.DefaultRun(ctx, cmd)
 	})
 	ctx = exec.NewContext(ctx, mockRun.Run)
 
diff --git a/autoroll/go/repo_manager/deps_repo_manager_test.go b/autoroll/go/repo_manager/deps_repo_manager_test.go
index c024148..9c01a9e 100644
--- a/autoroll/go/repo_manager/deps_repo_manager_test.go
+++ b/autoroll/go/repo_manager/deps_repo_manager_test.go
@@ -74,7 +74,7 @@
 	lastUpload := new(vcsinfo.LongCommit)
 	mockRun := &exec.CommandCollector{}
 	ctx := exec.NewContext(context.Background(), mockRun.Run)
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		if cmd.Name == "git" && cmd.Args[0] == "cl" {
 			if cmd.Args[1] == "upload" {
 				d, err := git.GitDir(cmd.Dir).Details(ctx, "HEAD")
@@ -98,7 +98,7 @@
 			assert.Equal(t, 2, len(splitDep))
 			assert.Equal(t, 40, len(splitDep[1]))
 		}
-		return exec.DefaultRun(cmd)
+		return exec.DefaultRun(ctx, cmd)
 	})
 
 	cleanup := func() {
diff --git a/autoroll/go/repo_manager/fuchsia_sdk_android_repo_manager_test.go b/autoroll/go/repo_manager/fuchsia_sdk_android_repo_manager_test.go
index f7d934c..1431aff 100644
--- a/autoroll/go/repo_manager/fuchsia_sdk_android_repo_manager_test.go
+++ b/autoroll/go/repo_manager/fuchsia_sdk_android_repo_manager_test.go
@@ -46,7 +46,7 @@
 
 	// Mock out repo commands.
 	mockRun := exec.CommandCollector{}
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		if strings.Contains(cmd.Name, "repo") {
 			return nil
 		} else if cmd.Name == "git" && strings.Contains(cmd.Dir, cfg.ChildPath) {
@@ -77,7 +77,7 @@
 			assert.NoError(t, ioutil.WriteFile(androidBp, []byte("hi"), os.ModePerm))
 			return nil
 		} else {
-			return exec.DefaultRun(cmd)
+			return exec.DefaultRun(ctx, cmd)
 		}
 	})
 	ctx := exec.NewContext(context.Background(), mockRun.Run)
diff --git a/autoroll/go/repo_manager/github_cipd_deps_repo_manager_test.go b/autoroll/go/repo_manager/github_cipd_deps_repo_manager_test.go
index 2621aea..fe52464 100644
--- a/autoroll/go/repo_manager/github_cipd_deps_repo_manager_test.go
+++ b/autoroll/go/repo_manager/github_cipd_deps_repo_manager_test.go
@@ -73,7 +73,7 @@
 	parent.Commit(context.Background())
 
 	mockRun := &exec.CommandCollector{}
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		if cmd.Name == "git" {
 			if cmd.Args[0] == "clone" || cmd.Args[0] == "fetch" || cmd.Args[0] == "reset" {
 				return nil
@@ -83,7 +83,7 @@
 				cmd.Args[1] = "origin/master"
 			}
 		}
-		return exec.DefaultRun(cmd)
+		return exec.DefaultRun(ctx, cmd)
 	})
 	ctx := exec.NewContext(context.Background(), mockRun.Run)
 
diff --git a/autoroll/go/repo_manager/github_deps_repo_manager_test.go b/autoroll/go/repo_manager/github_deps_repo_manager_test.go
index b3258d9..d9bfad8 100644
--- a/autoroll/go/repo_manager/github_deps_repo_manager_test.go
+++ b/autoroll/go/repo_manager/github_deps_repo_manager_test.go
@@ -88,13 +88,13 @@
 	parent.Commit(context.Background())
 
 	mockRun := &exec.CommandCollector{}
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		// Without this, the mock commands get confused with:
 		// "Could not switch upstream branch from refs/remotes/remote/master to refs/remotes/origin/master"
 		if strings.Contains(cmd.Name, "gclient") && (cmd.Args[0] == "sync" || cmd.Args[0] == "runhooks") {
 			return nil
 		}
-		return exec.DefaultRun(cmd)
+		return exec.DefaultRun(ctx, cmd)
 	})
 	ctx := exec.NewContext(context.Background(), mockRun.Run)
 
diff --git a/autoroll/go/repo_manager/github_repo_manager_test.go b/autoroll/go/repo_manager/github_repo_manager_test.go
index 0dfbac2..79aa61d 100644
--- a/autoroll/go/repo_manager/github_repo_manager_test.go
+++ b/autoroll/go/repo_manager/github_repo_manager_test.go
@@ -81,7 +81,7 @@
 	parent.Commit(context.Background())
 
 	mockRun := &exec.CommandCollector{}
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		if cmd.Name == "git" {
 			if cmd.Args[0] == "clone" || cmd.Args[0] == "fetch" {
 				return nil
@@ -91,7 +91,7 @@
 				cmd.Args[1] = "origin/master"
 			}
 		}
-		return exec.DefaultRun(cmd)
+		return exec.DefaultRun(ctx, cmd)
 	})
 	ctx := exec.NewContext(context.Background(), mockRun.Run)
 
diff --git a/autoroll/go/repo_manager/pre_upload_steps_test.go b/autoroll/go/repo_manager/pre_upload_steps_test.go
index 5d1f9fe..415ddc3 100644
--- a/autoroll/go/repo_manager/pre_upload_steps_test.go
+++ b/autoroll/go/repo_manager/pre_upload_steps_test.go
@@ -38,7 +38,7 @@
 	gitErr := error(nil)
 
 	mockRun := &exec.CommandCollector{}
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		pubCmd := "get"
 		dartCmd := "lib/main.dart --src ../../.. --out testing/out/licenses --golden testing/dir/ci/licenses_golden"
 		if cmd.Name == "testing/third_party/dart/tools/sdks/dart-sdk/bin/pub" && strings.Join(cmd.Args, " ") == pubCmd {
@@ -52,7 +52,7 @@
 				return gitErr
 			}
 		}
-		return exec.DefaultRun(cmd)
+		return exec.DefaultRun(ctx, cmd)
 	})
 	ctx := exec.NewContext(context.Background(), mockRun.Run)
 
diff --git a/ct/go/poller/poller_test.go b/ct/go/poller/poller_test.go
index c9bf7a6..c1bd3c2 100644
--- a/ct/go/poller/poller_test.go
+++ b/ct/go/poller/poller_test.go
@@ -88,7 +88,7 @@
 func TestChromiumPerfExecute(t *testing.T) {
 	unittest.SmallTest(t)
 	mockRun := exec.CommandCollector{}
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		runId := getRunId(t, cmd)
 		assertFileContents(t, filepath.Join(os.TempDir(), runId+".chromium.patch"),
 			"patches/abc.patch\n")
@@ -236,7 +236,7 @@
 	mockRun := exec.CommandCollector{}
 	ctx := exec.NewContext(context.Background(), mockRun.Run)
 	task := pendingLuaScriptTaskWithAggregator(ctx)
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		runId := getRunId(t, cmd)
 		assertFileContents(t, filepath.Join(os.TempDir(), runId+".lua"),
 			`print("lualualua")`)
@@ -270,7 +270,7 @@
 func TestLuaScriptExecuteWithoutAggregator(t *testing.T) {
 	unittest.SmallTest(t)
 	mockRun := exec.CommandCollector{}
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		runId := getRunId(t, cmd)
 		assertFileContents(t, filepath.Join(os.TempDir(), runId+".lua"),
 			`print("lualualua")`)
diff --git a/go/exec/exec.go b/go/exec/exec.go
index a3aa434..0883cd3 100644
--- a/go/exec/exec.go
+++ b/go/exec/exec.go
@@ -229,24 +229,19 @@
 	return nil
 }
 
-func waitSimple(command *Command, cmd *osexec.Cmd) error {
-	err := cmd.Wait()
-	if err != nil {
-		return fmt.Errorf("Command exited with %s: %s", err, DebugString(command))
-	}
-	return nil
-}
-
-func wait(command *Command, cmd *osexec.Cmd) error {
-	if command.Timeout == 0 {
-		return waitSimple(command, cmd)
+func wait(ctx context.Context, command *Command, cmd *osexec.Cmd) error {
+	if command.Timeout > 0 {
+		var cancel func()
+		ctx, cancel = context.WithTimeout(ctx, command.Timeout)
+		defer cancel()
 	}
 	done := make(chan error)
 	go func() {
 		done <- cmd.Wait()
 	}()
+	canceled := ctx.Done()
 	select {
-	case <-time.After(command.Timeout):
+	case <-canceled:
 		if command.Verbose != Silent {
 			sklog.Debugf("About to kill command '%s'", DebugString(command))
 		}
@@ -273,22 +268,22 @@
 }
 
 // DefaultRun can be passed to SetRunForTesting to go back to running commands as normal.
-func DefaultRun(command *Command) error {
+func DefaultRun(ctx context.Context, command *Command) error {
 	cmd := createCmd(command)
 	if err := start(command, cmd); err != nil {
 		return err
 	}
-	return wait(command, cmd)
+	return wait(ctx, command, cmd)
 }
 
 // execContext is a struct used for controlling the execution context of Commands.
 type execContext struct {
-	runFn func(*Command) error
+	runFn func(context.Context, *Command) error
 }
 
 // NewContext returns a context.Context instance which uses the given function
 // to run Commands.
-func NewContext(ctx context.Context, runFn func(*Command) error) context.Context {
+func NewContext(ctx context.Context, runFn func(context.Context, *Command) error) context.Context {
 	newCtx := &execContext{
 		runFn: runFn,
 	}
@@ -304,13 +299,13 @@
 }
 
 // See documentation for exec.Run.
-func (c *execContext) Run(command *Command) error {
-	return c.runFn(command)
+func (c *execContext) Run(ctx context.Context, command *Command) error {
+	return c.runFn(ctx, command)
 }
 
 // runSimpleCommand executes the given command.  Returns the combined stdout and stderr. May also
 // return an error if the command exited with a non-zero status or there is any other error.
-func (c *execContext) runSimpleCommand(command *Command) (string, error) {
+func (c *execContext) runSimpleCommand(ctx context.Context, command *Command) (string, error) {
 	output := bytes.Buffer{}
 	// We use a ThreadSafeWriter here because command.CombinedOutput may get
 	// wrapped with an io.MultiWriter if the caller set command.Stdout or
@@ -324,7 +319,7 @@
 	command.CombinedOutput = util.NewThreadSafeWriter(&output)
 	// Setting Verbose to Silent to maintain previous behavior.
 	command.Verbose = Silent
-	err := c.Run(command)
+	err := c.Run(ctx, command)
 	result := string(output.Bytes())
 	if err != nil {
 		return result, fmt.Errorf("%s; Stdout+Stderr:\n%s", err.Error(), result)
@@ -333,50 +328,50 @@
 }
 
 // See documentation for exec.RunSimple.
-func (c *execContext) RunSimple(commandLine string) (string, error) {
+func (c *execContext) RunSimple(ctx context.Context, commandLine string) (string, error) {
 	cmd := ParseCommand(commandLine)
-	return c.runSimpleCommand(&cmd)
+	return c.runSimpleCommand(ctx, &cmd)
 }
 
 // See documentation for exec.RunCommand.
-func (c *execContext) RunCommand(command *Command) (string, error) {
-	return c.runSimpleCommand(command)
+func (c *execContext) RunCommand(ctx context.Context, command *Command) (string, error) {
+	return c.runSimpleCommand(ctx, command)
 }
 
 // See documentation for exec.RunCwd.
-func (c *execContext) RunCwd(cwd string, args ...string) (string, error) {
+func (c *execContext) RunCwd(ctx context.Context, cwd string, args ...string) (string, error) {
 	command := &Command{
 		Name: args[0],
 		Args: args[1:],
 		Dir:  cwd,
 	}
-	return c.runSimpleCommand(command)
+	return c.runSimpleCommand(ctx, command)
 }
 
 // Run runs command and waits for it to finish. If any failure, returns non-nil. If a timeout was
 // specified, returns an error once the command has exceeded that timeout.
 func Run(ctx context.Context, command *Command) error {
-	return getCtx(ctx).Run(command)
+	return getCtx(ctx).Run(ctx, command)
 }
 
 // RunSimple executes the given command line string; the command being run is expected to not care
 // what its current working directory is. Returns the combined stdout and stderr. May also return
 // an error if the command exited with a non-zero status or there is any other error.
 func RunSimple(ctx context.Context, commandLine string) (string, error) {
-	return getCtx(ctx).RunSimple(commandLine)
+	return getCtx(ctx).RunSimple(ctx, commandLine)
 }
 
 // RunCommand executes the given command and returns the combined stdout and stderr. May also
 // return an error if the command exited with a non-zero status or there is any other error.
 func RunCommand(ctx context.Context, command *Command) (string, error) {
-	return getCtx(ctx).runSimpleCommand(command)
+	return getCtx(ctx).runSimpleCommand(ctx, command)
 }
 
 // RunCwd executes the given command in the given directory. Returns the combined stdout and
 // stderr. May also return an error if the command exited with a non-zero status or there is any
 // other error.
 func RunCwd(ctx context.Context, cwd string, args ...string) (string, error) {
-	return getCtx(ctx).RunCwd(cwd, args...)
+	return getCtx(ctx).RunCwd(ctx, cwd, args...)
 }
 
 // RunIndefinitely starts the command and then returns. Clients can listen for
diff --git a/go/exec/exec_linux.go b/go/exec/exec_linux.go
index a018ffc..4f93e05 100644
--- a/go/exec/exec_linux.go
+++ b/go/exec/exec_linux.go
@@ -3,20 +3,22 @@
 import (
 	"context"
 	"syscall"
+
+	"go.skia.org/infra/go/util"
 )
 
 // NoInterruptContext returns a context.Context instance which launches
 // subprocesses in a difference process group so that they are not killed when
 // this process is killed.
 //
-// This function is a no-op on Windows.
+// On Windows, this function just returns util.WithoutCancel(ctx).
 func NoInterruptContext(ctx context.Context) context.Context {
 	parent := getCtx(ctx)
-	runFn := func(c *Command) error {
+	runFn := func(ctx context.Context, c *Command) error {
 		c.SysProcAttr = &syscall.SysProcAttr{
 			Setpgid: true,
 		}
-		return parent.runFn(c)
+		return parent.runFn(ctx, c)
 	}
-	return NewContext(ctx, runFn)
+	return NewContext(util.WithoutCancel(ctx), runFn)
 }
diff --git a/go/exec/exec_test.go b/go/exec/exec_test.go
index 2b0d132..50181e0 100644
--- a/go/exec/exec_test.go
+++ b/go/exec/exec_test.go
@@ -328,7 +328,7 @@
 func TestInjection(t *testing.T) {
 	unittest.SmallTest(t)
 	var actualCommand *Command
-	ctx := NewContext(context.Background(), func(command *Command) error {
+	ctx := NewContext(context.Background(), func(ctx context.Context, command *Command) error {
 		actualCommand = command
 		return nil
 	})
diff --git a/go/exec/exec_testutil.go b/go/exec/exec_testutil.go
index 839530f..4e4fdbd 100644
--- a/go/exec/exec_testutil.go
+++ b/go/exec/exec_testutil.go
@@ -4,6 +4,7 @@
 package exec
 
 import (
+	"context"
 	"regexp"
 	"sync"
 )
@@ -22,7 +23,7 @@
 type CommandCollector struct {
 	mutex       sync.RWMutex
 	commands    []*Command
-	delegateRun func(*Command) error
+	delegateRun func(context.Context, *Command) error
 }
 
 func (c *CommandCollector) Commands() []*Command {
@@ -40,7 +41,7 @@
 	c.commands = nil
 }
 
-func (c *CommandCollector) SetDelegateRun(delegateRun func(*Command) error) {
+func (c *CommandCollector) SetDelegateRun(delegateRun func(context.Context, *Command) error) {
 	c.mutex.Lock()
 	defer c.mutex.Unlock()
 	c.delegateRun = delegateRun
@@ -49,7 +50,7 @@
 // Collects command into c and delegates to the function specified by SetDelegateRun. Returns nil
 // if SetDelegateRun has not been called. The command will be visible in Commands() before the
 // SetDelegateRun function is called.
-func (c *CommandCollector) Run(command *Command) error {
+func (c *CommandCollector) Run(ctx context.Context, command *Command) error {
 	c.mutex.Lock()
 	c.commands = append(c.commands, command)
 	delegateRun := c.delegateRun
@@ -57,7 +58,7 @@
 	if delegateRun == nil {
 		return nil
 	} else {
-		return delegateRun(command)
+		return delegateRun(ctx, command)
 	}
 }
 
@@ -94,7 +95,7 @@
 
 // Tries to match DebugString(command) against the regexps in the order of the calls to AddRule,
 // with the first matched giving the return value. Returns nil if no regexps match.
-func (m *MockRun) Run(command *Command) error {
+func (m *MockRun) Run(ctx context.Context, command *Command) error {
 	m.mutex.RLock()
 	defer m.mutex.RUnlock()
 	commandStr := DebugString(command)
diff --git a/go/exec/exec_windows.go b/go/exec/exec_windows.go
index 637b0cb..ef6c621 100644
--- a/go/exec/exec_windows.go
+++ b/go/exec/exec_windows.go
@@ -1,12 +1,16 @@
 package exec
 
-import "context"
+import (
+	"context"
+
+	"go.skia.org/infra/go/util"
+)
 
 // NoInterruptContext returns a context.Context instance which launches
 // subprocesses in a difference process group so that they are not killed when
 // this process is killed.
 //
-// This function is a no-op on Windows.
+// On Windows, this function just returns util.WithoutCancel(ctx).
 func NoInterruptContext(ctx context.Context) context.Context {
-	return ctx
+	return util.WithoutCancel(ctx)
 }
diff --git a/go/util/context.go b/go/util/context.go
new file mode 100644
index 0000000..cfd5137
--- /dev/null
+++ b/go/util/context.go
@@ -0,0 +1,33 @@
+package util
+
+import (
+	"context"
+	"time"
+)
+
+// withoutCancelContext is a context.Context implementation which is not
+// cancelable, even if the parent context is canceled.
+type withoutCancelContext struct {
+	context.Context
+}
+
+// See documentation for context.Context interface.
+func (ctx *withoutCancelContext) Deadline() (time.Time, bool) {
+	return time.Time{}, false
+}
+
+// See documentation for context.Context interface.
+func (ctx *withoutCancelContext) Done() <-chan struct{} {
+	return nil
+}
+
+// See documentation for context.Context interface.
+func (ctx *withoutCancelContext) Err() error {
+	return nil
+}
+
+// WithoutCancel returns a context.Context which cannot be canceled, even
+// if its parent is canceled.
+func WithoutCancel(ctx context.Context) context.Context {
+	return &withoutCancelContext{ctx}
+}
diff --git a/task_driver/go/lib/golang/golang_test.go b/task_driver/go/lib/golang/golang_test.go
index e3d5357..9513d8f 100644
--- a/task_driver/go/lib/golang/golang_test.go
+++ b/task_driver/go/lib/golang/golang_test.go
@@ -26,7 +26,7 @@
 		ctx = WithEnv(ctx, wd)
 		mockRun := &exec.CommandCollector{}
 		runCount := 0
-		mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+		mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 			runCount++
 
 			// Misc variables.
diff --git a/task_driver/go/td/context.go b/task_driver/go/td/context.go
index 93a965f..e5bd652 100644
--- a/task_driver/go/td/context.go
+++ b/task_driver/go/td/context.go
@@ -27,7 +27,7 @@
 
 	// execRun provides a Run function to be called by execCtx. This is used
 	// for testing, where we may want to mock out subprocess invocations.
-	execRun func(*exec.Command) error
+	execRun func(context.Context, *exec.Command) error
 }
 
 // getCtx retrieves the current Context. Panics if none exists.
@@ -71,7 +71,7 @@
 }
 
 // WithExecRunFn allows the Run function to be overridden for testing.
-func WithExecRunFn(ctx context.Context, run func(*exec.Command) error) context.Context {
+func WithExecRunFn(ctx context.Context, run func(context.Context, *exec.Command) error) context.Context {
 	return withChildCtx(ctx, &Context{
 		execRun: run,
 	})
diff --git a/task_driver/go/td/step.go b/task_driver/go/td/step.go
index 909186a..e62a02b 100644
--- a/task_driver/go/td/step.go
+++ b/task_driver/go/td/step.go
@@ -371,7 +371,7 @@
 // Return a context.Context associated with this Step. Any calls to exec which
 // use this Context will be attached to the Step.
 func execCtx(ctx context.Context) context.Context {
-	return exec.NewContext(ctx, func(cmd *exec.Command) error {
+	return exec.NewContext(ctx, func(ctx context.Context, cmd *exec.Command) error {
 		name := strings.Join(append([]string{cmd.Name}, cmd.Args...), " ")
 
 		// Merge the command's env into that of its parent.
@@ -398,7 +398,7 @@
 			StepData(ctx, DATA_TYPE_COMMAND, d)
 
 			// Run the command.
-			return getCtx(ctx).execRun(cmd)
+			return getCtx(ctx).execRun(ctx, cmd)
 		})
 	})
 }
diff --git a/task_driver/go/td/step_test.go b/task_driver/go/td/step_test.go
index e2153f4..18f7ce9 100644
--- a/task_driver/go/td/step_test.go
+++ b/task_driver/go/td/step_test.go
@@ -21,7 +21,7 @@
 func mockExec(ctx context.Context) (context.Context, *int) {
 	mockRun := &exec.CommandCollector{}
 	runCount := 0
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		runCount++
 		if cmd.Name == "true" {
 			return nil
@@ -321,7 +321,7 @@
 	expect := MergeEnv(os.Environ(), BASE_ENV)
 	expect = append(expect, "a=a", "b=b", "c=c", "d=d")
 	mockRun := &exec.CommandCollector{}
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		runCount++
 		assert.Equal(t, expect, cmd.Env)
 		return nil
diff --git a/task_scheduler/go/syncer/syncer_test.go b/task_scheduler/go/syncer/syncer_test.go
index f38d85f..d0f40e6 100644
--- a/task_scheduler/go/syncer/syncer_test.go
+++ b/task_scheduler/go/syncer/syncer_test.go
@@ -216,13 +216,13 @@
 
 	botUpdateCount := 0
 	mockRun := exec.CommandCollector{}
-	mockRun.SetDelegateRun(func(cmd *exec.Command) error {
+	mockRun.SetDelegateRun(func(ctx context.Context, cmd *exec.Command) error {
 		for _, arg := range cmd.Args {
 			if strings.Contains(arg, "bot_update") {
 				botUpdateCount++
 			}
 		}
-		return exec.DefaultRun(cmd)
+		return exec.DefaultRun(ctx, cmd)
 	})
 	ctx = exec.NewContext(context.Background(), mockRun.Run)