Skip to content

Commit

Permalink
[feature] add client side load balancing
Browse files Browse the repository at this point in the history
  • Loading branch information
nam committed Nov 17, 2020
1 parent dc6b53f commit 55f8d1d
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 1 deletion.
12 changes: 12 additions & 0 deletions runner/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,8 @@ type RunConfig struct {
// data
data []byte

// lbStrategy
lbStrategy string
// data func
dataFunc BinaryDataFunc

Expand Down Expand Up @@ -310,6 +312,16 @@ func WithBinaryData(data []byte) Option {
}
}

// WithClientLoadBalancing specifies the LB strategy to use
// The strategies has to be self written and pre defined
func WithClientLoadBalancing(strategy string) Option {
return func(o *RunConfig) error {
o.lbStrategy = strategy

return nil
}
}

// WithBinaryDataFunc specifies the binary data func which will be called on each request
// WithBinaryDataFunc(changeFunc)
func WithBinaryDataFunc(data func(mtd *desc.MethodDescriptor, callData *CallData) []byte) Option {
Expand Down
4 changes: 3 additions & 1 deletion runner/options_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -125,8 +125,9 @@ func TestRunConfig_newRunConfig(t *testing.T) {
WithDialTimeout(time.Duration(30*time.Second)),
WithName("asdf"),
WithCPUs(4),
WithBinaryDataFunc(changeFunc),
WithBinaryData([]byte("asdf1234foobar")),
WithBinaryDataFunc(changeFunc),
WithClientLoadBalancing(`{"loadBalancingPolicy":"round_robin"}`),
WithMetadataFromFile("../testdata/metadata.json"),
WithProtoset("testdata/bundle.protoset"),
)
Expand All @@ -152,6 +153,7 @@ func TestRunConfig_newRunConfig(t *testing.T) {
assert.Equal(t, 4, c.cpus)
assert.Equal(t, "asdf", c.name)
assert.Equal(t, []byte("asdf1234foobar"), c.data)
assert.Equal(t, `{"loadBalancingPolicy":"round_robin"}`, c.lbStrategy)
funcName1 := runtime.FuncForPC(reflect.ValueOf(changeFunc).Pointer()).Name()
funcName2 := runtime.FuncForPC(reflect.ValueOf(c.dataFunc).Pointer()).Name()
assert.Equal(t, funcName1, funcName2)
Expand Down
4 changes: 4 additions & 0 deletions runner/requester.go
Original file line number Diff line number Diff line change
Expand Up @@ -322,6 +322,10 @@ func (b *Requester) newClientConn(withStatsHandler bool) (*grpc.ClientConn, erro
grpc.MaxCallSendMsgSize(math.MaxInt32),
))

if b.config.lbStrategy != "" {
opts = append(opts, grpc.WithBalancerName(b.config.lbStrategy))
}

// create client connection
return grpc.DialContext(ctx, b.config.host, opts...)
}
Expand Down

0 comments on commit 55f8d1d

Please sign in to comment.