balancer: filter out grpclb addresses if balancer is not grpclb

Author: menghanl

Diff for: balancer_conn_wrappers.go

@@ -183,7 +183,7 @@ 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
+		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:])

Diff for: balancer_test.go renamed to balancer_v1_test.go

@@ -24,6 +24,8 @@ import (
 	"testing"
 	"time"
 
+	"google.golang.org/grpc/resolver/manual"
+
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/connectivity"
@@ -260,3 +262,104 @@ func testDoneLoads(t *testing.T, e env) {
 		t.Fatalf("b.doneInfo[0].ServerLoad = %v; want = %v", b.doneInfo[0].ServerLoad, testLoad)
 	}
 }
+
+const testBalancerKeepAddressesName = "testbalancer-keepingaddresses"
+
+// testBalancerKeepAddresses keeps the addresses in the builder instead of
+// creating SubConns.
+//
+// It's used to test the addresses balancer gets are correct.
+type testBalancerKeepAddresses struct {
+	addrs []resolver.Address
+}
+
+func (b *testBalancerKeepAddresses) Build(cc balancer.ClientConn, opt balancer.BuildOptions) balancer.Balancer {
+	return b
+}
+
+func (*testBalancerKeepAddresses) Name() string {
+	return testBalancerKeepAddressesName
+}
+
+func (b *testBalancerKeepAddresses) HandleResolvedAddrs(addrs []resolver.Address, err error) {
+	b.addrs = addrs
+}
+
+func (b *testBalancerKeepAddresses) HandleSubConnStateChange(sc balancer.SubConn, s connectivity.State) {
+	panic("not used")
+}
+
+func (b *testBalancerKeepAddresses) Close() {
+}
+
+// Make sure that non-grpclb balancers don't get grpclb addresses even if name
+// resolver sends them
+func (s) TestNonGRPCLBBalancerGetsNoGRPCLBAddress(t *testing.T) {
+	r, rcleanup := manual.GenerateAndRegisterManualResolver()
+	defer rcleanup()
+
+	b := &testBalancerKeepAddresses{}
+	balancer.Register(b)
+
+	cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(),
+		grpc.WithBalancerName(b.Name()))
+	if err != nil {
+		t.Fatalf("failed to dial: %v", err)
+	}
+	defer cc.Close()
+
+	if len(b.addrs) != 0 {
+		t.Fatalf("Before name resolving, balancer got addresses %v, want empty", b.addrs)
+	}
+
+	grpclbAddresses := []resolver.Address{{
+		Addr:       "grpc.lb.com",
+		Type:       resolver.GRPCLB,
+		ServerName: "grpc.lb.com",
+	}}
+
+	nonGRPCLBAddresses := []resolver.Address{{
+		Addr: "localhost",
+		Type: resolver.Backend,
+	}}
+
+	r.UpdateState(resolver.State{
+		Addresses: nonGRPCLBAddresses,
+	})
+	var i int
+	for i = 0; i < 1000; i++ {
+		if reflect.DeepEqual(b.addrs, nonGRPCLBAddresses) {
+			break
+		}
+		time.Sleep(time.Millisecond)
+	}
+	if i >= 1000 {
+		t.Fatalf("With only backend addresses, balancer got addresses %v, want %v", b.addrs, nonGRPCLBAddresses)
+	}
+
+	r.UpdateState(resolver.State{
+		Addresses: grpclbAddresses,
+	})
+	for i = 0; i < 1000; i++ {
+		if len(b.addrs) == 0 {
+			break
+		}
+		time.Sleep(time.Millisecond)
+	}
+	if i >= 1000 {
+		t.Fatalf("With only grpclb addresses, balancer got addresses %v, want empty", b.addrs)
+	}
+
+	r.UpdateState(resolver.State{
+		Addresses: append(grpclbAddresses, nonGRPCLBAddresses...),
+	})
+	for i = 0; i < 1000; i++ {
+		if reflect.DeepEqual(b.addrs, nonGRPCLBAddresses) {
+			break
+		}
+		time.Sleep(time.Millisecond)
+	}
+	if i >= 1000 {
+		t.Fatalf("With both backend and grpclb addresses, balancer got addresses %v, want %v", b.addrs, nonGRPCLBAddresses)
+	}
+}