From a66bce0e4b8ca242a64882f6e6e12f55b8aa0443 Mon Sep 17 00:00:00 2001
From: Scott Finlay <ynori7@yahoo.com>
Date: Mon, 7 Nov 2016 17:10:40 +0100
Subject: [PATCH] Adjusting arguments and return value of retry function and
 adding new retry logic to stream

---
 client/retry.go      | 8 +++++---
 client/rpc_client.go | 9 +++++++--
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/client/retry.go b/client/retry.go
index df74abec..59743980 100644
--- a/client/retry.go
+++ b/client/retry.go
@@ -1,8 +1,10 @@
 package client
 
-type RetryFunc func(err error) bool
+import "context"
+
+type RetryFunc func(ctx context.Context, req Request, retryCount int, err error) (bool, error)
 
 // always retry on error
-func alwaysRetry(err error) bool {
-	return true
+func alwaysRetry(ctx context.Context, req Request, retryCount int, err error) (bool, error) {
+	return true, err
 }
diff --git a/client/rpc_client.go b/client/rpc_client.go
index 0ad84805..ede9e899 100644
--- a/client/rpc_client.go
+++ b/client/rpc_client.go
@@ -300,7 +300,7 @@ func (r *rpcClient) Call(ctx context.Context, request Request, response interfac
 				return nil
 			}
 
-			if !callOpts.Retry(err) {
+			if retry, err := callOpts.Retry(ctx, request, i, err); !retry {
 				return err
 			}
 
@@ -405,7 +405,12 @@ func (r *rpcClient) Stream(ctx context.Context, request Request, opts ...CallOpt
 			if rsp.err == nil {
 				return rsp.stream, nil
 			}
-			grr = rsp.err
+
+			if retry, err := callOpts.Retry(ctx, request, i, rsp.err); !retry {
+				return nil, err
+			}
+
+			grr = err
 		}
 	}