From 07f59860ab6b4781fccf5f497d08e4312d59bd98 Mon Sep 17 00:00:00 2001
From: Ankur Gargi <ankur.gargi@gmail.com>
Date: Sun, 1 Nov 2020 12:17:22 -0500
Subject: [PATCH] Replace deprecated functions with new grpc version

---
 client/v3/balancer/balancer.go                   |  2 +-
 client/v3/balancer/picker/err.go                 |  6 ++----
 client/v3/balancer/picker/picker.go              |  2 +-
 client/v3/balancer/picker/roundrobin_balanced.go | 10 ++++++----
 client/v3/balancer/resolver/endpoint/endpoint.go |  5 +++--
 5 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/client/v3/balancer/balancer.go b/client/v3/balancer/balancer.go
index 3e63d6091e74..3f5a45310cec 100644
--- a/client/v3/balancer/balancer.go
+++ b/client/v3/balancer/balancer.go
@@ -247,7 +247,7 @@ func (bb *baseBalancer) HandleSubConnStateChange(sc balancer.SubConn, s grpcconn
 		bb.updatePicker()
 	}
 
-	bb.currentConn.UpdateBalancerState(bb.connectivityRecorder.GetCurrentState(), bb.picker)
+	bb.currentConn.UpdateState(balancer.State{ConnectivityState: bb.connectivityRecorder.GetCurrentState(), Picker: bb.picker})
 }
 
 func (bb *baseBalancer) updatePicker() {
diff --git a/client/v3/balancer/picker/err.go b/client/v3/balancer/picker/err.go
index f4b941d6529e..6e8242f59ca8 100644
--- a/client/v3/balancer/picker/err.go
+++ b/client/v3/balancer/picker/err.go
@@ -15,8 +15,6 @@
 package picker
 
 import (
-	"context"
-
 	"google.golang.org/grpc/balancer"
 )
 
@@ -34,6 +32,6 @@ func (ep *errPicker) String() string {
 	return ep.p.String()
 }
 
-func (ep *errPicker) Pick(context.Context, balancer.PickInfo) (balancer.SubConn, func(balancer.DoneInfo), error) {
-	return nil, nil, ep.err
+func (ep *errPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	return balancer.PickResult{}, nil
 }
diff --git a/client/v3/balancer/picker/picker.go b/client/v3/balancer/picker/picker.go
index bd1a5d25e8b8..2b08619bfd5c 100644
--- a/client/v3/balancer/picker/picker.go
+++ b/client/v3/balancer/picker/picker.go
@@ -24,7 +24,7 @@ import (
 
 // Picker defines balancer Picker methods.
 type Picker interface {
-	balancer.Picker
+	balancer.V2Picker
 	String() string
 }
 
diff --git a/client/v3/balancer/picker/roundrobin_balanced.go b/client/v3/balancer/picker/roundrobin_balanced.go
index e3971ecc4210..427fbbcbbfae 100644
--- a/client/v3/balancer/picker/roundrobin_balanced.go
+++ b/client/v3/balancer/picker/roundrobin_balanced.go
@@ -15,7 +15,6 @@
 package picker
 
 import (
-	"context"
 	"sync"
 
 	"go.uber.org/zap"
@@ -52,12 +51,13 @@ type rrBalanced struct {
 func (rb *rrBalanced) String() string { return rb.p.String() }
 
 // Pick is called for every client request.
-func (rb *rrBalanced) Pick(ctx context.Context, opts balancer.PickInfo) (balancer.SubConn, func(balancer.DoneInfo), error) {
+func (rb *rrBalanced) Pick(opts balancer.PickInfo) (balancer.PickResult, error) {
 	rb.mu.RLock()
 	n := len(rb.scs)
 	rb.mu.RUnlock()
 	if n == 0 {
-		return nil, nil, balancer.ErrNoSubConnAvailable
+		pr := balancer.PickResult{}
+		return pr, balancer.ErrNoSubConnAvailable
 	}
 
 	rb.mu.Lock()
@@ -91,5 +91,7 @@ func (rb *rrBalanced) Pick(ctx context.Context, opts balancer.PickInfo) (balance
 			rb.lg.Warn("balancer failed", fss...)
 		}
 	}
-	return sc, doneFunc, nil
+
+	pr := balancer.PickResult{SubConn: rb.scs[cur], Done: doneFunc}
+	return pr, nil
 }
diff --git a/client/v3/balancer/resolver/endpoint/endpoint.go b/client/v3/balancer/resolver/endpoint/endpoint.go
index 86992cbe6802..42b5177deb9f 100644
--- a/client/v3/balancer/resolver/endpoint/endpoint.go
+++ b/client/v3/balancer/resolver/endpoint/endpoint.go
@@ -65,7 +65,7 @@ func (e *ResolverGroup) addResolver(r *Resolver) {
 	addrs := epsToAddrs(e.endpoints...)
 	e.resolvers = append(e.resolvers, r)
 	e.mu.Unlock()
-	r.cc.NewAddress(addrs)
+	r.cc.UpdateState(resolver.State{Addresses: addrs})
 }
 
 func (e *ResolverGroup) removeResolver(r *Resolver) {
@@ -85,8 +85,9 @@ func (e *ResolverGroup) SetEndpoints(endpoints []string) {
 	addrs := epsToAddrs(endpoints...)
 	e.mu.Lock()
 	e.endpoints = endpoints
+	state := resolver.State{Addresses: addrs}
 	for _, r := range e.resolvers {
-		r.cc.NewAddress(addrs)
+		r.cc.UpdateState(state)
 	}
 	e.mu.Unlock()
 }