From bed5f388a8e23adf39dcc5d9ac37ce6fac879959 Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Fri, 6 Oct 2017 12:58:52 -0700
Subject: [PATCH] clientv3: add pinned() method to 'balancer'

Signed-off-by: Gyu-Ho Lee <gyuhox@gmail.com>
---
 clientv3/balancer.go | 8 ++++++++
 clientv3/retry.go    | 6 ++++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/clientv3/balancer.go b/clientv3/balancer.go
index deb0df1e5..35d8b3a4a 100644
--- a/clientv3/balancer.go
+++ b/clientv3/balancer.go
@@ -42,6 +42,8 @@ type balancer interface {
 
 	endpoint(host string) string
 	endpoints() []string
+	// pinned returns the current pinned endpoint.
+	pinned() string
 
 	// up is Up but includes whether the balancer will use the connection.
 	up(addr grpc.Address) (func(error), bool)
@@ -142,6 +144,12 @@ func (b *simpleBalancer) endpoints() []string {
 	return b.eps
 }
 
+func (b *simpleBalancer) pinned() string {
+	b.mu.RLock()
+	defer b.mu.RUnlock()
+	return b.pinAddr
+}
+
 func getHost2ep(eps []string) map[string]string {
 	hm := make(map[string]string, len(eps))
 	for i := range eps {
diff --git a/clientv3/retry.go b/clientv3/retry.go
index ddf0ea3c8..d6681441e 100644
--- a/clientv3/retry.go
+++ b/clientv3/retry.go
@@ -58,12 +58,13 @@ func (c *Client) newRetryWrapper(isStop retryStopErrFunc) retryRpcFunc {
 			case <-c.ctx.Done():
 				return c.ctx.Err()
 			}
+			pinned := c.balancer.pinned()
 			err := f(rpcCtx)
 			if err == nil {
 				return nil
 			}
 			if logger.V(4) {
-				logger.Infof("clientv3/retry: retry for error %v", err)
+				logger.Infof("clientv3/retry: error %v on pinned endpoint %s", err, pinned)
 			}
 			notify := c.balancer.ConnectNotify()
 			if s, ok := status.FromError(err); ok && s.Code() == codes.Unavailable {
@@ -86,12 +87,13 @@ func (c *Client) newRetryWrapper(isStop retryStopErrFunc) retryRpcFunc {
 func (c *Client) newAuthRetryWrapper() retryRpcFunc {
 	return func(rpcCtx context.Context, f rpcFunc) error {
 		for {
+			pinned := c.balancer.pinned()
 			err := f(rpcCtx)
 			if err == nil {
 				return nil
 			}
 			if logger.V(4) {
-				logger.Infof("clientv3/auth-retry: retry for error %v", err)
+				logger.Infof("clientv3/auth-retry: error %v on pinned endpoint %s", err, pinned)
 			}
 			// always stop retry on etcd errors other than invalid auth token
 			if rpctypes.Error(err) == rpctypes.ErrInvalidAuthToken {