cmd/go: bail out from script tests earlier when a timeout occurs

(I needed this to debug an accidental infinite-recursion I introduced
while revising CL 293689.)

Fixes #38768

Change-Id: I306122f15b5bbd2fc5e836b32fd4dd5992ea891e
Reviewed-on: https://go-review.googlesource.com/c/go/+/302052
Trust: Bryan C. Mills <bcmills@google.com>
Reviewed-by: Michael Matloob <matloob@golang.org>
Run-TryBot: Bryan C. Mills <bcmills@google.com>
TryBot-Result: Go Bot <gobot@golang.org>

1 files changed, 71 insertions, 61 deletions

diff --git a/src/cmd/go/script_test.go b/src/cmd/go/script_test.go
index dfaa40548e..1f38be8ee4 100644
--- a/src/cmd/go/script_test.go
+++ b/src/cmd/go/script_test.go
@@ -41,6 +41,33 @@ func TestScript(t *testing.T) {
 	testenv.MustHaveGoBuild(t)
 	testenv.SkipIfShortAndSlow(t)
 
+	var (
+		ctx         = context.Background()
+		gracePeriod = 100 * time.Millisecond
+	)
+	if deadline, ok := t.Deadline(); ok {
+		timeout := time.Until(deadline)
+
+		// If time allows, increase the termination grace period to 5% of the
+		// remaining time.
+		if gp := timeout / 20; gp > gracePeriod {
+			gracePeriod = gp
+		}
+
+		// When we run commands that execute subprocesses, we want to reserve two
+		// grace periods to clean up. We will send the first termination signal when
+		// the context expires, then wait one grace period for the process to
+		// produce whatever useful output it can (such as a stack trace). After the
+		// first grace period expires, we'll escalate to os.Kill, leaving the second
+		// grace period for the test function to record its output before the test
+		// process itself terminates.
+		timeout -= 2 * gracePeriod
+
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+		t.Cleanup(cancel)
+	}
+
 	files, err := filepath.Glob("testdata/script/*.txt")
 	if err != nil {
 		t.Fatal(err)
@@ -50,40 +77,51 @@ func TestScript(t *testing.T) {
 		name := strings.TrimSuffix(filepath.Base(file), ".txt")
 		t.Run(name, func(t *testing.T) {
 			t.Parallel()
-			ts := &testScript{t: t, name: name, file: file}
+			ctx, cancel := context.WithCancel(ctx)
+			ts := &testScript{
+				t:           t,
+				ctx:         ctx,
+				cancel:      cancel,
+				gracePeriod: gracePeriod,
+				name:        name,
+				file:        file,
+			}
 			ts.setup()
 			if !*testWork {
 				defer removeAll(ts.workdir)
 			}
 			ts.run()
+			cancel()
 		})
 	}
 }
 
 // A testScript holds execution state for a single test script.
 type testScript struct {
-	t          *testing.T
-	workdir    string            // temporary work dir ($WORK)
-	log        bytes.Buffer      // test execution log (printed at end of test)
-	mark       int               // offset of next log truncation
-	cd         string            // current directory during test execution; initially $WORK/gopath/src
-	name       string            // short name of test ("foo")
-	file       string            // full file name ("testdata/script/foo.txt")
-	lineno     int               // line number currently executing
-	line       string            // line currently executing
-	env        []string          // environment list (for os/exec)
-	envMap     map[string]string // environment mapping (matches env)
-	stdout     string            // standard output from last 'go' command; for 'stdout' command
-	stderr     string            // standard error from last 'go' command; for 'stderr' command
-	stopped    bool              // test wants to stop early
-	start      time.Time         // time phase started
-	background []*backgroundCmd  // backgrounded 'exec' and 'go' commands
+	t           *testing.T
+	ctx         context.Context
+	cancel      context.CancelFunc
+	gracePeriod time.Duration
+	workdir     string            // temporary work dir ($WORK)
+	log         bytes.Buffer      // test execution log (printed at end of test)
+	mark        int               // offset of next log truncation
+	cd          string            // current directory during test execution; initially $WORK/gopath/src
+	name        string            // short name of test ("foo")
+	file        string            // full file name ("testdata/script/foo.txt")
+	lineno      int               // line number currently executing
+	line        string            // line currently executing
+	env         []string          // environment list (for os/exec)
+	envMap      map[string]string // environment mapping (matches env)
+	stdout      string            // standard output from last 'go' command; for 'stdout' command
+	stderr      string            // standard error from last 'go' command; for 'stderr' command
+	stopped     bool              // test wants to stop early
+	start       time.Time         // time phase started
+	background  []*backgroundCmd  // backgrounded 'exec' and 'go' commands
 }
 
 type backgroundCmd struct {
 	want           simpleStatus
 	args           []string
-	cancel         context.CancelFunc
 	done           <-chan struct{}
 	err            error
 	stdout, stderr strings.Builder
@@ -109,6 +147,10 @@ var extraEnvKeys = []string{
 
 // setup sets up the test execution temporary directory and environment.
 func (ts *testScript) setup() {
+	if err := ts.ctx.Err(); err != nil {
+		ts.t.Fatalf("test interrupted during setup: %v", err)
+	}
+
 	StartProxy()
 	ts.workdir = filepath.Join(testTmpDir, "script-"+ts.name)
 	ts.check(os.MkdirAll(filepath.Join(ts.workdir, "tmp"), 0777))
@@ -200,9 +242,7 @@ func (ts *testScript) run() {
 		// On a normal exit from the test loop, background processes are cleaned up
 		// before we print PASS. If we return early (e.g., due to a test failure),
 		// don't print anything about the processes that were still running.
-		for _, bg := range ts.background {
-			bg.cancel()
-		}
+		ts.cancel()
 		for _, bg := range ts.background {
 			<-bg.done
 		}
@@ -275,6 +315,10 @@ Script:
 		fmt.Fprintf(&ts.log, "> %s\n", line)
 
 		for _, cond := range parsed.conds {
+			if err := ts.ctx.Err(); err != nil {
+				ts.fatalf("test interrupted: %v", err)
+			}
+
 			// Known conds are: $GOOS, $GOARCH, runtime.Compiler, and 'short' (for testing.Short).
 			//
 			// NOTE: If you make changes here, update testdata/script/README too!
@@ -356,9 +400,7 @@ Script:
 		}
 	}
 
-	for _, bg := range ts.background {
-		bg.cancel()
-	}
+	ts.cancel()
 	ts.cmdWait(success, nil)
 
 	// Final phase ended.
@@ -798,9 +840,7 @@ func (ts *testScript) cmdSkip(want simpleStatus, args []string) {
 
 	// Before we mark the test as skipped, shut down any background processes and
 	// make sure they have returned the correct status.
-	for _, bg := range ts.background {
-		bg.cancel()
-	}
+	ts.cancel()
 	ts.cmdWait(success, nil)
 
 	if len(args) == 1 {
@@ -1065,38 +1105,9 @@ func (ts *testScript) exec(command string, args ...string) (stdout, stderr strin
 func (ts *testScript) startBackground(want simpleStatus, command string, args ...string) (*backgroundCmd, error) {
 	done := make(chan struct{})
 	bg := &backgroundCmd{
-		want:   want,
-		args:   append([]string{command}, args...),
-		done:   done,
-		cancel: func() {},
-	}
-
-	ctx := context.Background()
-	gracePeriod := 100 * time.Millisecond
-	if deadline, ok := ts.t.Deadline(); ok {
-		timeout := time.Until(deadline)
-		// If time allows, increase the termination grace period to 5% of the
-		// remaining time.
-		if gp := timeout / 20; gp > gracePeriod {
-			gracePeriod = gp
-		}
-
-		// Send the first termination signal with two grace periods remaining.
-		// If it still hasn't finished after the first period has elapsed,
-		// we'll escalate to os.Kill with a second period remaining until the
-		// test deadline..
-		timeout -= 2 * gracePeriod
-
-		if timeout <= 0 {
-			// The test has less than the grace period remaining. There is no point in
-			// even starting the command, because it will be terminated immediately.
-			// Save the expense of starting it in the first place.
-			bg.err = context.DeadlineExceeded
-			close(done)
-			return bg, nil
-		}
-
-		ctx, bg.cancel = context.WithTimeout(ctx, timeout)
+		want: want,
+		args: append([]string{command}, args...),
+		done: done,
 	}
 
 	cmd := exec.Command(command, args...)
@@ -1105,12 +1116,11 @@ func (ts *testScript) startBackground(want simpleStatus, command string, args ..
 	cmd.Stdout = &bg.stdout
 	cmd.Stderr = &bg.stderr
 	if err := cmd.Start(); err != nil {
-		bg.cancel()
 		return nil, err
 	}
 
 	go func() {
-		bg.err = waitOrStop(ctx, cmd, stopSignal(), gracePeriod)
+		bg.err = waitOrStop(ts.ctx, cmd, stopSignal(), ts.gracePeriod)
 		close(done)
 	}()
 	return bg, nil