Add Cleanup option

leaks.go

@@ -55,6 +55,9 @@ func Find(options ...Option) error {
 	cur := stack.Current().ID()
 
 	opts := buildOpts(options...)
+	if opts.teardown != nil {
+		return fmt.Errorf("Cleanup can only be passed to VerifyNone or VerifyTestMain")
+	}
 	var stacks []stack.Stack
 	retry := true
 	for i := 0; retry; i++ {
@@ -79,12 +82,17 @@ type testHelper interface {
 //
 //	defer VerifyNone(t)
 func VerifyNone(t TestingT, options ...Option) {
+	opts := buildOpts(options...)
 	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 opts.teardown != nil {
+		opts.teardown(0)
+	}
 }

leaks_test.go

@@ -51,6 +51,10 @@ func TestFind(t *testing.T) {
 	// Once we unblock the goroutine, we shouldn't have leaks.
 	bg.unblock()
 	require.NoError(t, Find(), "Should find no leaks by default")
+
+	// Find can't take in Cleanup option
+	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) {
@@ -82,6 +86,13 @@ func TestVerifyNone(t *testing.T) {
 	VerifyNone(ft, testOptions())
 	require.NotEmpty(t, ft.errors, "Expect errors from VerifyNone on leaked goroutine")
 	bg.unblock()
+
+	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
+	teardown   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.teardown = o.teardown
 }
 
 // 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.
+// 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.teardown = 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,17 @@ 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...)
 
 	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)
+	if opts.teardown != nil {
+		opts.teardown(exitCode)
+	} else {
+		_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)
 }