From a62701e4aa1d276bec70311251d62a478404d63f Mon Sep 17 00:00:00 2001
From: Brandon Gao <reterclose@gmail.com>
Date: Tue, 28 Nov 2017 14:41:43 -0800
Subject: [PATCH] Eliminate data race in ccBalancerWrapper (#1688)

---
 balancer_conn_wrappers.go | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/balancer_conn_wrappers.go b/balancer_conn_wrappers.go
index c747e92b..395c4b6b 100644
--- a/balancer_conn_wrappers.go
+++ b/balancer_conn_wrappers.go
@@ -98,6 +98,7 @@ type ccBalancerWrapper struct {
 	resolverUpdateCh chan *resolverUpdate
 	done             chan struct{}
 
+	mu       sync.RWMutex
 	subConns map[*acBalancerWrapper]struct{}
 }
 
@@ -142,9 +143,11 @@ func (ccb *ccBalancerWrapper) watcher() {
 		select {
 		case <-ccb.done:
 			ccb.balancer.Close()
+			ccb.mu.RLock()
 			for acbw := range ccb.subConns {
 				ccb.cc.removeAddrConn(acbw.getAddrConn(), errConnDrain)
 			}
+			ccb.mu.RUnlock()
 			return
 		default:
 		}
@@ -195,7 +198,9 @@ func (ccb *ccBalancerWrapper) NewSubConn(addrs []resolver.Address, opts balancer
 	acbw.ac.mu.Lock()
 	ac.acbw = acbw
 	acbw.ac.mu.Unlock()
+	ccb.mu.Lock()
 	ccb.subConns[acbw] = struct{}{}
+	ccb.mu.Unlock()
 	return acbw, nil
 }
 
@@ -204,7 +209,9 @@ func (ccb *ccBalancerWrapper) RemoveSubConn(sc balancer.SubConn) {
 	if !ok {
 		return
 	}
+	ccb.mu.Lock()
 	delete(ccb.subConns, acbw)
+	ccb.mu.Unlock()
 	ccb.cc.removeAddrConn(acbw.getAddrConn(), errConnDrain)
 }