From 5133ad83b186ff140835cbae7d835b44321eeacc Mon Sep 17 00:00:00 2001
From: Anna Levenberg <annarose.levenberg@gmail.com>
Date: Fri, 10 Nov 2023 01:12:20 -0500
Subject: [PATCH] impl: add a retry with result function (#2837)

* impl: add a retry with result function

* fix ci errs
---
 pkg/util/util.go      | 21 +++++++++++++++++++++
 pkg/util/util_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/pkg/util/util.go b/pkg/util/util.go
index daba3e899..5303b0ea9 100644
--- a/pkg/util/util.go
+++ b/pkg/util/util.go
@@ -20,6 +20,7 @@ import (
 	"crypto/md5"
 	"crypto/sha256"
 	"encoding/hex"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"math"
@@ -197,6 +198,26 @@ func Retry(operation retryFunc, retryCount int, initialDelayMilliseconds int) er
 	return err
 }
 
+// Retry retries an operation with a return value
+func RetryWithResult[T any](operation func() (T, error), retryCount int, initialDelayMilliseconds int) (result T, err error) {
+	result, err = operation()
+	if err == nil {
+		return result, nil
+	}
+	for i := 0; i < retryCount; i++ {
+		sleepDuration := time.Millisecond * time.Duration(int(math.Pow(2, float64(i)))*initialDelayMilliseconds)
+		logrus.Warnf("Retrying operation after %s due to %v", sleepDuration, err)
+		time.Sleep(sleepDuration)
+
+		result, err = operation()
+		if err == nil {
+			return result, nil
+		}
+	}
+
+	return result, fmt.Errorf("unable to complete operation after %d attempts, last error: %w", retryCount, err)
+}
+
 func Lgetxattr(path string, attr string) ([]byte, error) {
 	// Start with a 128 length byte array
 	dest := make([]byte, 128)
diff --git a/pkg/util/util_test.go b/pkg/util/util_test.go
index 95cec8c78..59d445e32 100644
--- a/pkg/util/util_test.go
+++ b/pkg/util/util_test.go
@@ -64,3 +64,41 @@ func TestRetry(t *testing.T) {
 		t.Fatalf("Not expecting error: %v", err)
 	}
 }
+
+func makeRetryFuncWithResult(numFailures int) func() (int, error) {
+	i := -1
+
+	return func() (int, error) {
+		i++
+		if i < numFailures {
+			return i, fmt.Errorf("Failing with i=%v", i)
+		}
+		return i, nil
+	}
+}
+
+func TestRetryWithResult(t *testing.T) {
+	// test with a function that does not return an error
+	result, err := RetryWithResult(makeRetryFuncWithResult(0), 0, 10)
+	if err != nil || result != 0 {
+		t.Fatalf("Got result %d and error: %v", result, err)
+	}
+	result, err = RetryWithResult(makeRetryFuncWithResult(0), 3, 10)
+	if err != nil || result != 0 {
+		t.Fatalf("Got result %d and error: %v", result, err)
+	}
+
+	// test with a function that returns an error twice
+	result, err = RetryWithResult(makeRetryFuncWithResult(2), 0, 10)
+	if err == nil || result != 0 {
+		t.Fatalf("Got result %d and error: %v", result, err)
+	}
+	result, err = RetryWithResult(makeRetryFuncWithResult(2), 1, 10)
+	if err == nil || result != 1 {
+		t.Fatalf("Got result %d and error: %v", result, err)
+	}
+	result, err = RetryWithResult(makeRetryFuncWithResult(2), 2, 10)
+	if err != nil || result != 2 {
+		t.Fatalf("Got result %d and error: %v", result, err)
+	}
+}