From d7e4062a0e86790b90f46c6d9a65fb654a73d121 Mon Sep 17 00:00:00 2001
From: Scott Finlay <ynori7@yahoo.com>
Date: Thu, 3 Nov 2016 10:45:31 +0100
Subject: [PATCH] Adding the ability to specify a function to check if micro
 should retry a failed rpc call

---
 client/client.go       |  2 ++
 client/is_retriable.go | 22 ++++++++++++++++++++++
 client/options.go      | 19 +++++++++++++++----
 client/rpc_client.go   |  5 +++++
 4 files changed, 44 insertions(+), 4 deletions(-)
 create mode 100644 client/is_retriable.go

diff --git a/client/client.go b/client/client.go
index 723af2e0..5e8fe23d 100644
--- a/client/client.go
+++ b/client/client.go
@@ -68,6 +68,8 @@ var (
 	DefaultClient Client = newRpcClient()
 	// DefaultBackoff is the default backoff function for retries
 	DefaultBackoff = exponentialBackoff
+	// DefaultCheckIfRetriable is the default check-for-retry function for retries
+	DefaultCheckIfRetriable = AlwaysRetry
 	// DefaultRetries is the default number of times a request is tried
 	DefaultRetries = 1
 	// DefaultRequestTimeout is the default request timeout
diff --git a/client/is_retriable.go b/client/is_retriable.go
new file mode 100644
index 00000000..01fc8c47
--- /dev/null
+++ b/client/is_retriable.go
@@ -0,0 +1,22 @@
+package client
+
+import (
+	"github.com/micro/go-micro/errors"
+)
+
+type IsRetriableFunc func(err error) bool
+
+// always retry on error
+func AlwaysRetry(err error) bool {
+	return true
+}
+
+func Only500Errors(err error) bool {
+	errorData := errors.Parse(err.Error())
+
+	if(errorData.Code >= 500) {
+		return true
+	}
+
+	return false
+}
\ No newline at end of file
diff --git a/client/options.go b/client/options.go
index 7751d879..b5104990 100644
--- a/client/options.go
+++ b/client/options.go
@@ -43,6 +43,8 @@ type CallOptions struct {
 
 	// Backoff func
 	Backoff BackoffFunc
+	// Check if retriable func
+	CheckIfRetriable IsRetriableFunc
 	// Transport Dial Timeout
 	DialTimeout time.Duration
 	// Number of Call attempts
@@ -73,10 +75,11 @@ func newOptions(options ...Option) Options {
 	opts := Options{
 		Codecs: make(map[string]codec.NewCodec),
 		CallOptions: CallOptions{
-			Backoff:        DefaultBackoff,
-			Retries:        DefaultRetries,
-			RequestTimeout: DefaultRequestTimeout,
-			DialTimeout:    transport.DefaultDialTimeout,
+			Backoff:          DefaultBackoff,
+			CheckIfRetriable: DefaultCheckIfRetriable,
+			Retries:          DefaultRetries,
+			RequestTimeout:   DefaultRequestTimeout,
+			DialTimeout:      transport.DefaultDialTimeout,
 		},
 		PoolSize: DefaultPoolSize,
 		PoolTTL:  DefaultPoolTTL,
@@ -221,6 +224,14 @@ func WithBackoff(fn BackoffFunc) CallOption {
 	}
 }
 
+// WithCheckIfRetriable is a CallOption which overrides that which
+// set in Options.CallOptions
+func WithCheckIfRetriable(fn IsRetriableFunc) CallOption {
+	return func(o *CallOptions) {
+		o.CheckIfRetriable = fn
+	}
+}
+
 // WithRetries is a CallOption which overrides that which
 // set in Options.CallOptions
 func WithRetries(i int) CallOption {
diff --git a/client/rpc_client.go b/client/rpc_client.go
index 6e78db4a..b621eae2 100644
--- a/client/rpc_client.go
+++ b/client/rpc_client.go
@@ -299,6 +299,11 @@ func (r *rpcClient) Call(ctx context.Context, request Request, response interfac
 			if err == nil {
 				return nil
 			}
+
+			if !callOpts.CheckIfRetriable(err) {
+				return err
+			}
+
 			gerr = err
 		}
 	}