From 860a499f98b52cbb71c7d9602e4c83ed4ab3b213 Mon Sep 17 00:00:00 2001
From: donutloop <marcel.edmund.franke@gmail.com>
Date: Fri, 23 Feb 2024 03:04:38 +0100
Subject: [PATCH] Add custom backoff setter (#176)

Users should have the capability to customize the backoff pattern and accordingly adjust it within the retry mechanism.
---
 retry/retry.go              | 12 ++++++++++++
 retry/retry_example_test.go | 29 +++++++++++++++++++++++++++++
 retry/retry_test.go         | 28 ++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+)

diff --git a/retry/retry.go b/retry/retry.go
index 0210c76..7f3c766 100644
--- a/retry/retry.go
+++ b/retry/retry.go
@@ -44,6 +44,18 @@ func RetryTimes(n uint) Option {
 	}
 }
 
+// RetryWithCustomBackoff set abitary custom backoff strategy
+// todo: Add playground link
+func RetryWithCustomBackoff(backoffStrategy BackoffStrategy) Option {
+	if backoffStrategy == nil {
+		panic("programming error: backoffStrategy must be not nil")
+	}
+
+	return func(rc *RetryConfig) {
+		rc.backoffStrategy = backoffStrategy
+	}
+}
+
 // RetryWithLinearBackoff set linear strategy backoff
 // todo: Add playground link
 func RetryWithLinearBackoff(interval time.Duration) Option {
diff --git a/retry/retry_example_test.go b/retry/retry_example_test.go
index 81b26bb..3c9cbeb 100644
--- a/retry/retry_example_test.go
+++ b/retry/retry_example_test.go
@@ -51,6 +51,35 @@ func ExampleRetryWithLinearBackoff() {
 	// 3
 }
 
+type ExampleCustomBackoffStrategy struct {
+	interval time.Duration
+}
+
+func (c *ExampleCustomBackoffStrategy) CalculateInterval() time.Duration {
+	return c.interval + 1
+}
+
+func ExampleRetryWithCustomBackoff() {
+	number := 0
+	increaseNumber := func() error {
+		number++
+		if number == 3 {
+			return nil
+		}
+		return errors.New("error occurs")
+	}
+
+	err := Retry(increaseNumber, RetryWithCustomBackoff(&ExampleCustomBackoffStrategy{interval: time.Microsecond * 50}))
+	if err != nil {
+		return
+	}
+
+	fmt.Println(number)
+
+	// Output:
+	// 3
+}
+
 func ExampleRetryWithExponentialWithJitterBackoff() {
 	number := 0
 	increaseNumber := func() error {
diff --git a/retry/retry_test.go b/retry/retry_test.go
index f4d2e28..125c2bc 100644
--- a/retry/retry_test.go
+++ b/retry/retry_test.go
@@ -137,6 +137,34 @@ func TestRetryOneShotSucceeded(t *testing.T) {
 	assert.Equal(1, number)
 }
 
+func TestRetryWitCustomBackoffOneShotSucceeded(t *testing.T) {
+	t.Parallel()
+
+	assert := internal.NewAssert(t, "TestRetryWitCustomBackoffOneShotSucceeded")
+
+	var number int
+	increaseNumber := func() error {
+		number++
+		if number == DefaultRetryTimes {
+			return nil
+		}
+		return errors.New("error occurs")
+	}
+
+	err := Retry(increaseNumber, RetryWithCustomBackoff(&TestCustomBackoffStrategy{interval: time.Microsecond * 50}))
+
+	assert.IsNil(err)
+	assert.Equal(5, number)
+}
+
+type TestCustomBackoffStrategy struct {
+	interval time.Duration
+}
+
+func (c *TestCustomBackoffStrategy) CalculateInterval() time.Duration {
+	return c.interval + 1
+}
+
 func TestRetryWithExponentialWithJitterBackoffShiftOneShotSucceeded(t *testing.T) {
 	t.Parallel()