add CheckMain and support for t.Parallel()

README.md

@@ -5,8 +5,10 @@ Refactored, tested variant of the goroutine leak detector found in both
 `net/http` tests and the `cockroachdb` source tree.
 
 Takes a snapshot of running goroutines at the start of a test, and at the end -
-compares the two and *voila*. Ignores runtime/sys goroutines. Doesn't play nice
-with `t.Parallel()` right now, but there are plans to do so.
+compares the two and *voila*. Ignores runtime/sys goroutines. In order to work 
+cleanly with `t.Parallel()` you must use `CheckMain(m)` from within a `TestMain(m *testing.M)` 
+function, which will only run one copy of `Check` for all tests within the package.
+
 
 ### Installation
 

leaktest.go

@@ -12,10 +12,12 @@ package leaktest
 import (
 	"context"
 	"fmt"
+	"os"
 	"runtime"
 	"sort"
 	"strconv"
 	"strings"
+	"testing"
 	"time"
 )
 
@@ -94,13 +96,42 @@ type ErrorReporter interface {
 	Errorf(format string, args ...interface{})
 }
 
+type testReporter struct {
+	failed bool
+	msg    string
+}
+
+func (tr *testReporter) Errorf(format string, args ...interface{}) {
+	tr.failed = true
+	tr.msg = fmt.Sprintf(format, args)
+}
+
 // Check snapshots the currently-running goroutines and returns a
 // function to be run at the end of tests to see whether any
 // goroutines leaked, waiting up to 5 seconds in error conditions
 func Check(t ErrorReporter) func() {
 	return CheckTimeout(t, 5*time.Second)
 }
 
+// CheckMain can be used to Check many tests at once, or when they run in
+// parallel, in exchange for providing relatively worse information on
+// failures compared to Check
+func CheckMain(m *testing.M) {
+	t := &testReporter{}
+	fn := Check(t)
+
+	code := m.Run()
+	if code != 0 {
+		os.Exit(code)
+	}
+
+	fn()
+	if t.failed {
+		fmt.Println(t.msg)
+		os.Exit(1)
+	}
+}
+
 // CheckTimeout is the same as Check, but with a configurable timeout
 func CheckTimeout(t ErrorReporter, dur time.Duration) func() {
 	ctx, cancel := context.WithTimeout(context.Background(), dur)

leaktest_main_test.go

@@ -0,0 +1,17 @@
+package leaktest_test
+
+import (
+	"testing"
+
+	"github.com/fortytw2/leaktest"
+)
+
+func TestMain(m *testing.M) {
+	leaktest.CheckMain(m)
+}
+
+func TestSilly(t *testing.T) {
+	go func() {
+		select {}
+	}()
+}

leaktest_test.go

@@ -3,22 +3,11 @@ package leaktest
 import (
 	"context"
 	"errors"
-	"fmt"
 	"sync"
 	"testing"
 	"time"
 )
 
-type testReporter struct {
-	failed bool
-	msg    string
-}
-
-func (tr *testReporter) Errorf(format string, args ...interface{}) {
-	tr.failed = true
-	tr.msg = fmt.Sprintf(format, args)
-}
-
 var leakyFuncs = []func(){
 	// Infinite for loop
 	func() {