diff --git a/mock.go b/mock.go index d168849..d70f245 100644 --- a/mock.go +++ b/mock.go @@ -761,6 +761,11 @@ func (t *Trap) matches(c *apiCall) bool { func (t *Trap) Close() { t.mock.mu.Lock() defer t.mock.mu.Unlock() + select { + case <-t.done: + return // already closed + default: + } if t.unreleasedCalls != 0 { t.mock.tb.Helper() t.mock.tb.Errorf("trap Closed() with %d unreleased calls", t.unreleasedCalls) diff --git a/mock_test.go b/mock_test.go index 79e58e0..23e75b1 100644 --- a/mock_test.go +++ b/mock_test.go @@ -1,9 +1,13 @@ package quartz_test import ( + "bytes" "context" "errors" "fmt" + "os" + "runtime/pprof" + "strings" "testing" "time" @@ -403,7 +407,10 @@ func Test_UnreleasedCalls(t *testing.T) { _ = mClock.Now() }() - trap.MustWait(testCtx) // missing release + c := trap.MustWait(testCtx) // missing release + trap.Close() // detect unreleased call and fail + + c.Release(testCtx) // clean up goroutine }) } @@ -573,3 +580,55 @@ func TestTickerStop_Go123(t *testing.T) { // OK! } } + +func TestMain(m *testing.M) { + verifyNoLeakTestMain(m) +} + +func verifyNoLeakTestMain(m *testing.M) { + before := snapshot() + code := m.Run() + now := time.Now() + for { + after := snapshot() + if len(after) > len(before) { + // Allow test cleanup to settle. + if time.Since(now) < 200*time.Millisecond { + time.Sleep(50 * time.Millisecond) + continue + } + fmt.Fprintln(os.Stderr, "Possible goroutine leak(s):") + fmt.Fprintln(os.Stderr, diff(before, after)) + os.Exit(1) + } + os.Exit(code) + } +} + +func snapshot() []string { + var buf bytes.Buffer + _ = pprof.Lookup("goroutine").WriteTo(&buf, 2) + var clean []string + for _, s := range strings.Split(buf.String(), "\n\n") { + if !strings.Contains(s, "runtime/pprof") { + clean = append(clean, s) + } + } + return clean +} + +func diff(a, b []string) string { + m := make(map[string]int) + for _, s := range a { + m[s]++ + } + var leaks []string + for _, s := range b { + if m[s] > 0 { + m[s]-- + continue + } + leaks = append(leaks, s) + } + return strings.Join(leaks, "\n\n") +}