From 4a9ef9223dc08c6949441dd20d5c73c0aa1fc448 Mon Sep 17 00:00:00 2001
From: Carlo Alberto Ferraris <cafxx@mercari.com>
Date: Fri, 5 Feb 2021 12:46:25 +0900
Subject: [PATCH] add(cmd/ghz): expose --lb-strategy

Fixes #258
---
 cmd/ghz/main.go   | 9 +++++++++
 runner/config.go  | 1 +
 runner/options.go | 1 +
 3 files changed, 11 insertions(+)

diff --git a/cmd/ghz/main.go b/cmd/ghz/main.go
index e496e9e8..d4102728 100644
--- a/cmd/ghz/main.go
+++ b/cmd/ghz/main.go
@@ -254,6 +254,10 @@ var (
 	isEnableCompressionSet = false
 	enableCompression      = kingpin.Flag("enable-compression", "Enable Gzip compression on requests.").
 				Short('e').Default("false").IsSetByUser(&isEnableCompressionSet).Bool()
+
+	isLBStrategySet = false
+	lbStrategy      = kingpin.Flag("lb-strategy", "Client load balancing strategy.").
+			PlaceHolder(" ").IsSetByUser(&isLBStrategySet).String()
 )
 
 func main() {
@@ -471,6 +475,7 @@ func createConfigFromArgs(cfg *runner.Config) error {
 	cfg.CStepDuration = runner.Duration(*cStepDuration)
 	cfg.CMaxDuration = runner.Duration(*cMaxDuration)
 	cfg.CountErrors = *countErrors
+	cfg.LBStrategy = *lbStrategy
 
 	return nil
 }
@@ -646,6 +651,10 @@ func mergeConfig(dest *runner.Config, src *runner.Config) error {
 		dest.Host = src.Host
 	}
 
+	if isLBStrategySet {
+		dest.LBStrategy = src.LBStrategy
+	}
+
 	// load
 
 	if isAsyncSet {
diff --git a/runner/config.go b/runner/config.go
index 2d786b5d..917c9274 100644
--- a/runner/config.go
+++ b/runner/config.go
@@ -111,6 +111,7 @@ type Config struct {
 	LoadStep              int               `json:"load-step" toml:"load-step" yaml:"load-step"`
 	LoadStepDuration      Duration          `json:"load-step-duration" toml:"load-step-duration" yaml:"load-step-duration"`
 	LoadMaxDuration       Duration          `json:"load-max-duration" toml:"load-max-duration" yaml:"load-max-duration"`
+	LBStrategy            string            `json:"lb-strategy" toml:"lb-strategy" yaml:"lb-strategy"`
 }
 
 func checkData(data interface{}) error {
diff --git a/runner/options.go b/runner/options.go
index aa3f97d5..5c8598ae 100644
--- a/runner/options.go
+++ b/runner/options.go
@@ -1114,6 +1114,7 @@ func fromConfig(cfg *Config) []Option {
 		WithLoadStepDuration(time.Duration(cfg.LoadStepDuration)),
 		WithLoadEnd(cfg.LoadEnd),
 		WithLoadDuration(time.Duration(cfg.LoadMaxDuration)),
+		WithClientLoadBalancing(cfg.LBStrategy),
 		WithAsync(cfg.Async),
 		WithConcurrencySchedule(cfg.CSchedule),
 		WithConcurrencyStart(cfg.CStart),