From 55f8d1d2761cfe92abf51e90b1835b14fffcc2b4 Mon Sep 17 00:00:00 2001
From: nam <nam.nguyen@datadoghq.com>
Date: Tue, 17 Nov 2020 15:31:49 +0100
Subject: [PATCH] [feature] add client side load balancing

---
 runner/options.go      | 12 ++++++++++++
 runner/options_test.go |  4 +++-
 runner/requester.go    |  4 ++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/runner/options.go b/runner/options.go
index c5a1bf66..72afe9bb 100644
--- a/runner/options.go
+++ b/runner/options.go
@@ -66,6 +66,8 @@ type RunConfig struct {
 	// data
 	data []byte
 
+	// lbStrategy
+	lbStrategy string
 	// data func
 	dataFunc BinaryDataFunc
 
@@ -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 {
diff --git a/runner/options_test.go b/runner/options_test.go
index b4598212..5472f13e 100644
--- a/runner/options_test.go
+++ b/runner/options_test.go
@@ -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"),
 		)
@@ -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)
diff --git a/runner/requester.go b/runner/requester.go
index ca4894ee..5231a9e0 100644
--- a/runner/requester.go
+++ b/runner/requester.go
@@ -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...)
 }