Detect goroutines running for too long in testcontext (#685)

internal/testcontext/context.go

@@ -5,40 +5,93 @@ package testcontext
 
 import (
 	"context"
+	"fmt"
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"runtime"
+	"strings"
 	"sync"
-	"testing"
+	"time"
 
 	"golang.org/x/sync/errgroup"
 )
 
+const defaultTimeout = 3 * time.Minute
+
 // Context is a context that has utility methods for testing and waiting for asynchronous errors.
 type Context struct {
 	context.Context
+
+	timedctx context.Context
+	cancel   context.CancelFunc
+
 	group *errgroup.Group
-	test  testing.TB
+	test  TB
 
 	once      sync.Once
 	directory string
+
+	mu      sync.Mutex
+	running []caller
+}
+
+type caller struct {
+	pc   uintptr
+	file string
+	line int
+	ok   bool
+	done bool
+}
+
+// TB is a subset of testing.TB methods
+type TB interface {
+	Name() string
+	Helper()
+	Error(args ...interface{})
+	Fatal(args ...interface{})
 }
 
 // New creates a new test context
-func New(test testing.TB) *Context {
-	group, ctx := errgroup.WithContext(context.Background())
-	return &Context{
-		Context: ctx,
-		group:   group,
-		test:    test,
+func New(test TB) *Context {
+	return NewWithTimeout(test, defaultTimeout)
+}
+
+// NewWithTimeout creates a new test context with a given timeout
+func NewWithTimeout(test TB, timeout time.Duration) *Context {
+	timedctx, cancel := context.WithTimeout(context.Background(), timeout)
+	group, errctx := errgroup.WithContext(timedctx)
+
+	ctx := &Context{
+		Context:  errctx,
+		timedctx: timedctx,
+		cancel:   cancel,
+		group:    group,
+		test:     test,
 	}
+
+	return ctx
 }
 
 // Go runs fn in a goroutine.
 // Call Wait to check the result
 func (ctx *Context) Go(fn func() error) {
 	ctx.test.Helper()
-	ctx.group.Go(fn)
+
+	pc, file, line, ok := runtime.Caller(1)
+	ctx.mu.Lock()
+	index := len(ctx.running)
+	ctx.running = append(ctx.running, caller{pc, file, line, ok, false})
+	ctx.mu.Unlock()
+
+	ctx.group.Go(func() error {
+		defer func() {
+			ctx.mu.Lock()
+			ctx.running[index].done = true
+			ctx.mu.Unlock()
+		}()
+		return fn()
+	})
 }
 
 // Check calls fn and checks result
@@ -85,12 +138,50 @@ func (ctx *Context) Cleanup() {
 	ctx.test.Helper()
 
 	defer ctx.deleteTemporary()
-	err := ctx.group.Wait()
-	if err != nil {
-		ctx.test.Fatal(err)
+	defer ctx.cancel()
+
+	alldone := make(chan error, 1)
+	go func() {
+		alldone <- ctx.group.Wait()
+		defer close(alldone)
+	}()
+
+	select {
+	case <-ctx.timedctx.Done():
+		ctx.reportRunning()
+	case err := <-alldone:
+		if err != nil {
+			ctx.test.Fatal(err)
+		}
 	}
 }
 
+func (ctx *Context) reportRunning() {
+	ctx.mu.Lock()
+	defer ctx.mu.Unlock()
+
+	var problematic []caller
+	for _, caller := range ctx.running {
+		if !caller.done {
+			problematic = append(problematic, caller)
+		}
+	}
+
+	var message strings.Builder
+	message.WriteString("Test exceeded timeout")
+	if len(problematic) > 0 {
+		message.WriteString("\nsome goroutines are still running, did you forget to shut them down?")
+		for _, caller := range problematic {
+			fnname := ""
+			if fn := runtime.FuncForPC(caller.pc); fn != nil {
+				fnname = fn.Name()
+			}
+			fmt.Fprintf(&message, "\n%s:%d: %s", caller.file, caller.line, fnname)
+		}
+	}
+	ctx.test.Error(message.String())
+}
+
 // deleteTemporary tries to delete temporary directory
 func (ctx *Context) deleteTemporary() {
 	if ctx.directory == "" {

internal/testcontext/context_test.go

@@ -4,9 +4,12 @@
 package testcontext_test
 
 import (
+	"fmt"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+
 	"storj.io/storj/internal/testcontext"
 )
 
@@ -22,3 +25,47 @@ func TestBasic(t *testing.T) {
 	t.Log(ctx.Dir("a", "b", "c"))
 	t.Log(ctx.File("a", "w", "c.txt"))
 }
+
+func TestTimeout(t *testing.T) {
+	ok := testing.RunTests(nil, []testing.InternalTest{{
+		Name: "TimeoutFailure",
+		F: func(t *testing.T) {
+			ctx := testcontext.NewWithTimeout(t, 50*time.Millisecond)
+			defer ctx.Cleanup()
+
+			ctx.Go(func() error {
+				time.Sleep(time.Second)
+				return nil
+			})
+		},
+	}})
+
+	if ok {
+		t.Error("test should have failed")
+	}
+}
+
+func TestMessage(t *testing.T) {
+	var subtest test
+
+	ctx := testcontext.NewWithTimeout(&subtest, 50*time.Millisecond)
+	ctx.Go(func() error {
+		time.Sleep(time.Second)
+		return nil
+	})
+	ctx.Cleanup()
+
+	assert.Contains(t, subtest.errors[0], "Test exceeded timeout")
+	assert.Contains(t, subtest.errors[0], "some goroutines are still running")
+}
+
+type test struct {
+	errors []string
+	fatals []string
+}
+
+func (t *test) Name() string { return "Example" }
+func (t *test) Helper()      {}
+
+func (t *test) Error(args ...interface{}) { t.errors = append(t.errors, fmt.Sprint(args...)) }
+func (t *test) Fatal(args ...interface{}) { t.fatals = append(t.fatals, fmt.Sprint(args...)) }