Skip to content

Commit e6549e6

Browse files
authored
Add dial option to set balancer (grpc#1697)
WithBalancerName dial option specifies the name of the balancer to be used by the ClientConn. Service config updates can NOT override the balancer option.
1 parent 6610f9a commit e6549e6

File tree

10 files changed

+110
-45
lines changed

10 files changed

+110
-45
lines changed

balancer/roundrobin/roundrobin.go

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,9 +31,12 @@ import (
3131
"google.golang.org/grpc/resolver"
3232
)
3333

34+
// Name is the name of round_robin balancer.
35+
const Name = "round_robin"
36+
3437
// newBuilder creates a new roundrobin balancer builder.
3538
func newBuilder() balancer.Builder {
36-
return base.NewBalancerBuilder("round_robin", &rrPickerBuilder{})
39+
return base.NewBalancerBuilder(Name, &rrPickerBuilder{})
3740
}
3841

3942
func init() {

balancer/roundrobin/roundrobin_test.go

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ import (
2727

2828
"golang.org/x/net/context"
2929
"google.golang.org/grpc"
30-
"google.golang.org/grpc/balancer"
30+
"google.golang.org/grpc/balancer/roundrobin"
3131
"google.golang.org/grpc/codes"
3232
_ "google.golang.org/grpc/grpclog/glogger"
3333
"google.golang.org/grpc/peer"
@@ -38,8 +38,6 @@ import (
3838
"google.golang.org/grpc/test/leakcheck"
3939
)
4040

41-
var rr = balancer.Get("round_robin")
42-
4341
type testServer struct {
4442
testpb.TestServiceServer
4543
}
@@ -103,7 +101,7 @@ func TestOneBackend(t *testing.T) {
103101
}
104102
defer test.cleanup()
105103

106-
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerBuilder(rr))
104+
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerName(roundrobin.Name))
107105
if err != nil {
108106
t.Fatalf("failed to dial: %v", err)
109107
}
@@ -135,7 +133,7 @@ func TestBackendsRoundRobin(t *testing.T) {
135133
}
136134
defer test.cleanup()
137135

138-
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerBuilder(rr))
136+
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerName(roundrobin.Name))
139137
if err != nil {
140138
t.Fatalf("failed to dial: %v", err)
141139
}
@@ -194,7 +192,7 @@ func TestAddressesRemoved(t *testing.T) {
194192
}
195193
defer test.cleanup()
196194

197-
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerBuilder(rr))
195+
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerName(roundrobin.Name))
198196
if err != nil {
199197
t.Fatalf("failed to dial: %v", err)
200198
}
@@ -236,7 +234,7 @@ func TestCloseWithPendingRPC(t *testing.T) {
236234
}
237235
defer test.cleanup()
238236

239-
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerBuilder(rr))
237+
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerName(roundrobin.Name))
240238
if err != nil {
241239
t.Fatalf("failed to dial: %v", err)
242240
}
@@ -270,7 +268,7 @@ func TestNewAddressWhileBlocking(t *testing.T) {
270268
}
271269
defer test.cleanup()
272270

273-
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerBuilder(rr))
271+
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerName(roundrobin.Name))
274272
if err != nil {
275273
t.Fatalf("failed to dial: %v", err)
276274
}
@@ -319,7 +317,7 @@ func TestOneServerDown(t *testing.T) {
319317
}
320318
defer test.cleanup()
321319

322-
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerBuilder(rr), grpc.WithWaitForHandshake())
320+
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerName(roundrobin.Name), grpc.WithWaitForHandshake())
323321
if err != nil {
324322
t.Fatalf("failed to dial: %v", err)
325323
}
@@ -417,7 +415,7 @@ func TestAllServersDown(t *testing.T) {
417415
}
418416
defer test.cleanup()
419417

420-
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerBuilder(rr), grpc.WithWaitForHandshake())
418+
cc, err := grpc.Dial(r.Scheme()+":///test.server", grpc.WithInsecure(), grpc.WithBalancerName(roundrobin.Name), grpc.WithWaitForHandshake())
421419
if err != nil {
422420
t.Fatalf("failed to dial: %v", err)
423421
}

balancer_switching_test.go

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import (
2525
"time"
2626

2727
"golang.org/x/net/context"
28+
"google.golang.org/grpc/balancer/roundrobin"
2829
_ "google.golang.org/grpc/grpclog/glogger"
2930
"google.golang.org/grpc/resolver"
3031
"google.golang.org/grpc/resolver/manual"
@@ -132,6 +133,34 @@ func TestSwitchBalancer(t *testing.T) {
132133
}
133134
}
134135

136+
// Test that balancer specified by dial option will not be overridden.
137+
func TestBalancerDialOption(t *testing.T) {
138+
defer leakcheck.Check(t)
139+
r, rcleanup := manual.GenerateAndRegisterManualResolver()
140+
defer rcleanup()
141+
142+
numServers := 2
143+
servers, _, scleanup := startServers(t, numServers, math.MaxInt32)
144+
defer scleanup()
145+
146+
cc, err := Dial(r.Scheme()+":///test.server", WithInsecure(), WithCodec(testCodec{}), WithBalancerName(roundrobin.Name))
147+
if err != nil {
148+
t.Fatalf("failed to dial: %v", err)
149+
}
150+
defer cc.Close()
151+
r.NewAddress([]resolver.Address{{Addr: servers[0].addr}, {Addr: servers[1].addr}})
152+
// The init balancer is roundrobin.
153+
if err := checkRoundRobin(cc, servers); err != nil {
154+
t.Fatalf("check roundrobin returned non-nil error: %v", err)
155+
}
156+
// Switch to pickfirst.
157+
cc.handleServiceConfig(`{"loadBalancingPolicy": "pick_first"}`)
158+
// Balancer is still roundrobin.
159+
if err := checkRoundRobin(cc, servers); err != nil {
160+
t.Fatalf("check roundrobin returned non-nil error: %v", err)
161+
}
162+
}
163+
135164
// First addr update contains grpclb.
136165
func TestSwitchBalancerGRPCLBFirst(t *testing.T) {
137166
defer leakcheck.Check(t)
@@ -182,7 +211,7 @@ func TestSwitchBalancerGRPCLBFirst(t *testing.T) {
182211
r.NewAddress([]resolver.Address{{Addr: "backend"}})
183212
for i := 0; i < 5000; i++ {
184213
cc.mu.Lock()
185-
isPickFirst = cc.curBalancerName == pickfirstName
214+
isPickFirst = cc.curBalancerName == PickFirstBalancerName
186215
cc.mu.Unlock()
187216
if isPickFirst {
188217
break
@@ -210,7 +239,7 @@ func TestSwitchBalancerGRPCLBSecond(t *testing.T) {
210239
var isPickFirst bool
211240
for i := 0; i < 5000; i++ {
212241
cc.mu.Lock()
213-
isPickFirst = cc.curBalancerName == pickfirstName
242+
isPickFirst = cc.curBalancerName == PickFirstBalancerName
214243
cc.mu.Unlock()
215244
if isPickFirst {
216245
break
@@ -258,7 +287,7 @@ func TestSwitchBalancerGRPCLBSecond(t *testing.T) {
258287
r.NewAddress([]resolver.Address{{Addr: "backend"}})
259288
for i := 0; i < 5000; i++ {
260289
cc.mu.Lock()
261-
isPickFirst = cc.curBalancerName == pickfirstName
290+
isPickFirst = cc.curBalancerName == PickFirstBalancerName
262291
cc.mu.Unlock()
263292
if isPickFirst {
264293
break
@@ -352,7 +381,7 @@ func TestSwitchBalancerGRPCLBServiceConfig(t *testing.T) {
352381
var isPickFirst bool
353382
for i := 0; i < 5000; i++ {
354383
cc.mu.Lock()
355-
isPickFirst = cc.curBalancerName == pickfirstName
384+
isPickFirst = cc.curBalancerName == PickFirstBalancerName
356385
cc.mu.Unlock()
357386
if isPickFirst {
358387
break

clientconn.go

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,8 @@ type dialOptions struct {
9595
scChan <-chan ServiceConfig
9696
copts transport.ConnectOptions
9797
callOptions []CallOption
98-
// This is to support v1 balancer.
98+
// This is used by v1 balancer dial option WithBalancer to support v1
99+
// balancer, and also by WithBalancerName dial option.
99100
balancerBuilder balancer.Builder
100101
// This is to support grpclb.
101102
resolverBuilder resolver.Builder
@@ -200,7 +201,8 @@ func WithDecompressor(dc Decompressor) DialOption {
200201

201202
// WithBalancer returns a DialOption which sets a load balancer with the v1 API.
202203
// Name resolver will be ignored if this DialOption is specified.
203-
// Deprecated: use the new balancer APIs in balancer package instead.
204+
//
205+
// Deprecated: use the new balancer APIs in balancer package and WithBalancerName.
204206
func WithBalancer(b Balancer) DialOption {
205207
return func(o *dialOptions) {
206208
o.balancerBuilder = &balancerWrapperBuilder{
@@ -209,12 +211,21 @@ func WithBalancer(b Balancer) DialOption {
209211
}
210212
}
211213

212-
// WithBalancerBuilder is for testing only. Users using custom balancers should
213-
// register their balancer and use service config to choose the balancer to use.
214-
func WithBalancerBuilder(b balancer.Builder) DialOption {
215-
// TODO(bar) remove this when switching balancer is done.
214+
// WithBalancerName sets the balancer that the ClientConn will be initialized
215+
// with. Balancer registered with balancerName will be used. This function
216+
// panics if no balancer was registered by balancerName.
217+
//
218+
// The balancer cannot be overridden by balancer option specified by service
219+
// config.
220+
//
221+
// This is an EXPERIMENTAL API.
222+
func WithBalancerName(balancerName string) DialOption {
223+
builder := balancer.Get(balancerName)
224+
if builder == nil {
225+
panic(fmt.Sprintf("grpc.WithBalancerName: no balancer is registered for name %v", balancerName))
226+
}
216227
return func(o *dialOptions) {
217-
o.balancerBuilder = b
228+
o.balancerBuilder = builder
218229
}
219230
}
220231

@@ -670,9 +681,9 @@ func (cc *ClientConn) handleResolvedAddrs(addrs []resolver.Address, err error) {
670681

671682
cc.curAddresses = addrs
672683

673-
if cc.dopts.balancerBuilder != nil && cc.balancerWrapper == nil {
674-
cc.balancerWrapper = newCCBalancerWrapper(cc, cc.dopts.balancerBuilder, cc.balancerBuildOpts)
675-
} else {
684+
if cc.dopts.balancerBuilder == nil {
685+
// Only look at balancer types and switch balancer if balancer dial
686+
// option is not set.
676687
var isGRPCLB bool
677688
for _, a := range addrs {
678689
if a.Type == resolver.GRPCLB {
@@ -697,11 +708,16 @@ func (cc *ClientConn) handleResolvedAddrs(addrs []resolver.Address, err error) {
697708
// - the first time handling non-grpclb addresses
698709
// (curBalancerName="grpclb", preBalancerName="")
699710
if newBalancerName == "" {
700-
newBalancerName = pickfirstName
711+
newBalancerName = PickFirstBalancerName
701712
}
702713
}
703714
cc.switchBalancer(newBalancerName)
715+
} else if cc.balancerWrapper == nil {
716+
// Balancer dial option was set, and this is the first time handling
717+
// resolved addresses. Build a balancer with dopts.balancerBuilder.
718+
cc.balancerWrapper = newCCBalancerWrapper(cc, cc.dopts.balancerBuilder, cc.balancerBuildOpts)
704719
}
720+
705721
cc.balancerWrapper.handleResolvedAddrs(addrs, nil)
706722
}
707723

@@ -724,7 +740,7 @@ func (cc *ClientConn) switchBalancer(name string) {
724740

725741
grpclog.Infof("ClientConn switching balancer to %q", name)
726742
if cc.dopts.balancerBuilder != nil {
727-
grpclog.Infoln("ignoring balancer switching: WithBalancer DialOption used instead")
743+
grpclog.Infoln("ignoring balancer switching: Balancer DialOption used instead")
728744
return
729745
}
730746
// TODO(bar switching) change this to two steps: drain and close.

grpclb.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,10 @@ func (x *balanceLoadClientStream) Recv() (*lbpb.LoadBalanceResponse, error) {
8383
return m, nil
8484
}
8585

86+
func init() {
87+
balancer.Register(newLBBuilder())
88+
}
89+
8690
// newLBBuilder creates a builder for grpclb.
8791
func newLBBuilder() balancer.Builder {
8892
return NewLBBuilderWithFallbackTimeout(defaultFallbackTimeout)

grpclb/grpclb_test.go

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ import (
3636
"github.com/golang/protobuf/proto"
3737
"golang.org/x/net/context"
3838
"google.golang.org/grpc"
39+
"google.golang.org/grpc/balancer"
3940
"google.golang.org/grpc/codes"
4041
"google.golang.org/grpc/credentials"
4142
lbmpb "google.golang.org/grpc/grpclb/grpc_lb_v1/messages"
@@ -573,6 +574,24 @@ func TestBalancerDisconnects(t *testing.T) {
573574
t.Fatalf("No RPC sent to second backend after 1 second")
574575
}
575576

577+
type customGRPCLBBuilder struct {
578+
balancer.Builder
579+
name string
580+
}
581+
582+
func (b *customGRPCLBBuilder) Name() string {
583+
return b.name
584+
}
585+
586+
const grpclbCustomFallbackName = "grpclb_with_custom_fallback_timeout"
587+
588+
func init() {
589+
balancer.Register(&customGRPCLBBuilder{
590+
Builder: grpc.NewLBBuilderWithFallbackTimeout(100 * time.Millisecond),
591+
name: grpclbCustomFallbackName,
592+
})
593+
}
594+
576595
func TestFallback(t *testing.T) {
577596
defer leakcheck.Check(t)
578597

@@ -611,7 +630,7 @@ func TestFallback(t *testing.T) {
611630
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
612631
defer cancel()
613632
cc, err := grpc.DialContext(ctx, r.Scheme()+":///"+beServerName,
614-
grpc.WithBalancerBuilder(grpc.NewLBBuilderWithFallbackTimeout(100*time.Millisecond)),
633+
grpc.WithBalancerName(grpclbCustomFallbackName),
615634
grpc.WithTransportCredentials(&creds), grpc.WithDialer(fakeNameDialer))
616635
if err != nil {
617636
t.Fatalf("Failed to dial to the backend %v", err)

grpclb_remote_balancer.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -241,7 +241,7 @@ func (lb *lbBalancer) dialRemoteLB(remoteLBName string) {
241241
dopts = append(dopts, withContextDialer(lb.opt.Dialer))
242242
}
243243
// Explicitly set pickfirst as the balancer.
244-
dopts = append(dopts, WithBalancerBuilder(newPickfirstBuilder()))
244+
dopts = append(dopts, WithBalancerName(PickFirstBalancerName))
245245
dopts = append(dopts, withResolverBuilder(lb.manualResolver))
246246
// Dial using manualResolver.Scheme, which is a random scheme generated
247247
// when init grpclb. The target name is not important.

pickfirst.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@ import (
2626
"google.golang.org/grpc/resolver"
2727
)
2828

29-
const pickfirstName = "pick_first"
29+
// PickFirstBalancerName is the name of the pick_first balancer.
30+
const PickFirstBalancerName = "pick_first"
3031

3132
func newPickfirstBuilder() balancer.Builder {
3233
return &pickfirstBuilder{}
@@ -39,7 +40,7 @@ func (*pickfirstBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions
3940
}
4041

4142
func (*pickfirstBuilder) Name() string {
42-
return pickfirstName
43+
return PickFirstBalancerName
4344
}
4445

4546
type pickfirstBalancer struct {

0 commit comments

Comments
 (0)