Support retriable func condition

interceptors/retry/options.go

@@ -41,6 +41,9 @@ type BackoffFunc func(ctx context.Context, attempt uint) time.Duration
 // OnRetryCallback is the type of function called when a retry occurs.
 type OnRetryCallback func(ctx context.Context, attempt uint, err error)
 
+// RetriableFunc denotes a family of functions that control which error should be retried.
+type RetriableFunc func(err error) bool
+
 // Disable disables the retry behaviour on this call, or this interceptor.
 //
 // Its semantically the same to `WithMax`
@@ -100,13 +103,21 @@ func WithPerRetryTimeout(timeout time.Duration) CallOption {
 	}}
 }
 
+// WithRetriable sets which error should be retried.
+func WithRetriable(retriableFunc RetriableFunc) CallOption {
+	return CallOption{applyFunc: func(o *options) {
+		o.retriableFunc = retriableFunc
+	}}
+}
+
 type options struct {
 	max             uint
 	perCallTimeout  time.Duration
 	includeHeader   bool
 	codes           []codes.Code
 	backoffFunc     BackoffFunc
 	onRetryCallback OnRetryCallback
+	retriableFunc   RetriableFunc
 }
 
 // CallOption is a grpc.CallOption that is local to grpc_retry.

interceptors/retry/retry.go

@@ -267,6 +267,9 @@ func waitRetryBackoff(attempt uint, parentCtx context.Context, callOpts *options
 }
 
 func isRetriable(err error, callOpts *options) bool {
+	if callOpts.retriableFunc != nil {
+		return callOpts.retriableFunc(err)
+	}
 	errCode := status.Code(err)
 	if isContextError(err) {
 		// context errors are not retriable based on user settings.

interceptors/retry/retry_test.go

@@ -6,6 +6,7 @@ package retry
 import (
 	"context"
 	"io"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -178,6 +179,16 @@ func (s *RetrySuite) TestUnary_OverrideFromDialOpts() {
 	require.EqualValues(s.T(), 5, s.srv.requestCount(), "five requests should have been made")
 }
 
+func (s *RetrySuite) TestUnary_OverrideFromDialOpts2() {
+	s.srv.resetFailingConfiguration(5, codes.ResourceExhausted, noSleep) // default is 3 and retriable_errors
+	out, err := s.Client.Ping(s.SimpleCtx(), testpb.GoodPing, WithRetriable(func(err error) bool {
+		return strings.Contains(err.Error(), "maybeFailRequest")
+	}), WithMax(5))
+	require.NoError(s.T(), err, "the fifth invocation should succeed")
+	require.NotNil(s.T(), out, "Pong must be not nil")
+	require.EqualValues(s.T(), 5, s.srv.requestCount(), "five requests should have been made")
+}
+
 func (s *RetrySuite) TestUnary_OnRetryCallbackCalled() {
 	retryCallbackCount := 0
 
@@ -209,6 +220,16 @@ func (s *RetrySuite) TestServerStream_OverrideFromContext() {
 	require.EqualValues(s.T(), 5, s.srv.requestCount(), "three requests should have been made")
 }
 
+func (s *RetrySuite) TestServerStream_OverrideFromContext2() {
+	s.srv.resetFailingConfiguration(5, codes.ResourceExhausted, noSleep) // default is 3 and retriable_errors
+	stream, err := s.Client.PingList(s.SimpleCtx(), testpb.GoodPingList, WithRetriable(func(err error) bool {
+		return strings.Contains(err.Error(), "maybeFailRequest")
+	}), WithMax(5))
+	require.NoError(s.T(), err, "establishing the connection must always succeed")
+	s.assertPingListWasCorrect(stream)
+	require.EqualValues(s.T(), 5, s.srv.requestCount(), "three requests should have been made")
+}
+
 func (s *RetrySuite) TestServerStream_OnRetryCallbackCalled() {
 	retryCallbackCount := 0