Preserve error when return from failed precheck

ammario · ammario · commit d12fd6bf4583 · 2018-08-14T11:11:37.000-05:00
Resolves #23
diff --git a/retry.go b/retry.go
@@ -3,6 +3,7 @@ package retry
 
 import (
 	"context"
+	"fmt"
 	"math/rand"
 	"time"
 
@@ -237,18 +238,50 @@ func (r *Retry) Log(logFn func(error)) *Retry {
 	})
 }
 
+// Error combines the reason the retry cancelled, and the error
+// from the last call to Run.
+type Error struct {
+	Reason  error
+	LastRun error
+}
+
+// nilIfEmpty returns a nil error if e is effectively nil, or itself.
+func (e Error) nilIfEmpty() error {
+	if e.Reason == nil && e.LastRun == nil {
+		return nil
+	}
+	return e
+}
+
+func (e Error) Error() string {
+	switch {
+	case e.Reason == nil:
+		return e.LastRun.Error()
+	case e.LastRun == nil:
+		return e.Reason.Error()
+	default:
+		return fmt.Sprintf("retry failed because %v, last run error: %v", e.Reason, e.LastRun)
+	}
+}
+
+// Cause returns the error from the last run.
+func (e Error) Cause() error {
+	return e.LastRun
+}
+
 // Run runs the retry.
 // The retry must not be ran twice.
 func (r *Retry) Run(fn func() error) error {
+	var e Error
 	for {
-		err := r.preCheck()
-		if err != nil {
-			return err
+		e.Reason = r.preCheck()
+		if e.Reason != nil {
+			return e.nilIfEmpty()
 		}
 
-		err = fn()
-		if !r.postCheck(err) {
-			return err
+		e.LastRun = fn()
+		if !r.postCheck(e.LastRun) {
+			return e.nilIfEmpty()
 		}
 		time.Sleep(r.sleepDur())
 	}
diff --git a/retry_test.go b/retry_test.go
@@ -80,13 +80,15 @@ func TestRetry(t *testing.T) {
 
 			const sleep = time.Millisecond * 10
 
-			New(sleep).Timeout(sleep * 5).Run(func() error {
+			var errSomething = errors.New("something")
+			err := New(sleep).Timeout(sleep * 5).Run(func() error {
 				count++
-				return errors.Errorf("asdfasdf")
+				return errSomething
 			})
 
 			assert.Equal(t, 5, count)
 			assert.WithinDuration(t, start.Add(sleep*5), time.Now(), sleep)
+			assert.Equal(t, errSomething, errors.Cause(err))
 		})
 
 		t.Run("returns as soon as error is nil", func(t *testing.T) {
@@ -122,7 +124,7 @@ func TestRetry(t *testing.T) {
 			})
 
 		assert.Equal(t, 5, count)
-		assert.Equal(t, io.ErrShortWrite, err)
+		assert.Equal(t, io.ErrShortWrite, errors.Cause(err))
 
 	})
 
@@ -143,7 +145,7 @@ func TestRetry(t *testing.T) {
 		})
 
 		assert.Equal(t, 3, count)
-		assert.Equal(t, context.Canceled, errors.Cause(err))
+		assert.Equal(t, io.EOF, errors.Cause(err))
 	})
 
 	t.Run("Jitter", func(t *testing.T) {
