Make maxSleep and maxRetires configurable when building options (#94)

Make goleak options more flexible to adapt to users' variety scenarios

Signed-off-by: Kante Yin <[email protected]>

fix  #93

---------

Signed-off-by: Kante Yin <[email protected]>
Co-authored-by: Sung Yoon Whang <[email protected]>
Co-authored-by: Abhinav Gupta <[email protected]>

leaks.go

@@ -56,6 +56,9 @@ func Find(options ...Option) error {
 	cur := stack.Current().ID()
 
 	opts := buildOpts(options...)
+	if err := opts.validate(); err != nil {
+		return err
+	}
 	if opts.cleanup != nil {
 		return errors.New("Cleanup can only be passed to VerifyNone or VerifyTestMain")
 	}

leaks_test.go

@@ -36,7 +36,7 @@ var _ = TestingT(testing.TB(nil))
 // testOptions passes a shorter max sleep time, used so tests don't wait
 // ~1 second in cases where we expect Find to error out.
 func testOptions() Option {
-	return maxSleep(time.Millisecond)
+	return MaxSleepInterval(time.Millisecond)
 }
 
 func TestFind(t *testing.T) {
@@ -60,6 +60,16 @@ func TestFind(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")
 	})
+
+	t.Run("Find should return error when maxRetries is less than 0", func(t *testing.T) {
+		err := Find(MaxRetryAttempts(-1))
+		require.Error(t, err, "maxRetries should be greater than 0")
+	})
+
+	t.Run("Find should return error when maxSleep is less than 0s", func(t *testing.T) {
+		err := Find(MaxSleepInterval(time.Duration(-1)))
+		require.Error(t, err, "maxSleep should be greater than 0s")
+	})
 }
 
 func TestFindRetry(t *testing.T) {

options.go

@@ -21,6 +21,7 @@
 package goleak
 
 import (
+	"errors"
 	"strings"
 	"time"
 
@@ -32,10 +33,14 @@ type Option interface {
 	apply(*opts)
 }
 
-// We retry up to 20 times if we can't find the goroutine that
-// we are looking for. In between each attempt, we will sleep for
-// a short while to let any running goroutines complete.
-const _defaultRetries = 20
+const (
+	// We retry up to default 20 times if we can't find the goroutine that
+	// we are looking for.
+	_defaultRetryAttempts = 20
+	// In between each retry attempt, sleep for up to default 100 microseconds
+	// to let any running goroutine completes.
+	_defaultSleepInterval = 100 * time.Microsecond
+)
 
 type opts struct {
 	filters    []func(stack.Stack) bool
@@ -53,6 +58,17 @@ func (o *opts) apply(opts *opts) {
 	opts.cleanup = o.cleanup
 }
 
+// validate the options.
+func (o *opts) validate() error {
+	if o.maxRetries < 0 {
+		return errors.New("maxRetryAttempts should be greater than 0")
+	}
+	if o.maxSleep <= 0 {
+		return errors.New("maxSleepInterval should be greater than 0s")
+	}
+	return nil
+}
+
 // optionFunc lets us easily write options without a custom type.
 type optionFunc func(*opts)
 
@@ -91,12 +107,25 @@ func IgnoreCurrent() Option {
 	})
 }
 
-func maxSleep(d time.Duration) Option {
+// MaxSleepInterval sets the maximum sleep time in-between each retry attempt.
+// The sleep duration grows in an exponential backoff, to a maximum of the value specified here.
+// If not configured, default to 100 microseconds.
+func MaxSleepInterval(d time.Duration) Option {
 	return optionFunc(func(opts *opts) {
 		opts.maxSleep = d
 	})
 }
 
+// MaxRetryAttempts sets the retry upper limit.
+// When finding extra goroutines, we'll retry until all goroutines complete
+// or end up with the maximum retry attempts.
+// If not configured, default to 20 times.
+func MaxRetryAttempts(num int) Option {
+	return optionFunc(func(opts *opts) {
+		opts.maxRetries = num
+	})
+}
+
 func addFilter(f func(stack.Stack) bool) Option {
 	return optionFunc(func(opts *opts) {
 		opts.filters = append(opts.filters, f)
@@ -105,8 +134,8 @@ func addFilter(f func(stack.Stack) bool) Option {
 
 func buildOpts(options ...Option) *opts {
 	opts := &opts{
-		maxRetries: _defaultRetries,
-		maxSleep:   100 * time.Millisecond,
+		maxRetries: _defaultRetryAttempts,
+		maxSleep:   _defaultSleepInterval,
 	}
 	opts.filters = append(opts.filters,
 		isTestStack,
@@ -117,6 +146,7 @@ func buildOpts(options ...Option) *opts {
 	for _, option := range options {
 		option.apply(opts)
 	}
+
 	return opts
 }
 

options_test.go

@@ -65,10 +65,20 @@ func TestOptionsFilters(t *testing.T) {
 	require.Zero(t, countUnfiltered(), "blockedG should be filtered out. running: %v", stack.All())
 }
 
-func TestOptionsRetry(t *testing.T) {
+func TestBuildOptions(t *testing.T) {
+	// With default options.
 	opts := buildOpts()
-	opts.maxRetries = 50 // initial attempt + 50 retries = 11
-	opts.maxSleep = time.Millisecond
+	assert.Equal(t, _defaultSleepInterval, opts.maxSleep, "value of maxSleep not right")
+	assert.Equal(t, _defaultRetryAttempts, opts.maxRetries, "value of maxRetries not right")
+
+	// With customized options.
+	opts = buildOpts(MaxRetryAttempts(50), MaxSleepInterval(time.Microsecond))
+	assert.Equal(t, time.Microsecond, opts.maxSleep, "value of maxSleep not right")
+	assert.Equal(t, 50, opts.maxRetries, "value of maxRetries not right")
+}
+
+func TestOptionsRetry(t *testing.T) {
+	opts := buildOpts(MaxSleepInterval(time.Millisecond), MaxRetryAttempts(50)) // initial attempt + 50 retries = 51
 
 	for i := 0; i < 50; i++ {
 		assert.True(t, opts.retry(i), "Attempt %v/51 should allow retrying", i)