move type assertion

balancer_conn_wrappers.go

@@ -44,12 +44,13 @@ type exitIdle struct{}
 // ccBalancerWrapper is a wrapper on top of cc for balancers.
 // It implements balancer.ClientConn interface.
 type ccBalancerWrapper struct {
-	cc         *ClientConn
-	balancerMu sync.Mutex // synchronizes calls to the balancer
-	balancer   balancer.Balancer
-	updateCh   *buffer.Unbounded
-	closed     *grpcsync.Event
-	done       *grpcsync.Event
+	cc          *ClientConn
+	balancerMu  sync.Mutex // synchronizes calls to the balancer
+	balancer    balancer.Balancer
+	hasExitIdle bool
+	updateCh    *buffer.Unbounded
+	closed      *grpcsync.Event
+	done        *grpcsync.Event
 
 	mu       sync.Mutex
 	subConns map[*acBalancerWrapper]struct{}
@@ -65,6 +66,7 @@ func newCCBalancerWrapper(cc *ClientConn, b balancer.Builder, bopts balancer.Bui
 	}
 	go ccb.watcher()
 	ccb.balancer = b.Build(ccb, bopts)
+	_, ccb.hasExitIdle = ccb.balancer.(balancer.ExitIdle)
 	return ccb
 }
 
@@ -93,17 +95,12 @@ func (ccb *ccBalancerWrapper) watcher() {
 			case exitIdle:
 				if ccb.cc.GetState() == connectivity.Idle {
 					if ei, ok := ccb.balancer.(balancer.ExitIdle); ok {
+						// We already checked that the balancer implements
+						// ExitIdle before pushing the event to updateCh, but
+						// check conditionally again as defensive programming.
 						ccb.balancerMu.Lock()
 						ei.ExitIdle()
 						ccb.balancerMu.Unlock()
-					} else {
-						// Fallback path for LB policies that don't support
-						// ExitIdle.
-						ccb.cc.mu.Lock()
-						for ac := range ccb.cc.conns {
-							go ac.connect()
-						}
-						ccb.cc.mu.Unlock()
 					}
 				}
 			default:
@@ -138,8 +135,12 @@ func (ccb *ccBalancerWrapper) close() {
 	<-ccb.done.Done()
 }
 
-func (ccb *ccBalancerWrapper) exitIdle() {
+func (ccb *ccBalancerWrapper) exitIdle() bool {
+	if !ccb.hasExitIdle {
+		return false
+	}
 	ccb.updateCh.Put(exitIdle{})
+	return true
 }
 
 func (ccb *ccBalancerWrapper) handleSubConnStateChange(sc balancer.SubConn, s connectivity.State, err error) {
@@ -168,8 +169,8 @@ func (ccb *ccBalancerWrapper) updateClientConnState(ccs *balancer.ClientConnStat
 
 func (ccb *ccBalancerWrapper) resolverError(err error) {
 	ccb.balancerMu.Lock()
+	defer ccb.balancerMu.Unlock()
 	ccb.balancer.ResolverError(err)
-	ccb.balancerMu.Unlock()
 }
 
 func (ccb *ccBalancerWrapper) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {

clientconn.go

@@ -557,8 +557,11 @@ func (cc *ClientConn) GetState() connectivity.State {
 func (cc *ClientConn) Connect() {
 	cc.mu.Lock()
 	defer cc.mu.Unlock()
-	if cc.balancerWrapper != nil {
-		cc.balancerWrapper.exitIdle()
+	if cc.balancerWrapper != nil && cc.balancerWrapper.exitIdle() {
+		return
+	}
+	for ac := range cc.conns {
+		go ac.connect()
 	}
 }