From e62d389fb072e15ca4c6e3d42e283eff435c206b Mon Sep 17 00:00:00 2001
From: Clayton Coleman <ccoleman@redhat.com>
Date: Sat, 10 Nov 2018 11:05:30 -0500
Subject: [PATCH] util: Refactor Backoff to return the next step rather than
 sleeping

Allows consumers to use Backoff as a generator rather than have to
call ExponentialBackoff
---
 .../k8s.io/apimachinery/pkg/util/wait/wait.go | 62 ++++++++++++++-----
 .../apimachinery/pkg/util/wait/wait_test.go   | 45 ++++++++++++++
 2 files changed, 93 insertions(+), 14 deletions(-)

diff --git a/staging/src/k8s.io/apimachinery/pkg/util/wait/wait.go b/staging/src/k8s.io/apimachinery/pkg/util/wait/wait.go
index ca61168cd4a..590c17b4c59 100644
--- a/staging/src/k8s.io/apimachinery/pkg/util/wait/wait.go
+++ b/staging/src/k8s.io/apimachinery/pkg/util/wait/wait.go
@@ -173,10 +173,49 @@ type ConditionFunc func() (done bool, err error)
 
 // Backoff holds parameters applied to a Backoff function.
 type Backoff struct {
-	Duration time.Duration // the base duration
-	Factor   float64       // Duration is multiplied by factor each iteration
-	Jitter   float64       // The amount of jitter applied each iteration
-	Steps    int           // Exit with error after this many steps
+	// The initial duration.
+	Duration time.Duration
+	// Duration is multiplied by factor each iteration. Must be greater
+	// than or equal to zero.
+	Factor float64
+	// The amount of jitter applied each iteration. Jitter is applied after
+	// cap.
+	Jitter float64
+	// The number of steps before duration stops changing. If zero, initial
+	// duration is always used. Used for exponential backoff in combination
+	// with Factor.
+	Steps int
+	// The returned duration will never be greater than cap *before* jitter
+	// is applied. The actual maximum cap is `cap * (1.0 + jitter)`.
+	Cap time.Duration
+}
+
+// Step returns the next interval in the exponential backoff. This method
+// will mutate the provided backoff.
+func (b *Backoff) Step() time.Duration {
+	if b.Steps < 1 {
+		if b.Jitter > 0 {
+			return Jitter(b.Duration, b.Jitter)
+		}
+		return b.Duration
+	}
+	b.Steps--
+
+	duration := b.Duration
+
+	// calculate the next step
+	if b.Factor != 0 {
+		b.Duration = time.Duration(float64(b.Duration) * b.Factor)
+		if b.Cap > 0 && b.Duration > b.Cap {
+			b.Duration = b.Cap
+			b.Steps = 0
+		}
+	}
+
+	if b.Jitter > 0 {
+		duration = Jitter(duration, b.Jitter)
+	}
+	return duration
 }
 
 // ExponentialBackoff repeats a condition check with exponential backoff.
@@ -190,19 +229,14 @@ type Backoff struct {
 // If the condition never returns true, ErrWaitTimeout is returned. All other
 // errors terminate immediately.
 func ExponentialBackoff(backoff Backoff, condition ConditionFunc) error {
-	duration := backoff.Duration
-	for i := 0; i < backoff.Steps; i++ {
-		if i != 0 {
-			adjusted := duration
-			if backoff.Jitter > 0.0 {
-				adjusted = Jitter(duration, backoff.Jitter)
-			}
-			time.Sleep(adjusted)
-			duration = time.Duration(float64(duration) * backoff.Factor)
-		}
+	for backoff.Steps > 0 {
 		if ok, err := condition(); err != nil || ok {
 			return err
 		}
+		if backoff.Steps == 1 {
+			break
+		}
+		time.Sleep(backoff.Step())
 	}
 	return ErrWaitTimeout
 }
diff --git a/staging/src/k8s.io/apimachinery/pkg/util/wait/wait_test.go b/staging/src/k8s.io/apimachinery/pkg/util/wait/wait_test.go
index 2dfd2877756..24073bb1922 100644
--- a/staging/src/k8s.io/apimachinery/pkg/util/wait/wait_test.go
+++ b/staging/src/k8s.io/apimachinery/pkg/util/wait/wait_test.go
@@ -19,6 +19,7 @@ package wait
 import (
 	"errors"
 	"fmt"
+	"math/rand"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -499,3 +500,47 @@ func TestPollUntil(t *testing.T) {
 	// make sure we finished the poll
 	<-pollDone
 }
+
+func TestBackoff_Step(t *testing.T) {
+	tests := []struct {
+		initial *Backoff
+		want    []time.Duration
+	}{
+		{initial: &Backoff{Duration: time.Second, Steps: 0}, want: []time.Duration{time.Second, time.Second, time.Second}},
+		{initial: &Backoff{Duration: time.Second, Steps: 1}, want: []time.Duration{time.Second, time.Second, time.Second}},
+		{initial: &Backoff{Duration: time.Second, Factor: 1.0, Steps: 1}, want: []time.Duration{time.Second, time.Second, time.Second}},
+		{initial: &Backoff{Duration: time.Second, Factor: 2, Steps: 3}, want: []time.Duration{1 * time.Second, 2 * time.Second, 4 * time.Second}},
+		{initial: &Backoff{Duration: time.Second, Factor: 2, Steps: 3, Cap: 3 * time.Second}, want: []time.Duration{1 * time.Second, 2 * time.Second, 3 * time.Second}},
+		{initial: &Backoff{Duration: time.Second, Factor: 2, Steps: 2, Cap: 3 * time.Second, Jitter: 0.5}, want: []time.Duration{2 * time.Second, 3 * time.Second, 3 * time.Second}},
+		{initial: &Backoff{Duration: time.Second, Factor: 2, Steps: 6, Jitter: 4}, want: []time.Duration{1 * time.Second, 2 * time.Second, 4 * time.Second, 8 * time.Second, 16 * time.Second, 32 * time.Second}},
+	}
+	for seed := int64(0); seed < 5; seed++ {
+		for _, tt := range tests {
+			initial := *tt.initial
+			t.Run(fmt.Sprintf("%#v seed=%d", initial, seed), func(t *testing.T) {
+				rand.Seed(seed)
+				for i := 0; i < len(tt.want); i++ {
+					got := initial.Step()
+					t.Logf("[%d]=%s", i, got)
+					if initial.Jitter > 0 {
+						if got == tt.want[i] {
+							// this is statistically unlikely to happen by chance
+							t.Errorf("Backoff.Step(%d) = %v, no jitter", i, got)
+							continue
+						}
+						diff := float64(tt.want[i]-got) / float64(tt.want[i])
+						if diff > initial.Jitter {
+							t.Errorf("Backoff.Step(%d) = %v, want %v, outside range", i, got, tt.want)
+							continue
+						}
+					} else {
+						if got != tt.want[i] {
+							t.Errorf("Backoff.Step(%d) = %v, want %v", i, got, tt.want)
+							continue
+						}
+					}
+				}
+			})
+		}
+	}
+}