Add Cleanup option (#78)

This adds Cleanup option which can be passed to VerifyTestMain and
VerifyNone. This takes in a function that will be executed at the
end of the leak verification.

Internal Ref: GO-888
Fix #63

Co-authored-by: Abhinav Gupta <mail@abhinavg.net>
Co-authored-by: Abhinav Gupta <abg@uber.com>

leaks.go

@@ -21,6 +21,7 @@
 package goleak
 
 import (
+	"errors"
 	"fmt"
 
 	"go.uber.org/goleak/internal/stack"
@@ -55,6 +56,9 @@ func Find(options ...Option) error {
 	cur := stack.Current().ID()
 
 	opts := buildOpts(options...)
+	if opts.cleanup != nil {
+		return errors.New("Cleanup can only be passed to VerifyNone or VerifyTestMain")
+	}
 	var stacks []stack.Stack
 	retry := true
 	for i := 0; retry; i++ {
@@ -79,12 +83,20 @@ type testHelper interface {
 //
 //	defer VerifyNone(t)
 func VerifyNone(t TestingT, options ...Option) {
+	opts := buildOpts(options...)
+	var cleanup func(int)
+	cleanup, opts.cleanup = opts.cleanup, nil
+
 	if h, ok := t.(testHelper); ok {
 		// Mark this function as a test helper, if available.
 		h.Helper()
 	}
 
-	if err := Find(options...); err != nil {
+	if err := Find(opts); err != nil {
 		t.Error(err)
 	}
+
+	if cleanup != nil {
+		cleanup(0)
+	}
 }

leaks_test.go

@@ -40,17 +40,26 @@ func testOptions() Option {
 }
 
 func TestFind(t *testing.T) {
-	require.NoError(t, Find(), "Should find no leaks by default")
+	t.Run("Should find no leaks by default", func(t *testing.T) {
+		require.NoError(t, Find())
+	})
 
-	bg := startBlockedG()
-	err := Find(testOptions())
-	require.Error(t, err, "Should find leaks with leaked goroutine")
-	assert.Contains(t, err.Error(), "blockedG")
-	assert.Contains(t, err.Error(), "created by go.uber.org/goleak.startBlockedG")
-
-	// Once we unblock the goroutine, we shouldn't have leaks.
-	bg.unblock()
-	require.NoError(t, Find(), "Should find no leaks by default")
+	t.Run("Find leaks with leaked goroutine", func(t *testing.T) {
+		bg := startBlockedG()
+		err := Find(testOptions())
+		require.Error(t, err, "Should find leaks with leaked goroutine")
+		assert.Contains(t, err.Error(), "blockedG")
+		assert.Contains(t, err.Error(), "created by go.uber.org/goleak.startBlockedG")
+
+		// Once we unblock the goroutine, we shouldn't have leaks.
+		bg.unblock()
+		require.NoError(t, Find(), "Should find no leaks by default")
+	})
+
+	t.Run("Find can't take in Cleanup option", func(t *testing.T) {
+		err := Find(Cleanup(func(int) { assert.Fail(t, "this should not be called") }))
+		require.Error(t, err, "Should exit with invalid option")
+	})
 }
 
 func TestFindRetry(t *testing.T) {
@@ -74,14 +83,26 @@ func (ft *fakeT) Error(args ...interface{}) {
 }
 
 func TestVerifyNone(t *testing.T) {
-	ft := &fakeT{}
-	VerifyNone(ft)
-	require.Empty(t, ft.errors, "Expect no errors from VerifyNone")
+	t.Run("VerifyNone finds leaks", func(t *testing.T) {
+		ft := &fakeT{}
+		VerifyNone(ft)
+		require.Empty(t, ft.errors, "Expect no errors from VerifyNone")
+
+		bg := startBlockedG()
+		VerifyNone(ft, testOptions())
+		require.NotEmpty(t, ft.errors, "Expect errors from VerifyNone on leaked goroutine")
+		bg.unblock()
+	})
 
-	bg := startBlockedG()
-	VerifyNone(ft, testOptions())
-	require.NotEmpty(t, ft.errors, "Expect errors from VerifyNone on leaked goroutine")
-	bg.unblock()
+	t.Run("cleanup registered callback should be called", func(t *testing.T) {
+		ft := &fakeT{}
+		cleanupCalled := false
+		VerifyNone(ft, Cleanup(func(c int) {
+			assert.Equal(t, 0, c)
+			cleanupCalled = true
+		}))
+		require.True(t, cleanupCalled, "expect cleanup registered callback to be called")
+	})
 }
 
 func TestIgnoreCurrent(t *testing.T) {

options.go

@@ -41,6 +41,16 @@ type opts struct {
 	filters    []func(stack.Stack) bool
 	maxRetries int
 	maxSleep   time.Duration
+	cleanup    func(int)
+}
+
+// implement apply so that opts struct itself can be used as
+// an Option.
+func (o *opts) apply(opts *opts) {
+	opts.filters = o.filters
+	opts.maxRetries = o.maxRetries
+	opts.maxSleep = o.maxSleep
+	opts.cleanup = o.cleanup
 }
 
 // optionFunc lets us easily write options without a custom type.
@@ -57,6 +67,18 @@ func IgnoreTopFunction(f string) Option {
 	})
 }
 
+// Cleanup sets up a cleanup function that will be executed at the
+// end of the leak check.
+// When passed to [VerifyTestMain], the exit code passed to cleanupFunc
+// will be set to the exit code of TestMain.
+// When passed to [VerifyNone], the exit code will be set to 0.
+// This cannot be passed to [Find].
+func Cleanup(cleanupFunc func(exitCode int)) Option {
+	return optionFunc(func(opts *opts) {
+		opts.cleanup = cleanupFunc
+	})
+}
+
 // IgnoreCurrent records all current goroutines when the option is created, and ignores
 // them in any future Find/Verify calls.
 func IgnoreCurrent() Option {
@@ -98,23 +120,23 @@ func buildOpts(options ...Option) *opts {
 	return opts
 }
 
-func (vo *opts) filter(s stack.Stack) bool {
-	for _, filter := range vo.filters {
+func (o *opts) filter(s stack.Stack) bool {
+	for _, filter := range o.filters {
 		if filter(s) {
 			return true
 		}
 	}
 	return false
 }
 
-func (vo *opts) retry(i int) bool {
-	if i >= vo.maxRetries {
+func (o *opts) retry(i int) bool {
+	if i >= o.maxRetries {
 		return false
 	}
 
 	d := time.Duration(int(time.Microsecond) << uint(i))
-	if d > vo.maxSleep {
-		d = vo.maxSleep
+	if d > o.maxSleep {
+		d = o.maxSleep
 	}
 	time.Sleep(d)
 	return true

testmain.go

@@ -51,13 +51,19 @@ type TestingM interface {
 // for any goroutine leaks and fail the tests if any leaks were found.
 func VerifyTestMain(m TestingM, options ...Option) {
 	exitCode := m.Run()
+	opts := buildOpts(options...)
+
+	var cleanup func(int)
+	cleanup, opts.cleanup = opts.cleanup, nil
+	if cleanup == nil {
+		cleanup = _osExit
+	}
+	defer func() { cleanup(exitCode) }()
 
 	if exitCode == 0 {
-		if err := Find(options...); err != nil {
+		if err := Find(opts); err != nil {
 			fmt.Fprintf(_osStderr, "goleak: Errors on successful test run: %v\n", err)
 			exitCode = 1
 		}
 	}
-
-	_osExit(exitCode)
 }

testmain_test.go

@@ -75,4 +75,13 @@ func TestVerifyTestMain(t *testing.T) {
 	VerifyTestMain(dummyTestMain(0))
 	assert.Equal(t, 0, <-exitCode, "Expect no errors without leaks")
 	assert.NotContains(t, <-stderr, "goleak: Errors", "No errors on successful run without leaks")
+
+	cleanupCalled := false
+	cleanupExitcode := 0
+	VerifyTestMain(dummyTestMain(3), Cleanup(func(ec int) {
+		cleanupCalled = true
+		cleanupExitcode = ec
+	}))
+	assert.True(t, cleanupCalled)
+	assert.Equal(t, 3, cleanupExitcode)
 }