feat(util): try to wait with context in retry

cpu: Intel(R) Core(TM) i7-10700 CPU @ 2.90GHz BenchmarkWait BenchmarkWait/success BenchmarkWait/success-4 432140625 2.793 ns/op 0 B/op 0 allocs/op BenchmarkWait/retry-3 BenchmarkWait/retry-3-4 97395789 12.27 ns/op 0 B/op 0 allocs/op BenchmarkContext BenchmarkContext/success BenchmarkContext/success-4 449174043 2.641 ns/op 0 B/op 0 allocs/op BenchmarkContext/retry-3 BenchmarkContext/retry-3-4 3177039 380.2 ns/op 232 B/op 3 allocs/op @formatter:off Closes #3758 Signed-off-by: slasher <shenjie1@oppo.com>
2026-08-02 02:00:56 +00:00 · 2025-04-09 14:14:38 +08:00 · 2025-04-09 14:14:38 +08:00 · 65889a5213
commit 65889a5213
parent 11983f9888
2 changed files with 187 additions and 3 deletions
--- a/blobstore/util/retry/retry.go
+++ b/blobstore/util/retry/retry.go
@ -15,6 +15,7 @@
 package retry

 import (
+	"context"
 	"errors"
 	"time"
 )
@ -28,20 +29,49 @@ var (

 // Retryer is an interface retry on a specific function.
 type Retryer interface {
+	// Reset the delay as beginning.
+	Reset()
 	// On performs a retry on function, until it doesn't return any error.
 	On(func() error) error
+	// OnContext On or context done.
+	OnContext(context.Context, func() error) error
 	// RuptOn performs a retry on function, until it doesn't return any error or interrupt.
 	RuptOn(func() (bool, error)) error
+	// RuptOnContext RuptOn or context done.
+	RuptOnContext(context.Context, func() (bool, error)) error
 }

+// retry non-thread-safer implements Retryer
 type retry struct {
 	attempts  int
+	reset     func()
 	nextDelay func() uint32
 }

+type internalContext struct{}
+
+func (internalContext) Deadline() (deadline time.Time, ok bool) { return }
+func (internalContext) Done() <-chan struct{}                   { return nil }
+func (internalContext) Err() error                              { return nil }
+func (internalContext) Value(key interface{}) interface{}       { return nil }
+
+func (r *retry) Reset() {
+	if r.reset != nil {
+		r.reset()
+	}
+}
+
 // On implements Retryer.On.
 func (r *retry) On(caller func() error) error {
+	return r.OnContext(internalContext{}, caller)
+}
+
+// OnContext implements Retryer.OnContext.
+func (r *retry) OnContext(ctx context.Context, caller func() error) error {
 	var lastErr error
+	var timer *time.Timer
+	var ctxDone bool
+
 	attempt := 1
 	for attempt <= r.attempts {
 		if lastErr = caller(); lastErr == nil {
@ -52,15 +82,29 @@ func (r *retry) On(caller func() error) error {
 		if attempt >= r.attempts {
 			break
 		}
-		time.Sleep(time.Duration(r.nextDelay()) * time.Millisecond)
 		attempt++
+
+		if timer, ctxDone = r.waitContext(ctx, timer); ctxDone {
+			break
+		}
+	}
+	if timer != nil {
+		timer.Stop()
 	}
 	return lastErr
 }

 // RuptOn implements Retryer.RuptOn.
 func (r *retry) RuptOn(caller func() (bool, error)) error {
+	return r.RuptOnContext(internalContext{}, caller)
+}
+
+// RuptOnContext implements Retryer.RuptOnContext.
+func (r *retry) RuptOnContext(ctx context.Context, caller func() (bool, error)) error {
 	var lastErr error
+	var timer *time.Timer
+	var ctxDone bool
+
 	attempt := 1
 	for attempt <= r.attempts {
 		interrupted, err := caller()
@ -79,12 +123,54 @@ func (r *retry) RuptOn(caller func() (bool, error)) error {
 		if attempt >= r.attempts {
 			break
 		}
-		time.Sleep(time.Duration(r.nextDelay()) * time.Millisecond)
 		attempt++
+
+		if timer, ctxDone = r.waitContext(ctx, timer); ctxDone {
+			break
+		}
+	}
+	if timer != nil {
+		timer.Stop()
 	}
 	return lastErr
 }

+func (r *retry) waitContext(ctx context.Context, timer *time.Timer) (*time.Timer, bool) {
+	var wait <-chan time.Time
+
+	next := time.Duration(r.nextDelay()) * time.Millisecond
+	if _, internal := ctx.(internalContext); internal {
+		if next > 0 {
+			time.Sleep(next)
+		}
+		return timer, false
+	}
+
+	if next > 0 {
+		if timer == nil {
+			timer = time.NewTimer(next)
+		} else {
+			timer.Reset(next)
+		}
+		wait = timer.C
+	}
+
+	if wait != nil {
+		select {
+		case <-wait:
+			return timer, false
+		case <-ctx.Done():
+			return timer, true
+		}
+	}
+	select {
+	case <-ctx.Done():
+		return timer, true
+	default:
+		return timer, false
+	}
+}
+
 // Timed returns a retry with fixed interval delay.
 func Timed(attempts int, delay uint32) Retryer {
 	return &retry{
@ -100,6 +186,9 @@ func ExponentialBackoff(attempts int, expDelay uint32) Retryer {
 	next := expDelay
 	return &retry{
 		attempts: attempts,
+		reset: func() {
+			next = expDelay
+		},
 		nextDelay: func() uint32 {
 			r := next
 			next += expDelay
--- a/blobstore/util/retry/retry_test.go
+++ b/blobstore/util/retry/retry_test.go
@ -15,6 +15,7 @@
 package retry_test

 import (
+	"context"
 	"errors"
 	"testing"
 	"time"
@ -96,7 +97,8 @@ func TestRetryExhausted(t *testing.T) {
 func TestRetryExponentialBackoff(t *testing.T) {
 	st := time.Now()
 	called := 0
-	err := retry.ExponentialBackoff(7, 50).On(func() error {
+	r := retry.ExponentialBackoff(7, 50)
+	err := r.On(func() error {
 		called++
 		return errTestOnly
 	})
@ -104,8 +106,64 @@ func TestRetryExponentialBackoff(t *testing.T) {

 	require.ErrorIs(t, err, errTestOnly)
 	v := int64(duration / time.Millisecond)
+	require.Equal(t, 7, called)
 	require.Less(t, int64(1000), v, "duration: ", v)
 	require.Greater(t, int64(1250), v, "duration: ", v)
+
+	r.Reset()
+	st = time.Now()
+	err = r.On(func() error {
+		called++
+		return errTestOnly
+	})
+	duration = time.Since(st)
+
+	require.ErrorIs(t, err, errTestOnly)
+	v = int64(duration / time.Millisecond)
+	require.Equal(t, 14, called)
+	require.Less(t, int64(1000), v, "duration: ", v)
+	require.Greater(t, int64(1250), v, "duration: ", v)
+}
+
+func TestRetryContext(t *testing.T) {
+	{
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		err := retry.Timed(10, 0).OnContext(ctx, func() error {
+			return errTestOnly
+		})
+		require.ErrorIs(t, errTestOnly, err)
+	}
+	{
+		ctx, cancel := context.WithCancel(context.Background())
+		err := retry.Timed(10, 0).OnContext(ctx, func() error {
+			return errTestOnly
+		})
+		require.ErrorIs(t, errTestOnly, err)
+		cancel()
+	}
+	{
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		called := 0
+		err := retry.Timed(10, 400).OnContext(ctx, func() error {
+			called++
+			return errTestOnly
+		})
+		require.Equal(t, 3, called) // 0, 400, 800
+		require.ErrorIs(t, errTestOnly, err)
+		cancel()
+	}
+	{
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		called := 0
+		err := retry.ExponentialBackoff(10, 300).RuptOnContext(ctx, func() (bool, error) {
+			called++
+			return false, errTestOnly
+		})
+		require.Equal(t, 3, called) // 0, 300, 900
+		require.ErrorIs(t, errTestOnly, err)
+		cancel()
+	}
 }

 func TestRetryInterrupted(t *testing.T) {
@ -172,3 +230,40 @@ func TestRetryInterruptedError(t *testing.T) {
 	v = int64(duration / time.Millisecond)
 	require.Less(t, int64(190), v, "duration: ", v)
 }
+
+func BenchmarkWait(b *testing.B) {
+	b.Run("success", func(b *testing.B) {
+		r := retry.Timed(3, 1000)
+		fn := func() error { return nil }
+		for ii := 0; ii < b.N; ii++ {
+			r.On(fn)
+		}
+	})
+	b.Run("retry-3", func(b *testing.B) {
+		r := retry.Timed(3, 0)
+		fn := func() error { return errTestOnly }
+		for ii := 0; ii < b.N; ii++ {
+			r.On(fn)
+		}
+	})
+}
+
+func BenchmarkContext(b *testing.B) {
+	b.Run("success", func(b *testing.B) {
+		r := retry.Timed(3, 1000)
+		ctx := context.Background()
+		fn := func() error { return nil }
+		for ii := 0; ii < b.N; ii++ {
+			r.OnContext(ctx, fn)
+		}
+	})
+	b.Run("retry-3", func(b *testing.B) {
+		r := retry.Timed(3, 1000)
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		fn := func() error { return errTestOnly }
+		for ii := 0; ii < b.N; ii++ {
+			r.OnContext(ctx, fn)
+		}
+	})
+}