resolver: add State fields to support error handling (#2951)

balancer_conn_wrappers.go

@@ -25,6 +25,7 @@ import (
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal/grpcsync"
 	"google.golang.org/grpc/resolver"
 )
 
@@ -86,10 +87,10 @@ func (b *scStateUpdateBuffer) get() <-chan *scStateUpdate {
 // It implements balancer.ClientConn interface.
 type ccBalancerWrapper struct {
 	cc               *ClientConn
+	balancerMu       sync.Mutex // synchronizes calls to the balancer
 	balancer         balancer.Balancer
 	stateChangeQueue *scStateUpdateBuffer
-	ccUpdateCh       chan *balancer.ClientConnState
-	done             chan struct{}
+	done             *grpcsync.Event
 
 	mu       sync.Mutex
 	subConns map[*acBalancerWrapper]struct{}
@@ -99,8 +100,7 @@ func newCCBalancerWrapper(cc *ClientConn, b balancer.Builder, bopts balancer.Bui
 	ccb := &ccBalancerWrapper{
 		cc:               cc,
 		stateChangeQueue: newSCStateUpdateBuffer(),
-		ccUpdateCh:       make(chan *balancer.ClientConnState, 1),
-		done:             make(chan struct{}),
+		done:             grpcsync.NewEvent(),
 		subConns:         make(map[*acBalancerWrapper]struct{}),
 	}
 	go ccb.watcher()
@@ -115,34 +115,20 @@ func (ccb *ccBalancerWrapper) watcher() {
 		select {
 		case t := <-ccb.stateChangeQueue.get():
 			ccb.stateChangeQueue.load()
-			select {
-			case <-ccb.done:
-				ccb.balancer.Close()
-				return
-			default:
+			if ccb.done.HasFired() {
+				break
 			}
+			ccb.balancerMu.Lock()
 			if ub, ok := ccb.balancer.(balancer.V2Balancer); ok {
 				ub.UpdateSubConnState(t.sc, balancer.SubConnState{ConnectivityState: t.state})
 			} else {
 				ccb.balancer.HandleSubConnStateChange(t.sc, t.state)
 			}
-		case s := <-ccb.ccUpdateCh:
-			select {
-			case <-ccb.done:
-				ccb.balancer.Close()
-				return
-			default:
-			}
-			if ub, ok := ccb.balancer.(balancer.V2Balancer); ok {
-				ub.UpdateClientConnState(*s)
-			} else {
-				ccb.balancer.HandleResolvedAddrs(s.ResolverState.Addresses, nil)
-			}
-		case <-ccb.done:
+			ccb.balancerMu.Unlock()
+		case <-ccb.done.Done():
 		}
 
-		select {
-		case <-ccb.done:
+		if ccb.done.HasFired() {
 			ccb.balancer.Close()
 			ccb.mu.Lock()
 			scs := ccb.subConns
@@ -153,14 +139,12 @@ func (ccb *ccBalancerWrapper) watcher() {
 			}
 			ccb.UpdateBalancerState(connectivity.Connecting, nil)
 			return
-		default:
 		}
-		ccb.cc.firstResolveEvent.Fire()
 	}
 }
 
 func (ccb *ccBalancerWrapper) close() {
-	close(ccb.done)
+	ccb.done.Fire()
 }
 
 func (ccb *ccBalancerWrapper) handleSubConnStateChange(sc balancer.SubConn, s connectivity.State) {
@@ -180,24 +164,22 @@ func (ccb *ccBalancerWrapper) handleSubConnStateChange(sc balancer.SubConn, s co
 	})
 }
 
-func (ccb *ccBalancerWrapper) updateClientConnState(ccs *balancer.ClientConnState) {
-	if ccb.cc.curBalancerName != grpclbName {
-		// Filter any grpclb addresses since we don't have the grpclb balancer.
-		s := &ccs.ResolverState
-		for i := 0; i < len(s.Addresses); {
-			if s.Addresses[i].Type == resolver.GRPCLB {
-				copy(s.Addresses[i:], s.Addresses[i+1:])
-				s.Addresses = s.Addresses[:len(s.Addresses)-1]
-				continue
-			}
-			i++
-		}
+func (ccb *ccBalancerWrapper) updateClientConnState(ccs *balancer.ClientConnState) error {
+	ccb.balancerMu.Lock()
+	defer ccb.balancerMu.Unlock()
+	if ub, ok := ccb.balancer.(balancer.V2Balancer); ok {
+		return ub.UpdateClientConnState(*ccs)
 	}
-	select {
-	case <-ccb.ccUpdateCh:
-	default:
+	ccb.balancer.HandleResolvedAddrs(ccs.ResolverState.Addresses, nil)
+	return nil
+}
+
+func (ccb *ccBalancerWrapper) resolverError(err error) {
+	if ub, ok := ccb.balancer.(balancer.V2Balancer); ok {
+		ccb.balancerMu.Lock()
+		ub.ResolverError(err)
+		ccb.balancerMu.Unlock()
 	}
-	ccb.ccUpdateCh <- ccs
 }
 
 func (ccb *ccBalancerWrapper) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {