From 2435eb9ecdddc69f62acc1a4086cddc7fcc09fb6 Mon Sep 17 00:00:00 2001
From: Xiang Li <xiangli.cs@gmail.com>
Date: Wed, 14 Sep 2016 13:52:16 +0800
Subject: [PATCH] clientv3: balancer panics when call up after close

Fix the issue by adding a simple guard varable.
---
 clientv3/balancer.go | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/clientv3/balancer.go b/clientv3/balancer.go
index b7fba6a20..7e8d957f9 100644
--- a/clientv3/balancer.go
+++ b/clientv3/balancer.go
@@ -45,6 +45,8 @@ type simpleBalancer struct {
 	// pinAddr is the currently pinned address; set to the empty string on
 	// intialization and shutdown.
 	pinAddr string
+
+	closed bool
 }
 
 func newSimpleBalancer(eps []string) *simpleBalancer {
@@ -74,15 +76,25 @@ func (b *simpleBalancer) ConnectNotify() <-chan struct{} {
 
 func (b *simpleBalancer) Up(addr grpc.Address) func(error) {
 	b.mu.Lock()
+	defer b.mu.Unlock()
+
+	// gRPC might call Up after it called Close. We add this check
+	// to "fix" it up at application layer. Or our simplerBalancer
+	// might panic since b.upc is closed.
+	if b.closed {
+		return func(err error) {}
+	}
+
 	if len(b.upEps) == 0 {
 		// notify waiting Get()s and pin first connected address
 		close(b.upc)
 		b.pinAddr = addr.Addr
 	}
 	b.upEps[addr.Addr] = struct{}{}
-	b.mu.Unlock()
+
 	// notify client that a connection is up
 	b.readyOnce.Do(func() { close(b.readyc) })
+
 	return func(err error) {
 		b.mu.Lock()
 		delete(b.upEps, addr.Addr)
@@ -128,13 +140,19 @@ func (b *simpleBalancer) Notify() <-chan []grpc.Address { return b.notifyCh }
 
 func (b *simpleBalancer) Close() error {
 	b.mu.Lock()
+	defer b.mu.Unlock()
+	// In case gRPC calls close twice. TODO: remove the checking
+	// when we are sure that gRPC wont call close twice.
+	if b.closed {
+		return nil
+	}
+	b.closed = true
 	close(b.notifyCh)
 	// terminate all waiting Get()s
 	b.pinAddr = ""
 	if len(b.upEps) == 0 {
 		close(b.upc)
 	}
-	b.mu.Unlock()
 	return nil
 }