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, 33 insertions, 5 deletions

diff --git a/src/testing/testing_test.go b/src/testing/testing_test.go
index 45e44683b4..9954f9af8c 100644
--- a/src/testing/testing_test.go
+++ b/src/testing/testing_test.go
@@ -5,14 +5,42 @@
 package testing_test
 
 import (
+	"fmt"
 	"os"
+	"runtime"
 	"testing"
+	"time"
 )
 
-// This is exactly what a test would do without a TestMain.
-// It's here only so that there is at least one package in the
-// standard library with a TestMain, so that code is executed.
-
 func TestMain(m *testing.M) {
-	os.Exit(m.Run())
+	g0 := runtime.NumGoroutine()
+
+	code := m.Run()
+	if code != 0 {
+		os.Exit(code)
+	}
+
+	// Check that there are no goroutines left behind.
+	t0 := time.Now()
+	stacks := make([]byte, 1<<20)
+	for {
+		g1 := runtime.NumGoroutine()
+		if g1 == g0 {
+			return
+		}
+		stacks = stacks[:runtime.Stack(stacks, true)]
+		time.Sleep(50 * time.Millisecond)
+		if time.Since(t0) > 2*time.Second {
+			fmt.Fprintf(os.Stderr, "Unexpected leftover goroutines detected: %v -> %v\n%s\n", g0, g1, stacks)
+			os.Exit(1)
+		}
+	}
+}
+
+func TestContextCancel(t *testing.T) {
+	ctx := t.Context()
+	// Tests we don't leak this goroutine:
+	go func() {
+		<-ctx.Done()
+	}()
 }