use the existing manual resolver

Author: easwars

balancer/grpclb/grpclb.go

@@ -132,11 +132,6 @@ func (b *lbBuilder) Name() string {
 }
 
 func (b *lbBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions) balancer.Balancer {
-	// This generates a manual resolver builder with a fixed scheme. This
-	// scheme will be used to dial to remote LB, so we can send filtered
-	// address updates to remote LB ClientConn using this manual resolver.
-	r := &lbManualResolver{scheme: "grpclb-internal", ccb: cc}
-
 	lb := &lbBalancer{
 		cc:              newLBCacheClientConn(cc),
 		dialTarget:      opt.Target.Endpoint(),
@@ -145,7 +140,10 @@ func (b *lbBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions) bal
 		fallbackTimeout: b.fallbackTimeout,
 		doneCh:          make(chan struct{}),
 
-		manualResolver: r,
+		// This generates a manual resolver builder with a fixed scheme. This
+		// scheme will be used to dial to remote LB, so we can send filtered
+		// address updates to remote LB ClientConn using this manual resolver.
+		manualResolver: newManualResolver("grpclb-internal", cc),
 		subConns:       make(map[resolver.Address]balancer.SubConn),
 		scStates:       make(map[balancer.SubConn]connectivity.State),
 		picker:         base.NewErrPicker(balancer.ErrNoSubConnAvailable),

balancer/grpclb/grpclb_util.go

@@ -25,6 +25,7 @@ import (
 
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/resolver/manual"
 )
 
 // The parent ClientConn should re-resolve when grpclb loses connection to the
@@ -60,50 +61,41 @@ import (
 // so when grpclb client lose contact with remote balancers, the parent
 // ClientConn's resolver will re-resolve.
 type lbManualResolver struct {
-	scheme string
-	ccr    resolver.ClientConn
-	ccb    balancer.ClientConn
-
-	// Cache the most recently received state update via UpdateState(), and push
-	// the same when Build() is invoked. This ensures that this manual resolver
-	// can handle restarts when channel idleness comes into the picture.
-	stateMu         sync.Mutex
-	lastState       resolver.State
-	stateUpdateRecv bool
-}
-
-func (r *lbManualResolver) Build(_ resolver.Target, cc resolver.ClientConn, _ resolver.BuildOptions) (resolver.Resolver, error) {
-	r.ccr = cc
-
-	r.stateMu.Lock()
-	if r.stateUpdateRecv {
-		r.ccr.UpdateState(r.lastState)
+	// These fields are setup at creation time and are read-only after that, and
+	// therefore need not be protected with a mutex.
+	resolver.Builder                      // The underlying manual resolver builder.
+	resolver.Resolver                     // The underlying manual resolver.
+	ccb               balancer.ClientConn // ClientConn passed to the parent channel.
+
+	// The resolver.ClientConn is updated everytime the resolver on the channel
+	// is restarted, and this happens when the channel exits IDLE.
+	mu  sync.Mutex
+	ccr resolver.ClientConn
+}
+
+func newManualResolver(scheme string, ccb balancer.ClientConn) *lbManualResolver {
+	mr := manual.NewBuilderWithScheme(scheme)
+	r := &lbManualResolver{
+		Builder:  mr,
+		Resolver: mr,
+		ccb:      ccb,
 	}
-	r.stateMu.Unlock()
 
-	return r, nil
-}
-
-func (r *lbManualResolver) Scheme() string {
-	return r.scheme
-}
-
-// ResolveNow calls resolveNow on the parent ClientConn.
-func (r *lbManualResolver) ResolveNow(o resolver.ResolveNowOptions) {
-	r.ccb.ResolveNow(o)
+	mr.BuildCallback = func(_ resolver.Target, ccr resolver.ClientConn, _ resolver.BuildOptions) {
+		r.mu.Lock()
+		r.ccr = ccr
+		r.mu.Unlock()
+	}
+	mr.ResolveNowCallback = func(o resolver.ResolveNowOptions) {
+		r.ccb.ResolveNow(o)
+	}
+	return r
 }
 
-// Close is a noop for Resolver.
-func (*lbManualResolver) Close() {}
-
-// UpdateState calls cc.UpdateState.
 func (r *lbManualResolver) UpdateState(s resolver.State) {
-	r.stateMu.Lock()
-	r.lastState = s
-	r.stateUpdateRecv = true
-	r.stateMu.Unlock()
-
+	r.mu.Lock()
 	r.ccr.UpdateState(s)
+	r.mu.Unlock()
 }
 
 const subConnCacheTime = time.Second * 10