From 25424a72a88a26d679553e791a93b3f111a1e589 Mon Sep 17 00:00:00 2001
From: Dustin Decker <dustin.decker@getcruise.com>
Date: Tue, 12 May 2020 15:09:38 -0700
Subject: [PATCH] Add sleep option for buffer retries

Signed-off-by: Dustin Decker <dustin.decker@getcruise.com>
---
 buffer/buffer.go     | 29 ++++++++++++++++++++++++++---
 buffer/retry_test.go |  9 ++++++---
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/buffer/buffer.go b/buffer/buffer.go
index 404d7996..3ecb043d 100644
--- a/buffer/buffer.go
+++ b/buffer/buffer.go
@@ -26,11 +26,14 @@ Examples of a buffering middleware:
   // Will do the same as above, but with responses
   buffer.New(handler,
     buffer.MemResponseBodyBytes(2 * 1024 * 1024),
-    buffer.MaxResponseBodyBytes(10 * 1024 * 1024))
+	buffer.MaxResponseBodyBytes(10 * 1024 * 1024))
+
 
   // Buffer will replay the request if the handler returns error at least 3 times
-  // before returning the response
-  buffer.New(handler, buffer.Retry(`IsNetworkError() && Attempts() <= 2`))
+  // before returning the response, sleeping for 1 second in between attempts
+  buffer.New(handler,
+	buffer.Retry(`IsNetworkError() && Attempts() <= 2`),
+	buffer.RetrySleep(time.Second))
 
 */
 package buffer
@@ -43,6 +46,7 @@ import (
 	"net"
 	"net/http"
 	"reflect"
+	"time"
 
 	"github.com/mailgun/multibuf"
 	log "github.com/sirupsen/logrus"
@@ -56,6 +60,8 @@ const (
 	DefaultMaxBodyBytes = -1
 	// DefaultMaxRetryAttempts Maximum retry attempts
 	DefaultMaxRetryAttempts = 10
+	// DefaultRetrySleep specifies the default sleep between retries
+	DefaultRetrySleep = time.Second * 0
 )
 
 var errHandler utils.ErrorHandler = &SizeErrHandler{}
@@ -69,6 +75,8 @@ type Buffer struct {
 	maxResponseBodyBytes int64
 	memResponseBodyBytes int64
 
+	retrySleep time.Duration
+
 	retryPredicate hpredicate
 
 	next       http.Handler
@@ -88,6 +96,8 @@ func New(next http.Handler, setters ...optSetter) (*Buffer, error) {
 		maxResponseBodyBytes: DefaultMaxBodyBytes,
 		memResponseBodyBytes: DefaultMemBodyBytes,
 
+		retrySleep: DefaultRetrySleep,
+
 		log: log.StandardLogger(),
 	}
 	for _, s := range setters {
@@ -202,6 +212,17 @@ func MemResponseBodyBytes(m int64) optSetter {
 	}
 }
 
+// RetrySleep sets sleep duration between retries
+func RetrySleep(t time.Duration) optSetter {
+	return func(b *Buffer) error {
+		if t < 0 {
+			return fmt.Errorf("sleep duration should be >= 0 got %s", t)
+		}
+		b.retrySleep = t
+		return nil
+	}
+}
+
 // Wrap sets the next handler to be called by buffer handler.
 func (b *Buffer) Wrap(next http.Handler) error {
 	b.next = next
@@ -317,6 +338,8 @@ func (b *Buffer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 
 		outreq = b.copyRequest(req, body, totalSize)
 		b.log.Debugf("vulcand/oxy/buffer: retry Request(%v %v) attempt %v", req.Method, req.URL, attempt)
+
+		time.Sleep(b.retrySleep)
 	}
 }
 
diff --git a/buffer/retry_test.go b/buffer/retry_test.go
index bb2880e3..62cf7771 100644
--- a/buffer/retry_test.go
+++ b/buffer/retry_test.go
@@ -4,6 +4,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -37,7 +38,7 @@ func TestRetryOnError(t *testing.T) {
 	})
 	defer srv.Close()
 
-	lb, rt := newBufferMiddleware(t, `IsNetworkError() && Attempts() <= 2`)
+	lb, rt := newBufferMiddleware(t, `IsNetworkError() && Attempts() <= 2`, RetrySleep(time.Second))
 
 	proxy := httptest.NewServer(rt)
 	defer proxy.Close()
@@ -72,7 +73,7 @@ func TestRetryExceedAttempts(t *testing.T) {
 	assert.Equal(t, http.StatusBadGateway, re.StatusCode)
 }
 
-func newBufferMiddleware(t *testing.T, p string) (*roundrobin.RoundRobin, *Buffer) {
+func newBufferMiddleware(t *testing.T, p string, setters ...optSetter) (*roundrobin.RoundRobin, *Buffer) {
 	// forwarder will proxy the request to whatever destination
 	fwd, err := forward.New()
 	require.NoError(t, err)
@@ -81,8 +82,10 @@ func newBufferMiddleware(t *testing.T, p string) (*roundrobin.RoundRobin, *Buffe
 	lb, err := roundrobin.New(fwd)
 	require.NoError(t, err)
 
+	setters = append(setters, Retry(p), MemRequestBodyBytes(1))
+
 	// stream handler will forward requests to redirect, make sure it uses files
-	st, err := New(lb, Retry(p), MemRequestBodyBytes(1))
+	st, err := New(lb, setters...)
 	require.NoError(t, err)
 
 	return lb, st