Allow specifying custom retry behavior in retry middleware (#77)

Add an optional GetRetryParameters function member to RetryOptions. When defined, this will be called before each retry attempt and allow specifying whether to retry, the timeout, and a duration to wait before retrying.

Also add a simple unit test for the retry middleware.

pkg/middleware/retry.go

@@ -25,10 +25,11 @@ import (
 )
 
 type RetryOptions struct {
-	MaxAttempts   int
-	Timeout       time.Duration
-	Backoff       time.Duration
-	IsRecoverable func(err error) bool
+	MaxAttempts        int
+	Timeout            time.Duration
+	Backoff            time.Duration
+	IsRecoverable      func(err error) bool
+	GetRetryParameters func(err error, attempt int) (retry bool, timeout time.Duration, waitTime time.Duration) // will override the MaxAttempts, Timeout and Backoff parameters
 }
 
 func WithRPCRetries(opt RetryOptions) psrpc.ClientOption {
@@ -62,26 +63,51 @@ func isTimeout(err error) bool {
 	return e.Code() == psrpc.DeadlineExceeded || e.Code() == psrpc.Unavailable
 }
 
+func getRetryWithBackoffParameters(o RetryOptions) func(err error, attempt int) (retry bool, timeout time.Duration, waitTime time.Duration) {
+	timeout := o.Timeout
+
+	return func(err error, attempt int) (bool, time.Duration, time.Duration) {
+		if !o.IsRecoverable(err) || attempt == o.MaxAttempts {
+			return false, 0, 0
+		}
+
+		timeout += o.Backoff
+
+		return true, timeout, 0
+	}
+}
+
 func retry(opt RetryOptions, done <-chan struct{}, fn func(timeout time.Duration) error) error {
-	attempt := 1
 	timeout := opt.Timeout
+	attempt := 1
 	if opt.IsRecoverable == nil {
 		opt.IsRecoverable = isTimeout
 	}
 
+	if opt.GetRetryParameters == nil {
+		opt.GetRetryParameters = getRetryWithBackoffParameters(opt)
+	}
+
 	for {
 		err := fn(timeout)
-		if err == nil || !opt.IsRecoverable(err) || attempt == opt.MaxAttempts {
-			return err
+		if err == nil {
+			return nil
 		}
-		select {
-		case <-done:
-			return psrpc.ErrRequestCanceled
-		default:
+
+		var retry bool
+		var waitTime time.Duration
+		retry, timeout, waitTime = opt.GetRetryParameters(err, attempt)
+		if !retry {
+			return err
 		}
 
 		attempt++
-		timeout += opt.Backoff
+
+		select {
+		case <-done:
+			return psrpc.ErrRequestCanceled
+		case <-time.After(waitTime):
+		}
 	}
 }
 

pkg/middleware/retry_test.go

@@ -0,0 +1,118 @@
+package middleware
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/livekit/psrpc"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+)
+
+func TestRetryBackoff(t *testing.T) {
+	ro := RetryOptions{
+		MaxAttempts: 3,
+		Timeout:     100 * time.Millisecond,
+		Backoff:     200 * time.Millisecond,
+	}
+
+	var timeouts []time.Duration
+	getClientRpcHandler := func(errors []error) func(ctx context.Context, req proto.Message, opts ...psrpc.RequestOption) (res proto.Message, err error) {
+		timeouts = nil
+		attempt := 0
+
+		return func(ctx context.Context, req proto.Message, opts ...psrpc.RequestOption) (res proto.Message, err error) {
+			o := &psrpc.RequestOpts{}
+
+			for _, f := range opts {
+				f(o)
+			}
+
+			timeouts = append(timeouts, o.Timeout)
+
+			err = errors[attempt]
+			attempt++
+
+			return nil, err
+		}
+	}
+
+	t.Run("TestSuccess", func(t *testing.T) {
+		ro.IsRecoverable = func(err error) bool { return true }
+		ri := NewRPCRetryInterceptor(ro)
+
+		errs := []error{nil}
+		h := ri(psrpc.RPCInfo{}, getClientRpcHandler(errs))
+		h(context.Background(), nil)
+
+		require.Equal(t, 1, len(timeouts))
+		require.Equal(t, ro.Timeout, timeouts[0])
+	})
+
+	t.Run("TestFailureAllErrorsRetryable", func(t *testing.T) {
+		ro.IsRecoverable = func(err error) bool { return true }
+		ri := NewRPCRetryInterceptor(ro)
+
+		errs := make([]error, 3)
+		for i, _ := range errs {
+			errs[i] = errors.New("test error")
+		}
+		h := ri(psrpc.RPCInfo{}, getClientRpcHandler(errs))
+		h(context.Background(), nil)
+
+		require.Equal(t, ro.MaxAttempts, len(timeouts))
+
+		expectedTimeout := ro.Timeout
+		for _, timeout := range timeouts {
+			require.Equal(t, expectedTimeout, timeout)
+			expectedTimeout += ro.Backoff
+		}
+	})
+
+	t.Run("TestFailureNoErrorRetryable", func(t *testing.T) {
+		ro.IsRecoverable = func(err error) bool { return false }
+		ri := NewRPCRetryInterceptor(ro)
+
+		errs := []error{errors.New("test error")}
+		h := ri(psrpc.RPCInfo{}, getClientRpcHandler(errs))
+		h(context.Background(), nil)
+
+		require.Equal(t, 1, len(timeouts))
+		require.Equal(t, ro.Timeout, timeouts[0])
+	})
+
+	t.Run("TestCustomParameters", func(t *testing.T) {
+		lastTry := time.Now()
+
+		ro.GetRetryParameters = func(err error, attempt int) (retry bool, timeout time.Duration, waitTime time.Duration) {
+			if attempt > 1 {
+				now := time.Now()
+				require.InDelta(t, 500*time.Millisecond, now.Sub(lastTry), float64(20*time.Millisecond), "Retry didn't wait for required interval")
+				lastTry = now
+			}
+
+			if attempt == 3 {
+				return false, 0, 0
+			}
+
+			return true, 100 * time.Millisecond, 500 * time.Millisecond
+		}
+
+		ri := NewRPCRetryInterceptor(ro)
+
+		errs := make([]error, 3)
+		for i, _ := range errs {
+			errs[i] = errors.New("test error")
+		}
+		h := ri(psrpc.RPCInfo{}, getClientRpcHandler(errs))
+		h(context.Background(), nil)
+
+		require.Equal(t, ro.MaxAttempts, len(timeouts))
+
+		for _, timeout := range timeouts {
+			require.Equal(t, 100*time.Millisecond, timeout)
+		}
+	})
+}