From a2986ebb9c0cd23dd5a56b78698acb41cf87d58e Mon Sep 17 00:00:00 2001
From: Daniel Nephin <dnephin@hashicorp.com>
Date: Fri, 23 Apr 2021 17:01:05 -0400
Subject: [PATCH] sdk/retry: support ending the iteration early

I've found this feature to be very useful in https://pkg.go.dev/gotest.tools/v3/poll#WaitOn

I have encountered a few cases where I wanted that same support, so this commit adds it.
---
 sdk/testutil/retry/retry.go      | 31 ++++++++++++++++++++++++-------
 sdk/testutil/retry/retry_test.go | 21 +++++++++++++++++++++
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/sdk/testutil/retry/retry.go b/sdk/testutil/retry/retry.go
index b59bd1c4c3..e2b3efc1fc 100644
--- a/sdk/testutil/retry/retry.go
+++ b/sdk/testutil/retry/retry.go
@@ -34,6 +34,7 @@ type Failer interface {
 // R provides context for the retryer.
 type R struct {
 	fail   bool
+	done   bool
 	output []string
 }
 
@@ -77,6 +78,12 @@ func (r *R) log(s string) {
 	r.output = append(r.output, decorate(s))
 }
 
+// Stop retrying, and fail the test with the specified error.
+func (r *R) Stop(err error) {
+	r.log(err.Error())
+	r.done = true
+}
+
 func decorate(s string) string {
 	_, file, line, ok := runtime.Caller(3)
 	if ok {
@@ -120,6 +127,16 @@ func dedup(a []string) string {
 func run(r Retryer, t Failer, f func(r *R)) {
 	t.Helper()
 	rr := &R{}
+
+	fail := func() {
+		t.Helper()
+		out := dedup(rr.output)
+		if out != "" {
+			t.Log(out)
+		}
+		t.FailNow()
+	}
+
 	for r.Continue() {
 		func() {
 			defer func() {
@@ -129,17 +146,17 @@ func run(r Retryer, t Failer, f func(r *R)) {
 			}()
 			f(rr)
 		}()
-		if !rr.fail {
+
+		switch {
+		case rr.done:
+			fail()
+			return
+		case !rr.fail:
 			return
 		}
 		rr.fail = false
 	}
-
-	out := dedup(rr.output)
-	if out != "" {
-		t.Log(out)
-	}
-	t.FailNow()
+	fail()
 }
 
 // DefaultFailer provides default retry.Run() behavior for unit tests.
diff --git a/sdk/testutil/retry/retry_test.go b/sdk/testutil/retry/retry_test.go
index 95186374e0..31923a0bfb 100644
--- a/sdk/testutil/retry/retry_test.go
+++ b/sdk/testutil/retry/retry_test.go
@@ -1,6 +1,7 @@
 package retry
 
 import (
+	"fmt"
 	"testing"
 	"time"
 
@@ -52,15 +53,35 @@ func TestRunWith(t *testing.T) {
 		require.Equal(t, 3, iter)
 		require.Equal(t, 1, ft.fails)
 	})
+
+	t.Run("Stop ends the retrying", func(t *testing.T) {
+		ft := &fakeT{}
+		iter := 0
+		RunWith(&Counter{Count: 5, Wait: time.Millisecond}, ft, func(r *R) {
+			iter++
+			if iter == 2 {
+				r.Stop(fmt.Errorf("do not proceed"))
+			}
+			r.Fatalf("not yet")
+		})
+
+		require.Equal(t, 2, iter)
+		require.Equal(t, 1, ft.fails)
+		require.Len(t, ft.out, 1)
+		require.Contains(t, ft.out[0], "not yet\n")
+		require.Contains(t, ft.out[0], "do not proceed\n")
+	})
 }
 
 type fakeT struct {
 	fails int
+	out   []string
 }
 
 func (f *fakeT) Helper() {}
 
 func (f *fakeT) Log(args ...interface{}) {
+	f.out = append(f.out, fmt.Sprint(args...))
 }
 
 func (f *fakeT) FailNow() {