testing: add T.Context method

From the doc comment:

Context returns the context for the current test or benchmark.
The context is cancelled when the test or benchmark finishes.
A goroutine started during a test or benchmark can wait for the
context's Done channel to become readable as a signal that the
test or benchmark is over, so that the goroutine can exit.

Fixes #16221.
Fixes #17552.

Change-Id: I657df946be2c90048cc74615436c77c7d9d1226c
Reviewed-on: https://go-review.googlesource.com/31724
Reviewed-by: Rob Pike <r@golang.org>

1 files changed, 29 insertions, 6 deletions

diff --git a/src/testing/testing.go b/src/testing/testing.go
index 31290aaec0..01f5da31d7 100644
--- a/src/testing/testing.go
+++ b/src/testing/testing.go
@@ -204,6 +204,7 @@ package testing
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"flag"
 	"fmt"
@@ -261,12 +262,14 @@ type common struct {
 	mu         sync.RWMutex // guards output, failed, and done.
 	output     []byte       // Output generated by test or benchmark.
 	w          io.Writer    // For flushToParent.
-	chatty     bool         // A copy of the chatty flag.
-	ran        bool         // Test or benchmark (or one of its subtests) was executed.
-	failed     bool         // Test or benchmark has failed.
-	skipped    bool         // Test of benchmark has been skipped.
-	finished   bool         // Test function has completed.
-	done       bool         // Test is finished and all subtests have completed.
+	ctx        context.Context
+	cancel     context.CancelFunc
+	chatty     bool // A copy of the chatty flag.
+	ran        bool // Test or benchmark (or one of its subtests) was executed.
+	failed     bool // Test or benchmark has failed.
+	skipped    bool // Test of benchmark has been skipped.
+	finished   bool // Test function has completed.
+	done       bool // Test is finished and all subtests have completed.
 	hasSub     bool
 	raceErrors int // number of races detected during test
 
@@ -280,6 +283,13 @@ type common struct {
 	sub      []*T      // Queue of subtests to be run in parallel.
 }
 
+func (c *common) parentContext() context.Context {
+	if c == nil || c.parent == nil || c.parent.ctx == nil {
+		return context.Background()
+	}
+	return c.parent.ctx
+}
+
 // Short reports whether the -test.short flag is set.
 func Short() bool {
 	return *short
@@ -376,6 +386,7 @@ func fmtDuration(d time.Duration) string {
 
 // TB is the interface common to T and B.
 type TB interface {
+	Context() context.Context
 	Error(args ...interface{})
 	Errorf(format string, args ...interface{})
 	Fail()
@@ -423,6 +434,15 @@ func (c *common) Name() string {
 	return c.name
 }
 
+// Context returns the context for the current test or benchmark.
+// The context is cancelled when the test or benchmark finishes.
+// A goroutine started during a test or benchmark can wait for the
+// context's Done channel to become readable as a signal that the
+// test or benchmark is over, so that the goroutine can exit.
+func (c *common) Context() context.Context {
+	return c.ctx
+}
+
 func (c *common) setRan() {
 	if c.parent != nil {
 		c.parent.setRan()
@@ -599,6 +619,9 @@ type InternalTest struct {
 }
 
 func tRunner(t *T, fn func(t *T)) {
+	t.ctx, t.cancel = context.WithCancel(t.parentContext())
+	defer t.cancel()
+
 	// When this goroutine is done, either because fn(t)
 	// returned normally or because a test failure triggered
 	// a call to runtime.Goexit, record the duration and send