From 5de0a452cf47e3a232a97dacd4ef42661a837c68 Mon Sep 17 00:00:00 2001
From: Daniel Nephin <dnephin@hashicorp.com>
Date: Wed, 21 Apr 2021 17:04:38 -0400
Subject: [PATCH] sdk/retry: a few small debug improvements

On a few occasions I've had to read timeout stack traces for tests and
noticed that retry.Run runs the function in a goroutine. This makes
debuging a timeout more difficult because the gourinte of the retryable
function is disconnected from the stack of the actual test. It requires
searching through the entire stack trace to find the other goroutine.

By using panic instead of runtime.Goexit() we remove the need for a
separate goroutine.

Also a few other small improvements:
* add `R.Helper` so that an assertion function can be used with both
  testing.T and retry.R.
* Pass t to `Retryer.NextOr`, and call `t.Helper` in a number of places
  so that the line number reported by `t.Log` is the line in the test
  where `retry.Run` was called, instead of some line in `retry.go` that
  is not relevant to the failure.
* improve the implementation of `dedup` by removing the need to iterate
  twice. Instad track the lines and skip any duplicate when writing to
  the buffer.
---
 sdk/testutil/retry/retry.go      | 51 ++++++++++++++++++--------------
 sdk/testutil/retry/retry_test.go |  2 +-
 2 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/sdk/testutil/retry/retry.go b/sdk/testutil/retry/retry.go
index 603c47346a..09f845abe9 100644
--- a/sdk/testutil/retry/retry.go
+++ b/sdk/testutil/retry/retry.go
@@ -17,7 +17,6 @@ import (
 	"fmt"
 	"runtime"
 	"strings"
-	"sync"
 	"time"
 )
 
@@ -38,9 +37,13 @@ type R struct {
 	output []string
 }
 
+func (r *R) Helper() {}
+
+var runFailed = struct{}{}
+
 func (r *R) FailNow() {
 	r.fail = true
-	runtime.Goexit()
+	panic(runFailed)
 }
 
 func (r *R) Fatal(args ...interface{}) {
@@ -93,6 +96,7 @@ func Run(t Failer, f func(r *R)) {
 }
 
 func RunWith(r Retryer, t Failer, f func(r *R)) {
+	t.Helper()
 	run(r, t, f)
 }
 
@@ -100,22 +104,21 @@ func dedup(a []string) string {
 	if len(a) == 0 {
 		return ""
 	}
-	m := map[string]int{}
-	for _, s := range a {
-		m[s] = m[s] + 1
-	}
+	seen := map[string]struct{}{}
 	var b bytes.Buffer
 	for _, s := range a {
-		if _, ok := m[s]; ok {
-			b.WriteString(s)
-			b.WriteRune('\n')
-			delete(m, s)
+		if _, ok := seen[s]; ok {
+			continue
 		}
+		seen[s] = struct{}{}
+		b.WriteString(s)
+		b.WriteRune('\n')
 	}
 	return b.String()
 }
 
 func run(r Retryer, t Failer, f func(r *R)) {
+	t.Helper()
 	rr := &R{}
 	fail := func() {
 		t.Helper()
@@ -125,19 +128,19 @@ func run(r Retryer, t Failer, f func(r *R)) {
 		}
 		t.FailNow()
 	}
-	for r.NextOr(fail) {
-		var wg sync.WaitGroup
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
+	for r.NextOr(t, fail) {
+		func() {
+			defer func() {
+				if p := recover(); p != nil && p != runFailed {
+					panic(p)
+				}
+			}()
 			f(rr)
 		}()
-		wg.Wait()
-		if rr.fail {
-			rr.fail = false
-			continue
+		if !rr.fail {
+			return
 		}
-		break
+		rr.fail = false
 	}
 }
 
@@ -161,7 +164,7 @@ func ThreeTimes() *Counter {
 type Retryer interface {
 	// NextOr returns true if the operation should be repeated.
 	// Otherwise, it calls fail and returns false.
-	NextOr(fail func()) bool
+	NextOr(t Failer, fail func()) bool
 }
 
 // Counter repeats an operation a given number of
@@ -173,7 +176,8 @@ type Counter struct {
 	count int
 }
 
-func (r *Counter) NextOr(fail func()) bool {
+func (r *Counter) NextOr(t Failer, fail func()) bool {
+	t.Helper()
 	if r.count == r.Count {
 		fail()
 		return false
@@ -196,7 +200,8 @@ type Timer struct {
 	stop time.Time
 }
 
-func (r *Timer) NextOr(fail func()) bool {
+func (r *Timer) NextOr(t Failer, fail func()) bool {
+	t.Helper()
 	if r.stop.IsZero() {
 		r.stop = time.Now().Add(r.Timeout)
 		return true
diff --git a/sdk/testutil/retry/retry_test.go b/sdk/testutil/retry/retry_test.go
index db9c4aa84d..f58f8eb92f 100644
--- a/sdk/testutil/retry/retry_test.go
+++ b/sdk/testutil/retry/retry_test.go
@@ -22,7 +22,7 @@ func TestRetryer(t *testing.T) {
 			var iters, fails int
 			fail := func() { fails++ }
 			start := time.Now()
-			for tt.r.NextOr(fail) {
+			for tt.r.NextOr(t, fail) {
 				iters++
 			}
 			dur := time.Since(start)