Use context.Context for configurable timeouts

.travis.yml

@@ -1,8 +1,7 @@
 language: go
 go:
- - 1.5.3
- - 1.6.3
  - 1.7
+ - 1.8
  - tip
 
 script:

README.md

@@ -1,24 +1,31 @@
-Leaktest [![Build Status](https://travis-ci.org/fortytw2/leaktest.svg?branch=master)](https://travis-ci.org/fortytw2/leaktest)
+Leaktest [![Build Status](https://travis-ci.org/fortytw2/leaktest.svg?branch=master)](https://travis-ci.org/fortytw2/leaktest) [![codecov](https://codecov.io/gh/fortytw2/leaktest/branch/master/graph/badge.svg)](https://codecov.io/gh/fortytw2/leaktest)
 ------
 
-Refactored, tested variant of the goroutine leak detector found in both `net/http` tests and the `cockroachdb`
-source tree.
+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.
 
 ### Installation
 
+Go 1.7+
+
 ```
 go get -u github.com/fortytw2/leaktest
 ```
 
+Go 1.5/1.6 need to use the tag `v1.0.0`, as newer versions depend on
+`context.Context`. 
+
 ### Example
 
-This test fails, because it leaks a goroutine :o
+These tests fail, because they leak a goroutine
 
 ```go
+// Default "Check" will poll for 5 seconds to check that all
+// goroutines are cleaned up
 func TestPool(t *testing.T) {
 	defer leaktest.Check(t)()
 
@@ -28,6 +35,29 @@ func TestPool(t *testing.T) {
         }
     }()
 }
+
+// Helper function to timeout after X duration
+func TestPoolTimeout(t *testing.T) {
+	defer leaktest.CheckTimeout(t, time.Second)()
+
+    go func() {
+        for {
+            time.Sleep(time.Second)
+        }
+    }()
+}
+
+// Use Go 1.7+ context.Context for cancellation
+func TestPoolContext(t *testing.T) {
+    ctx, _ := context.WithTimeout(context.Background(), time.Second)
+	defer leaktest.CheckContext(ctx, t)()
+
+    go func() {
+        for {
+            time.Sleep(time.Second)
+        }
+    }()
+}
 ```
 
 

leaktest.go

@@ -10,6 +10,7 @@
 package leaktest
 
 import (
+	"context"
 	"runtime"
 	"sort"
 	"strings"
@@ -59,34 +60,52 @@ type ErrorReporter interface {
 
 // Check snapshots the currently-running goroutines and returns a
 // function to be run at the end of tests to see whether any
-// goroutines leaked.
+// goroutines leaked, waiting up to 5 seconds in error conditions
 func Check(t ErrorReporter) func() {
+	return CheckTimeout(t, 5*time.Second)
+}
+
+// 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)
+	fn := CheckContext(ctx, t)
+	return func() {
+		fn()
+		cancel()
+	}
+}
+
+// CheckContext is the same as Check, but uses a context.Context for
+// cancellation and timeout control
+func CheckContext(ctx context.Context, t ErrorReporter) func() {
 	orig := map[string]bool{}
 	for _, g := range interestingGoroutines() {
 		orig[g] = true
 	}
 	return func() {
-		// Loop, waiting for goroutines to shut down.
-		// Wait up to 5 seconds, but finish as quickly as possible.
-		deadline := time.Now().Add(5 * time.Second)
+		var leaked []string
 		for {
-			var leaked []string
-			for _, g := range interestingGoroutines() {
-				if !orig[g] {
-					leaked = append(leaked, g)
+			select {
+			case <-ctx.Done():
+				t.Errorf("leaktest: timed out checking goroutines")
+			default:
+				leaked = make([]string, 0)
+				for _, g := range interestingGoroutines() {
+					if !orig[g] {
+						leaked = append(leaked, g)
+					}
 				}
-			}
-			if len(leaked) == 0 {
-				return
-			}
-			if time.Now().Before(deadline) {
-				time.Sleep(50 * time.Millisecond)
+				if len(leaked) == 0 {
+					return
+				}
+				// don't spin needlessly
+				time.Sleep(time.Millisecond * 50)
 				continue
 			}
-			for _, g := range leaked {
-				t.Errorf("Leaked goroutine: %v", g)
-			}
-			return
+			break
+		}
+		for _, g := range leaked {
+			t.Errorf("leaktest: leaked goroutine: %v", g)
 		}
 	}
 }

leaktest_test.go

@@ -1,6 +1,7 @@
 package leaktest
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"testing"
@@ -61,13 +62,12 @@ var leakyFuncs = []func(){
 }
 
 func TestCheck(t *testing.T) {
-
 	// this works because the running goroutine is left running at the
 	// start of the next test case - so the previous leaks don't affect the
 	// check for the next one
 	for i, fn := range leakyFuncs {
 		checker := &testReporter{}
-		snapshot := Check(checker)
+		snapshot := CheckTimeout(checker, time.Second)
 		go fn()
 
 		snapshot()
@@ -78,6 +78,8 @@ func TestCheck(t *testing.T) {
 }
 
 func TestEmptyLeak(t *testing.T) {
-	defer Check(t)()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	defer CheckContext(ctx, t)()
 	time.Sleep(time.Second)
 }