Add withinitbalancer dial option

service config will not override dial option balancer

balancer_switching_test.go

@@ -131,3 +131,31 @@ func TestSwitchBalancer(t *testing.T) {
 		t.Fatalf("check pickfirst returned non-nil error: %v", err)
 	}
 }
+
+// Test that balancer specified by dial option will not be overridden.
+func TestBalancerDialOption(t *testing.T) {
+	defer leakcheck.Check(t)
+	r, rcleanup := manual.GenerateAndRegisterManualResolver()
+	defer rcleanup()
+
+	numServers := 2
+	servers, _, scleanup := startServers(t, numServers, math.MaxInt32)
+	defer scleanup()
+
+	cc, err := Dial(r.Scheme()+":///test.server", WithInsecure(), WithCodec(testCodec{}), WithBalancerName("round_robin"))
+	if err != nil {
+		t.Fatalf("failed to dial: %v", err)
+	}
+	defer cc.Close()
+	r.NewAddress([]resolver.Address{{Addr: servers[0].addr}, {Addr: servers[1].addr}})
+	// The init balancer is roundrobin.
+	if err := checkRoundRobin(cc, servers); err != nil {
+		t.Fatalf("check roundrobin returned non-nil error: %v", err)
+	}
+	// Switch to pickfirst.
+	cc.handleServiceConfig(`{"loadBalancingPolicy": "pick_first"}`)
+	// Balancer is still roundrobin.
+	if err := checkRoundRobin(cc, servers); err != nil {
+		t.Fatalf("check roundrobin returned non-nil error: %v", err)
+	}
+}

clientconn.go

@@ -97,6 +97,8 @@ type dialOptions struct {
 	callOptions []CallOption
 	// This is to support v1 balancer.
 	balancerBuilder balancer.Builder
+	// The balancer to be used. Can not be overridden by service config.
+	balancerName string
 	// This is to support grpclb.
 	resolverBuilder  resolver.Builder
 	waitForHandshake bool
@@ -198,7 +200,7 @@ func WithDecompressor(dc Decompressor) DialOption {
 
 // WithBalancer returns a DialOption which sets a load balancer with the v1 API.
 // Name resolver will be ignored if this DialOption is specified.
-// Deprecated: use the new balancer APIs in balancer package instead.
+// Deprecated: use the new balancer APIs in balancer package and WithBalancerName.
 func WithBalancer(b Balancer) DialOption {
 	return func(o *dialOptions) {
 		o.balancerBuilder = &balancerWrapperBuilder{
@@ -216,6 +218,15 @@ func WithBalancerBuilder(b balancer.Builder) DialOption {
 	}
 }
 
+// WithBalancerName sets the balancer that the ClientConn will be initialized
+// with. The balancer can not be overridden by balancer option specified by
+// service config.
+func WithBalancerName(balancerName string) DialOption {
+	return func(o *dialOptions) {
+		o.balancerName = balancerName
+	}
+}
+
 // withResolverBuilder is only for grpclb.
 func withResolverBuilder(b resolver.Builder) DialOption {
 	return func(o *dialOptions) {
@@ -661,7 +672,10 @@ func (cc *ClientConn) handleResolvedAddrs(addrs []resolver.Address, err error) {
 	if cc.balancerWrapper == nil {
 		// First time handling resolved addresses. Build a balancer use either
 		// the builder specified by dial option, or pickfirst.
-		builder := cc.dopts.balancerBuilder
+		builder := balancer.Get(cc.dopts.balancerName)
+		if builder == nil && cc.dopts.balancerBuilder != nil {
+			builder = cc.dopts.balancerBuilder
+		}
 		if builder == nil {
 			// No customBalancer was specified by DialOption, and this is the first
 			// time handling resolved addresses, create a pickfirst balancer.
@@ -684,8 +698,8 @@ func (cc *ClientConn) switchBalancer(name string) {
 	}
 	grpclog.Infof("ClientConn switching balancer to %q", name)
 
-	if cc.dopts.balancerBuilder != nil {
-		grpclog.Infoln("ignoring service config balancer configuration: WithBalancer DialOption used instead")
+	if cc.dopts.balancerName != "" || cc.dopts.balancerBuilder != nil {
+		grpclog.Infoln("ignoring service config balancer configuration: Balancer DialOption used instead")
 		return
 	}
 

grpclb_remote_balancer.go

@@ -241,7 +241,7 @@ func (lb *lbBalancer) dialRemoteLB(remoteLBName string) {
 		dopts = append(dopts, withContextDialer(lb.opt.Dialer))
 	}
 	// Explicitly set pickfirst as the balancer.
-	dopts = append(dopts, WithBalancerBuilder(newPickfirstBuilder()))
+	dopts = append(dopts, WithBalancerName("pick_first"))
 	dopts = append(dopts, withResolverBuilder(lb.manualResolver))
 	// Dial using manualResolver.Scheme, which is a random scheme generated
 	// when init grpclb. The target name is not important.