From f5b85321c2baa4925a7492cf06a09d6ca19e14da Mon Sep 17 00:00:00 2001
From: michaeljguarino <mguarino46@gmail.com>
Date: Tue, 4 Apr 2023 01:40:00 -0400
Subject: [PATCH] Implement generic retry (#15)

Simple, pluggable retry system.  Supports constant and exponential retries for now
---
 algorithms/math.go  | 13 +++++++++++++
 retry/algorithms.go | 47 +++++++++++++++++++++++++++++++++++++++++++++
 retry/retry.go      | 24 +++++++++++++++++++++++
 retry/retry_test.go | 34 ++++++++++++++++++++++++++++++++
 4 files changed, 118 insertions(+)
 create mode 100644 algorithms/math.go
 create mode 100644 retry/algorithms.go
 create mode 100644 retry/retry.go
 create mode 100644 retry/retry_test.go

diff --git a/algorithms/math.go b/algorithms/math.go
new file mode 100644
index 0000000..9c96ed7
--- /dev/null
+++ b/algorithms/math.go
@@ -0,0 +1,13 @@
+package algorithms
+
+type Numeric interface {
+	int | int8 | int16 | int32 | int64 | float64 | float32
+}
+
+func Max[T Numeric](a, b T) T {
+	if a > b {
+		return a
+	}
+
+	return b
+}
diff --git a/retry/algorithms.go b/retry/algorithms.go
new file mode 100644
index 0000000..0c9a1d8
--- /dev/null
+++ b/retry/algorithms.go
@@ -0,0 +1,47 @@
+package retry
+
+import (
+	"time"
+
+	"github.com/pluralsh/polly/algorithms"
+)
+
+type BackoffAlgorithm interface {
+	Backoff(iter int) time.Duration
+	Continue() bool
+}
+
+type Exponential struct {
+	mult  float64
+	max   float64
+	start float64
+}
+
+func (exp *Exponential) Backoff(iter int) time.Duration {
+	dur := algorithms.Max(exp.start*exp.mult, exp.max)
+	exp.start = dur
+	return time.Duration(exp.start)
+}
+
+func (exp *Exponential) Continue() bool {
+	if exp.max <= exp.start {
+		return false
+	}
+
+	return true
+}
+
+type Constant struct {
+	max   int
+	dur   int
+	count int
+}
+
+func (con *Constant) Backoff(iter int) time.Duration {
+	con.count++
+	return time.Duration(con.dur)
+}
+
+func (con *Constant) Continue() bool {
+	return con.count < con.max
+}
diff --git a/retry/retry.go b/retry/retry.go
new file mode 100644
index 0000000..3e1943e
--- /dev/null
+++ b/retry/retry.go
@@ -0,0 +1,24 @@
+package retry
+
+import "time"
+
+func NewExponential(start, max int64, mult float64) *Exponential {
+	return &Exponential{start: float64(start), max: float64(max), mult: mult}
+}
+
+func NewConstant(duration, count int) *Constant {
+	return &Constant{dur: duration, max: count}
+}
+
+func Retry[T any](algo BackoffAlgorithm, fun func() (T, error)) (res T, err error) {
+	iter := 0
+	for {
+		res, err = fun()
+		if err != nil || !algo.Continue() {
+			return
+		}
+
+		iter++
+		time.Sleep(algo.Backoff(iter))
+	}
+}
diff --git a/retry/retry_test.go b/retry/retry_test.go
new file mode 100644
index 0000000..bfd2b29
--- /dev/null
+++ b/retry/retry_test.go
@@ -0,0 +1,34 @@
+package retry
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestConstantRetry(t *testing.T) {
+	algo := NewConstant(2, 5)
+	failer := func() (int, error) { return 0, fmt.Errorf("whoops") }
+
+	_, err := Retry(algo, failer)
+	assert.Error(t, err)
+
+	success := func() (int, error) { return 2, nil }
+	res, err := Retry(algo, success)
+	assert.NoError(t, err)
+	assert.Equal(t, res, 2)
+}
+
+func TestExponentialRetry(t *testing.T) {
+	algo := NewExponential(2, 5, 1.2)
+	failer := func() (int, error) { return 0, fmt.Errorf("whoops") }
+
+	_, err := Retry(algo, failer)
+	assert.Error(t, err)
+
+	success := func() (int, error) { return 2, nil }
+	res, err := Retry(algo, success)
+	assert.NoError(t, err)
+	assert.Equal(t, res, 2)
+}