Merge #148551

148551: util: support max duration in retry loop r=stevendanna a=kev-cao

This commit teaches `Retry` to also respect a specified maximum duration. Once the maximum duration has elapsed since the `Start` of a retry loop, no more retry attempts will be made.

Informs: #148028

Release note: None

Co-authored-by: Kevin Cao <[email protected]>

Author: craig[bot]

pkg/kv/kvserver/asim/state/state.go

@@ -291,6 +291,10 @@ func (m *ManualSimClock) Since(t time.Time) time.Duration {
 	return m.Now().Sub(t)
 }
 
+func (m *ManualSimClock) Until(t time.Time) time.Duration {
+	return t.Sub(m.Now())
+}
+
 func (m *ManualSimClock) NewTimer() timeutil.TimerI {
 	panic("unimplemented")
 }

pkg/util/retry/BUILD.bazel

@@ -25,7 +25,9 @@ go_test(
     ],
     embed = [":retry"],
     deps = [
+        "//pkg/testutils/skip",
         "//pkg/util/log",
+        "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",
     ],

pkg/util/retry/retry.go

@@ -21,10 +21,22 @@ type Options struct {
 	InitialBackoff time.Duration // Default retry backoff interval
 	MaxBackoff     time.Duration // Maximum retry backoff interval
 	Multiplier     float64       // Default backoff constant
-	// Maximum number of retries; attempts = MaxRetries + 1. (0 for infinite)
-	MaxRetries          int
-	RandomizationFactor float64         // Randomize the backoff interval by constant
+	// Randomize the backoff interval by constant. Set to -1 to disable.
+	RandomizationFactor float64
 	Closer              <-chan struct{} // Optionally end retry loop channel close
+	// Maximum number of retries; attempts = MaxRetries + 1. (0 for infinite)
+	MaxRetries int
+	// MaxDuration is the maximum duration for which the retry loop will make
+	// attempts. Once the deadline has elapsed, the loop will stop attempting
+	// retries.
+	// The loop will run for at least one iteration. (0 for infinite)
+	MaxDuration time.Duration
+	// PreemptivelyCancel indicates whether the retry loop should cancel itself if
+	// it determines that the next backoff would exceed the MaxDuration.
+	PreemptivelyCancel bool
+	// Clock is used to control the time source for the retry loop. Intended for
+	// testing purposes. Should be nil in production code.
+	Clock timeutil.TimeSource
 }
 
 // Retry implements the public methods necessary to control an exponential-
@@ -34,6 +46,13 @@ type Retry struct {
 	ctx            context.Context
 	currentAttempt int
 	isReset        bool
+	deadline       time.Time // Deadline for the retry loop if MaxDuration is set.
+
+	// Testing hook that is called when the retry loop is waiting for the backoff.
+	// If no max duration is set, deadline will be zero. Passes in the actual wait
+	// time between each retry attempt (accounting for max duration).
+	// Set here instead of options to allow Options to be compared.
+	backingOffHook func(backoff time.Duration)
 }
 
 // Start returns a new Retry initialized to some default values. The Retry can
@@ -55,10 +74,15 @@ func StartWithCtx(ctx context.Context, opts Options) Retry {
 	}
 	if opts.RandomizationFactor == 0 {
 		opts.RandomizationFactor = 0.15
+	} else if opts.RandomizationFactor < 0 {
+		opts.RandomizationFactor = 0 // Disable randomization.
 	}
 	if opts.Multiplier == 0 {
 		opts.Multiplier = 2
 	}
+	if opts.Clock == nil {
+		opts.Clock = timeutil.DefaultTimeSource{}
+	}
 
 	var r Retry
 	r.opts = opts
@@ -87,8 +111,14 @@ func (r *Retry) Reset() {
 func (r *Retry) mustReset() {
 	r.currentAttempt = 0
 	r.isReset = true
+	if r.opts.MaxDuration != 0 {
+		r.deadline = r.opts.Clock.Now().Add(r.opts.MaxDuration)
+	} else {
+		r.deadline = time.Time{}
+	}
 }
 
+// retryIn returns the duration to wait before the next retry attempt.
 func (r Retry) retryIn() time.Duration {
 	backoff := float64(r.opts.InitialBackoff) * math.Pow(r.opts.Multiplier, float64(r.currentAttempt))
 	if maxBackoff := float64(r.opts.MaxBackoff); backoff > maxBackoff {
@@ -102,6 +132,23 @@ func (r Retry) retryIn() time.Duration {
 	return time.Duration(backoff - delta + rand.Float64()*(2*delta+1))
 }
 
+// calcDurationScopedBackoff calculates the duration to wait before the next
+// attempt, taking into account the MaxDuration option. It returns the computed
+// backoff duration, the actual wait duration (if the backoff exceeds the max
+// duration), and a boolean indicating whether the retry should be attempted.
+func (r Retry) calcDurationScopedBackoff() (time.Duration, time.Duration, bool) {
+	backoff := r.retryIn()
+	actualWait := backoff
+	shouldAttempt := true
+	if r.opts.MaxDuration != 0 && !r.opts.Clock.Now().Add(backoff).Before(r.deadline) {
+		// If the backoff would exceed the deadline, we return the remaining time
+		// until the deadline instead.
+		shouldAttempt = false
+		actualWait = max(r.deadline.Sub(r.opts.Clock.Now()), 0)
+	}
+	return backoff, actualWait, shouldAttempt
+}
+
 // Next returns whether the retry loop should continue, and blocks for the
 // appropriate length of time before yielding back to the caller.
 //
@@ -115,50 +162,107 @@ func (r *Retry) Next() bool {
 		r.isReset = false
 		return true
 	}
+	if r.retryLimitReached() {
+		return false
+	}
 
-	if r.opts.MaxRetries > 0 && r.currentAttempt >= r.opts.MaxRetries {
+	backoff, actualWait, shouldAttempt := r.calcDurationScopedBackoff()
+
+	if !shouldAttempt && r.opts.PreemptivelyCancel {
+		log.VEventf(
+			r.ctx, 2 /* level */, "preemptively canceling retry loop as backoff would exceed MaxDuration",
+		)
 		return false
 	}
 
-	// Wait before retry.
-	d := r.retryIn()
-	if d > 0 {
-		log.VEventfDepth(r.ctx, 1 /* depth */, 2 /* level */, "will retry after %s", d)
+	timer := r.opts.Clock.NewTimer()
+	timer.Reset(actualWait)
+	timerCh := timer.Ch()
+	defer timer.Stop()
+
+	log.VEventfDepth(r.ctx, 1 /* depth */, 2 /* level */, "will retry after %s", backoff)
+
+	if r.backingOffHook != nil {
+		r.backingOffHook(actualWait)
 	}
+
 	select {
-	case <-time.After(d):
-		r.currentAttempt++
-		return true
 	case <-r.opts.Closer:
 		return false
 	case <-r.ctx.Done():
 		return false
+	case <-timerCh:
+		if shouldAttempt {
+			r.currentAttempt++
+		}
+		return shouldAttempt
 	}
 }
 
-// closedC is returned from Retry.NextCh whenever a retry
-// can begin immediately.
-var closedC = func() chan time.Time {
-	c := make(chan time.Time)
+func (r *Retry) retryLimitReached() bool {
+	return (r.opts.MaxRetries > 0 && r.currentAttempt >= r.opts.MaxRetries) ||
+		(r.opts.MaxDuration > 0 && !r.opts.Clock.Now().Before(r.deadline))
+}
+
+// immediateCh creates a channel that is immediately written to with the
+// provided value and then closed.
+func immediateCh(v bool) chan bool {
+	c := make(chan bool, 1)
+	c <- v
 	close(c)
 	return c
-}()
+}
+
+var closedCh = immediateCh(false)
 
 // NextCh returns a channel which will receive when the next retry
-// interval has expired.
-func (r *Retry) NextCh() <-chan time.Time {
+// interval has expired. If the received value is true, it indicates a retry
+// should be made. If the received value is false, it indicates that no retry
+// should be made.
+// Note: This does not respect the Closer or context cancellation and it is the
+// caller's responsibility to manage the lifecycle.
+func (r *Retry) NextCh() <-chan bool {
 	if r.isReset {
 		r.isReset = false
-		return closedC
+		return immediateCh(true)
+	}
+	if r.retryLimitReached() {
+		return closedCh
+	}
+
+	backoff, actualWait, shouldAttempt := r.calcDurationScopedBackoff()
+
+	if !shouldAttempt && r.opts.PreemptivelyCancel {
+		log.VEventf(
+			r.ctx, 2 /* level */, "preemptively canceling retry loop as backoff would exceed MaxDuration",
+		)
+		return closedCh
+	}
+
+	timer := r.opts.Clock.NewTimer()
+	timer.Reset(actualWait)
+	timerCh := timer.Ch()
+
+	if r.backingOffHook != nil {
+		r.backingOffHook(actualWait)
 	}
-	r.currentAttempt++
-	if r.opts.MaxRetries > 0 && r.currentAttempt > r.opts.MaxRetries {
-		return nil
+
+	log.VEventfDepth(r.ctx, 1 /* depth */, 2 /* level */, "will retry after %s", backoff)
+
+	ch := make(chan bool, 1)
+	if shouldAttempt {
+		r.currentAttempt++
 	}
-	return time.After(r.retryIn())
+	go func() {
+		defer timer.Stop()
+		<-timerCh
+		ch <- shouldAttempt
+	}()
+
+	return ch
 }
 
-// CurrentAttempt returns the current attempt
+// CurrentAttempt returns the current attempt (0-based index)
 func (r *Retry) CurrentAttempt() int {
 	return r.currentAttempt
 }