Merge pull request avast#91 from craigpastro/do-with-data

Add DoWithData function

Author: JaSei

.github/workflows/workflow.yaml

@@ -1,14 +1,18 @@
 name: Go
 
-on: [push]
+on:
+  push:
+  pull_request:
+    branches:
+      - master
 
 jobs:
   golangci-lint:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/setup-go@v3
+      - uses: actions/setup-go@v4
         with:
-          go-version: 1.17
+          go-version: 1.18
       - uses: actions/checkout@v3
       - name: golangci-lint
         uses: golangci/golangci-lint-action@v3
@@ -27,7 +31,7 @@ jobs:
     steps:
       - uses: actions/checkout@v3
       - name: Setup Go ${{ matrix.go-version }}
-        uses: actions/setup-go@v3
+        uses: actions/setup-go@v4
         with:
           go-version: ${{ matrix.go-version }}
           check-latest: true

README.md

@@ -35,8 +35,36 @@ http get with retry:
     		return nil
     	},
     )
+    if err != nil {
+    	// handle error
+    }
+
+    fmt.Println(string(body))
 
-    fmt.Println(body)
+http get with retry with data:
+
+    url := "http://example.com"
+
+    body, err := retry.DoWithData(
+    	func() ([]byte, error) {
+    		resp, err := http.Get(url)
+    		if err != nil {
+    			return nil, err
+    		}
+    		defer resp.Body.Close()
+    		body, err := ioutil.ReadAll(resp.Body)
+    		if err != nil {
+    			return nil, err
+    		}
+
+    		return body, nil
+    	},
+    )
+    if err != nil {
+    	// handle error
+    }
+
+    fmt.Println(string(body))
 
 [next examples](https://github.com/avast/retry-go/tree/master/examples)
 
@@ -94,6 +122,12 @@ BackOffDelay is a DelayType which increases delay between consecutive retries
 func Do(retryableFunc RetryableFunc, opts ...Option) error
 ```
 
+#### func  DoWithData
+
+```go
+func DoWithData[T any](retryableFunc RetryableFuncWithData[T], opts ...Option) (T, error)
+```
+
 #### func  FixedDelay
 
 ```go
@@ -383,6 +417,14 @@ type RetryableFunc func() error
 
 Function signature of retryable function
 
+#### type RetryableFuncWithData
+
+```go
+type RetryableFuncWithData[T any] func() (T, error)
+```
+
+Function signature of retryable function with data
+
 #### type Timer
 
 ```go

go.mod

@@ -1,5 +1,11 @@
 module github.com/avast/retry-go/v4
 
-go 1.13
+go 1.18
 
 require github.com/stretchr/testify v1.8.2
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

retry.go

@@ -25,8 +25,36 @@ http get with retry:
 			return nil
 		},
 	)
+	if err != nil {
+		// handle error
+	}
+
+	fmt.Println(string(body))
+
+http get with retry with data:
+
+	url := "http://example.com"
+
+	body, err := retry.DoWithData(
+		func() ([]byte, error) {
+			resp, err := http.Get(url)
+			if err != nil {
+				return nil, err
+			}
+			defer resp.Body.Close()
+			body, err := ioutil.ReadAll(resp.Body)
+			if err != nil {
+				return nil, err
+			}
+
+			return body, nil
+		},
+	)
+	if err != nil {
+		// handle error
+	}
 
-	fmt.Println(body)
+	fmt.Println(string(body))
 
 [next examples](https://github.com/avast/retry-go/tree/master/examples)
 
@@ -72,6 +100,9 @@ import (
 // Function signature of retryable function
 type RetryableFunc func() error
 
+// Function signature of retryable function with data
+type RetryableFuncWithData[T any] func() (T, error)
+
 // Default timer is a wrapper around time.After
 type timerImpl struct{}
 
@@ -80,7 +111,17 @@ func (t *timerImpl) After(d time.Duration) <-chan time.Time {
 }
 
 func Do(retryableFunc RetryableFunc, opts ...Option) error {
+	retryableFuncWithData := func() (any, error) {
+		return nil, retryableFunc()
+	}
+
+	_, err := DoWithData(retryableFuncWithData, opts...)
+	return err
+}
+
+func DoWithData[T any](retryableFunc RetryableFuncWithData[T], opts ...Option) (T, error) {
 	var n uint
+	var emptyT T
 
 	// default
 	config := newDefaultRetryConfig()
@@ -91,30 +132,33 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 	}
 
 	if err := config.context.Err(); err != nil {
-		return err
+		return emptyT, err
 	}
 
 	// Setting attempts to 0 means we'll retry until we succeed
 	if config.attempts == 0 {
-		for err := retryableFunc(); err != nil; err = retryableFunc() {
+		for {
+			t, err := retryableFunc()
+			if err == nil {
+				return t, nil
+			}
+
 			if !IsRecoverable(err) {
-				return err
+				return emptyT, err
 			}
 
 			if !config.retryIf(err) {
-				return err
+				return emptyT, err
 			}
 
 			n++
 			config.onRetry(n, err)
 			select {
 			case <-config.timer.After(delay(config, n, err)):
 			case <-config.context.Done():
-				return config.context.Err()
+				return emptyT, config.context.Err()
 			}
 		}
-
-		return nil
 	}
 
 	errorLog := Error{}
@@ -126,9 +170,9 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 
 	shouldRetry := true
 	for shouldRetry {
-		err := retryableFunc()
+		t, err := retryableFunc()
 		if err == nil {
-			return nil
+			return t, nil
 		}
 
 		errorLog = append(errorLog, unpackUnrecoverable(err))
@@ -156,22 +200,20 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 		case <-config.timer.After(delay(config, n, err)):
 		case <-config.context.Done():
 			if config.lastErrorOnly {
-				return config.context.Err()
+				return emptyT, config.context.Err()
 			}
-			n++
 
-			return append(errorLog, config.context.Err())
+			return emptyT, append(errorLog, config.context.Err())
 		}
 
 		n++
 		shouldRetry = shouldRetry && n < config.attempts
 	}
 
 	if config.lastErrorOnly {
-		return errorLog.Unwrap()
+		return emptyT, errorLog.Unwrap()
 	}
-
-	return errorLog
+	return emptyT, errorLog
 }
 
 func newDefaultRetryConfig() *Config {

retry_test.go

@@ -11,14 +11,15 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestDoAllFailed(t *testing.T) {
+func TestDoWithDataAllFailed(t *testing.T) {
 	var retrySum uint
-	err := Do(
-		func() error { return errors.New("test") },
+	v, err := DoWithData(
+		func() (int, error) { return 7, errors.New("test") },
 		OnRetry(func(n uint, err error) { retrySum += n }),
 		Delay(time.Nanosecond),
 	)
 	assert.Error(t, err)
+	assert.Equal(t, 0, v)
 
 	expectedErrorFormat := `All attempts fail:
 #1: test
@@ -44,7 +45,19 @@ func TestDoFirstOk(t *testing.T) {
 	)
 	assert.NoError(t, err)
 	assert.Equal(t, uint(0), retrySum, "no retry")
+}
+
+func TestDoWithDataFirstOk(t *testing.T) {
+	returnVal := 1
 
+	var retrySum uint
+	val, err := DoWithData(
+		func() (int, error) { return returnVal, nil },
+		OnRetry(func(n uint, err error) { retrySum += n }),
+	)
+	assert.NoError(t, err)
+	assert.Equal(t, returnVal, val)
+	assert.Equal(t, uint(0), retrySum, "no retry")
 }
 
 func TestRetryIf(t *testing.T) {
@@ -530,7 +543,7 @@ func BenchmarkDo(b *testing.B) {
 	testError := errors.New("test error")
 
 	for i := 0; i < b.N; i++ {
-		Do(
+		_ = Do(
 			func() error {
 				return testError
 			},
@@ -540,9 +553,23 @@ func BenchmarkDo(b *testing.B) {
 	}
 }
 
+func BenchmarkDoWithData(b *testing.B) {
+	testError := errors.New("test error")
+
+	for i := 0; i < b.N; i++ {
+		_, _ = DoWithData(
+			func() (int, error) {
+				return 0, testError
+			},
+			Attempts(10),
+			Delay(0),
+		)
+	}
+}
+
 func BenchmarkDoNoErrors(b *testing.B) {
 	for i := 0; i < b.N; i++ {
-		Do(
+		_ = Do(
 			func() error {
 				return nil
 			},
@@ -552,6 +579,18 @@ func BenchmarkDoNoErrors(b *testing.B) {
 	}
 }
 
+func BenchmarkDoWithDataNoErrors(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		_, _ = DoWithData(
+			func() (int, error) {
+				return 0, nil
+			},
+			Attempts(10),
+			Delay(0),
+		)
+	}
+}
+
 func TestIsRecoverable(t *testing.T) {
 	err := errors.New("err")
 	assert.True(t, IsRecoverable(err))