Browse Source

clientv3/balancer: only notify healthy addresses

Signed-off-by: Gyu-Ho Lee <gyuhox@gmail.com>
Gyu-Ho Lee 8 years ago
parent
commit
103efd922b
2 changed files with 7 additions and 11 deletions
  1. 3 7
      clientv3/health_balancer.go
  2. 4 4
      clientv3/health_balancer_test.go

+ 3 - 7
clientv3/health_balancer.go

@@ -367,19 +367,15 @@ func (b *healthBalancer) notifyAddrs(msg notifyMsg) {
 		}
 	}
 	b.mu.RLock()
-	addrs := b.addrs
 	pinAddr := b.pinAddr
 	downc := b.downc
 	b.mu.RUnlock()
+	addrs, hostPorts := b.liveAddrs()
 
 	var waitDown bool
 	if pinAddr != "" {
-		waitDown = true
-		for _, a := range addrs {
-			if a.Addr == pinAddr {
-				waitDown = false
-			}
-		}
+		_, ok := hostPorts[pinAddr]
+		waitDown = !ok
 	}
 
 	select {

+ 4 - 4
clientv3/health_balancer_test.go

@@ -64,8 +64,8 @@ func TestBalancerGetUnblocking(t *testing.T) {
 	}
 
 	down1(errors.New("error"))
-	if addrs := <-hb.Notify(); len(addrs) != len(endpoints) {
-		t.Errorf("closing the only connection should triggered balancer to send the all endpoints via Notify chan so that we can establish a connection")
+	if addrs := <-hb.Notify(); len(addrs) != len(endpoints)-1 { // we call down on one endpoint
+		t.Errorf("closing the only connection should triggered balancer to send the %d endpoints via Notify chan so that we can establish a connection", len(endpoints)-1)
 	}
 	down2(errors.New("error"))
 	_, _, err = hb.Get(context.Background(), unblockingOpts)
@@ -119,8 +119,8 @@ func TestBalancerGetBlocking(t *testing.T) {
 	}
 
 	down1(errors.New("error"))
-	if addrs := <-hb.Notify(); len(addrs) != len(endpoints) {
-		t.Errorf("closing the only connection should triggered balancer to send the all endpoints via Notify chan so that we can establish a connection")
+	if addrs := <-hb.Notify(); len(addrs) != len(endpoints)-1 { // we call down on one endpoint
+		t.Errorf("closing the only connection should triggered balancer to send the %d endpoints via Notify chan so that we can establish a connection", len(endpoints)-1)
 	}
 	down2(errors.New("error"))
 	ctx, cancel = context.WithTimeout(context.Background(), time.Millisecond*100)