added load balancer option (#49)

talariadb · Nov 18, 2020 · 6ccf06d · 6ccf06d
1 parent 14075e3
commit 6ccf06d
Show file tree

Hide file tree

Showing 3 changed files with 104 additions and 89 deletions.
diff --git a/client/golang/client.go b/client/golang/client.go
@@ -6,7 +6,6 @@ package client
 import (
 	"context"
 	"errors"
-	"fmt"
 	"time"
 
 	pb "github.com/kelindar/talaria/proto"
@@ -70,7 +69,11 @@ func (c *Client) connect() error {
 	timeoutCtx, cancel := context.WithTimeout(context.Background(), c.netconf.DialTimeout)
 	defer cancel()
 
-	dialOptions := []grpc.DialOption{}
+	var dialOptions []grpc.DialOption
+	if c.netconf.LoadBalancer != "" {
+		dialOptions = append(dialOptions, grpc.WithBalancerName(c.netconf.LoadBalancer))
+	}
+
 	if !c.netconf.NonBlocking {
 		dialOptions = append(dialOptions, grpc.WithBlock())
 	}
@@ -83,7 +86,6 @@ func (c *Client) connect() error {
 
 	conn, err := grpc.DialContext(timeoutCtx, c.netconf.Address, dialOptions...)
 	if err != nil {
-		fmt.Println("err", err)
 		return ErrUnableToConnect
 	}
 	c.ingress = pb.NewIngressClient(conn)

diff --git a/client/golang/client_opts.go b/client/golang/client_opts.go
@@ -1,59 +1,68 @@
-// Copyright 2019-2020 Grabtaxi Holdings PTE LTE (GRAB), All rights reserved.
-// Use of this source code is governed by an MIT-style license that can be found in the LICENSE file
-
-package client
-
-import (
-	"time"
-
-	"github.com/myteksi/hystrix-go/hystrix"
-	"google.golang.org/grpc/credentials"
-)
-
-// Option is a functional parameter used to configure the client.
-type Option func(client *Client)
-
-// netconf defines connection pool configuration for a gRPC service
-type netconf struct {
-	CircuitOptions map[string]hystrix.CommandConfig
-	Address        string
-	DialTimeout    time.Duration
-	Credentials    credentials.TransportCredentials
-	NonBlocking    bool // once set to true, the client will be returned before connection gets ready
-}
-
-// WithNetwork specifies the configuration for a connection.
-func WithNetwork(dialTimeout time.Duration) Option {
-	return func(client *Client) {
-		client.netconf.DialTimeout = dialTimeout
-	}
-}
-
-// WithCircuit specifies the configuration for the circuit breaker.
-func WithCircuit(timeout time.Duration, maxConcurrent, errorThresholdPercent int) Option {
-	return func(client *Client) {
-		if client.netconf.CircuitOptions == nil {
-			client.netconf.CircuitOptions = make(map[string]hystrix.CommandConfig, 1)
-		}
-
-		client.netconf.CircuitOptions[commandName] = hystrix.CommandConfig{
-			Timeout:               int(timeout / time.Millisecond),
-			MaxConcurrentRequests: maxConcurrent,
-			ErrorPercentThreshold: errorThresholdPercent,
-		}
-	}
-}
-
-// WithCredential specfies the configuration for the gRPC credentials based on TLS
-// with this one being set, connection will be created in secure manner
-func WithCredential(credentials credentials.TransportCredentials) Option {
-	return func(client *Client) {
-		client.netconf.Credentials = credentials
-	}
-}
-
-func WithNonBlock() Option {
-	return func(client *Client) {
-		client.netconf.NonBlocking = true
-	}
-}
+// Copyright 2019-2020 Grabtaxi Holdings PTE LTE (GRAB), All rights reserved.
+// Use of this source code is governed by an MIT-style license that can be found in the LICENSE file
+
+package client
+
+import (
+	"time"
+
+	"github.com/myteksi/hystrix-go/hystrix"
+	"google.golang.org/grpc/credentials"
+)
+
+// Option is a functional parameter used to configure the client.
+type Option func(client *Client)
+
+// netconf defines connection pool configuration for a gRPC service
+type netconf struct {
+	CircuitOptions map[string]hystrix.CommandConfig // Circuit configuration
+	Address        string                           // Endpoint of the server
+	DialTimeout    time.Duration                    // Dial timaout
+	Credentials    credentials.TransportCredentials // Transport credentials to use
+	NonBlocking    bool                             // once set to true, the client will be returned before connection gets ready
+	LoadBalancer   string                           // gRPC load balancing strategy
+}
+
+// WithNetwork specifies the configuration for a connection.
+func WithNetwork(dialTimeout time.Duration) Option {
+	return func(client *Client) {
+		client.netconf.DialTimeout = dialTimeout
+	}
+}
+
+// WithCircuit specifies the configuration for the circuit breaker.
+func WithCircuit(timeout time.Duration, maxConcurrent, errorThresholdPercent int) Option {
+	return func(client *Client) {
+		if client.netconf.CircuitOptions == nil {
+			client.netconf.CircuitOptions = make(map[string]hystrix.CommandConfig, 1)
+		}
+
+		client.netconf.CircuitOptions[commandName] = hystrix.CommandConfig{
+			Timeout:               int(timeout / time.Millisecond),
+			MaxConcurrentRequests: maxConcurrent,
+			ErrorPercentThreshold: errorThresholdPercent,
+		}
+	}
+}
+
+// WithCredential specfies the configuration for the gRPC credentials based on TLS
+// with this one being set, connection will be created in secure manner
+func WithCredential(credentials credentials.TransportCredentials) Option {
+	return func(client *Client) {
+		client.netconf.Credentials = credentials
+	}
+}
+
+// WithNonBlock creates a non-blocking gRPC dial()
+func WithNonBlock() Option {
+	return func(client *Client) {
+		client.netconf.NonBlocking = true
+	}
+}
+
+// WithLoadBalancer specifies the load balancer to use
+func WithLoadBalancer(name string) Option {
+	return func(client *Client) {
+		client.netconf.LoadBalancer = name
+	}
+}
diff --git a/client/golang/client_opts_test.go b/client/golang/client_opts_test.go
@@ -1,27 +1,31 @@
-package client
-
-import (
-	"testing"
-	"time"
-
-	"github.com/myteksi/hystrix-go/hystrix"
-
-	"github.com/stretchr/testify/assert"
-)
-
-func TestWithNetwork(t *testing.T) {
-	client, _ := Dial("invalid", WithNonBlock(), WithNetwork(10*time.Second))
-	assert.Equal(t, 10*time.Second, client.netconf.DialTimeout)
-}
-
-func TestWithCircuit(t *testing.T) {
-	client, _ := Dial("invalid", WithNonBlock(), WithCircuit(200*time.Millisecond, 500, 50))
-	expectedCircuit := map[string]hystrix.CommandConfig{
-		commandName: {
-			Timeout:               200,
-			MaxConcurrentRequests: 500,
-			ErrorPercentThreshold: 50,
-		},
-	}
-	assert.Equal(t, expectedCircuit, client.netconf.CircuitOptions)
-}
+package client
+
+import (
+	"testing"
+	"time"
+
+	"github.com/myteksi/hystrix-go/hystrix"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWithNetwork(t *testing.T) {
+	client, _ := Dial("invalid", WithNonBlock(), WithNetwork(10*time.Second))
+	assert.Equal(t, 10*time.Second, client.netconf.DialTimeout)
+}
+
+func TestWithLoadBalancer(t *testing.T) {
+	client, _ := Dial("invalid", WithNonBlock(), WithLoadBalancer("round_robin"))
+	assert.Equal(t, "round_robin", client.netconf.LoadBalancer)
+}
+
+func TestWithCircuit(t *testing.T) {
+	client, _ := Dial("invalid", WithNonBlock(), WithCircuit(200*time.Millisecond, 500, 50))
+	expectedCircuit := map[string]hystrix.CommandConfig{
+		commandName: {
+			Timeout:               200,
+			MaxConcurrentRequests: 500,
+			ErrorPercentThreshold: 50,
+		},
+	}
+	assert.Equal(t, expectedCircuit, client.netconf.CircuitOptions)
+}