From ab9881471a54d26d6d5a7ee61d3c453672d5e179 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Fri, 4 May 2018 16:35:39 +0400
Subject: [PATCH] [tm-bench] give user ability to change rpc function (#91)

Closes #17
---
 tm-bench/main.go       | 16 +++++++++++-----
 tm-bench/transacter.go | 28 +++++++++++++++-------------
 2 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/tm-bench/main.go b/tm-bench/main.go
index 91e67ba7f..6bbe4f220 100644
--- a/tm-bench/main.go
+++ b/tm-bench/main.go
@@ -29,19 +29,20 @@ type statistics struct {
 func main() {
 	var duration, txsRate, connections int
 	var verbose bool
-	var outputFormat string
+	var outputFormat, broadcastTxMethod string
 
 	flag.IntVar(&connections, "c", 1, "Connections to keep open per endpoint")
 	flag.IntVar(&duration, "T", 10, "Exit after the specified amount of time in seconds")
 	flag.IntVar(&txsRate, "r", 1000, "Txs per second to send in a connection")
 	flag.StringVar(&outputFormat, "output-format", "plain", "Output format: plain or json")
+	flag.StringVar(&broadcastTxMethod, "broadcast-tx-method", "async", "Broadcast method: async (no guarantees; fastest), sync (ensures tx is checked) or commit (ensures tx is checked and committed; slowest)")
 	flag.BoolVar(&verbose, "v", false, "Verbose output")
 
 	flag.Usage = func() {
 		fmt.Println(`Tendermint blockchain benchmarking tool.
 
 Usage:
-	tm-bench [-c 1] [-T 10] [-r 1000] [endpoints] [-output-format <plain|json>]
+	tm-bench [-c 1] [-T 10] [-r 1000] [endpoints] [-output-format <plain|json> [-broadcast-tx-method <async|sync|commit>]]
 
 Examples:
 	tm-bench localhost:46657`)
@@ -75,6 +76,11 @@ Examples:
 		fmt.Printf("Running %ds test @ %s\n", duration, flag.Arg(0))
 	}
 
+	if broadcastTxMethod != "async" && broadcastTxMethod != "sync" && broadcastTxMethod != "commit" {
+		fmt.Fprintln(os.Stderr, "broadcast-tx-method should be either 'sync', 'async' or 'commit'.")
+		os.Exit(1)
+	}
+
 	endpoints := strings.Split(flag.Arg(0), ",")
 
 	client := tmrpc.NewHTTP(endpoints[0], "/websocket")
@@ -86,7 +92,7 @@ Examples:
 	timeStart := time.Now()
 	logger.Info("Time started", "t", timeStart)
 
-	transacters := startTransacters(endpoints, connections, txsRate)
+	transacters := startTransacters(endpoints, connections, txsRate, "broadcast_tx_"+broadcastTxMethod)
 
 	select {
 	case <-time.After(time.Duration(duration) * time.Second):
@@ -164,11 +170,11 @@ func secondsSinceTimeStart(timeStart, timePassed time.Time) int64 {
 	return int64(timePassed.Sub(timeStart).Seconds())
 }
 
-func startTransacters(endpoints []string, connections int, txsRate int) []*transacter {
+func startTransacters(endpoints []string, connections, txsRate int, broadcastTxMethod string) []*transacter {
 	transacters := make([]*transacter, len(endpoints))
 
 	for i, e := range endpoints {
-		t := newTransacter(e, connections, txsRate)
+		t := newTransacter(e, connections, txsRate, broadcastTxMethod)
 		t.SetLogger(logger)
 		if err := t.Start(); err != nil {
 			fmt.Fprintln(os.Stderr, err)
diff --git a/tm-bench/transacter.go b/tm-bench/transacter.go
index a86a777f9..d5316d9f4 100644
--- a/tm-bench/transacter.go
+++ b/tm-bench/transacter.go
@@ -31,9 +31,10 @@ const (
 )
 
 type transacter struct {
-	Target      string
-	Rate        int
-	Connections int
+	Target            string
+	Rate              int
+	Connections       int
+	BroadcastTxMethod string
 
 	conns   []*websocket.Conn
 	wg      sync.WaitGroup
@@ -42,13 +43,14 @@ type transacter struct {
 	logger log.Logger
 }
 
-func newTransacter(target string, connections int, rate int) *transacter {
+func newTransacter(target string, connections, rate int, broadcastTxMethod string) *transacter {
 	return &transacter{
-		Target:      target,
-		Rate:        rate,
-		Connections: connections,
-		conns:       make([]*websocket.Conn, connections),
-		logger:      log.NewNopLogger(),
+		Target:            target,
+		Rate:              rate,
+		Connections:       connections,
+		BroadcastTxMethod: broadcastTxMethod,
+		conns:             make([]*websocket.Conn, connections),
+		logger:            log.NewNopLogger(),
 	}
 }
 
@@ -151,19 +153,19 @@ func (t *transacter) sendLoop(connIndex int) {
 			for i := 0; i < t.Rate; i++ {
 				// each transaction embeds connection index, tx number and hash of the hostname
 				tx := generateTx(connIndex, txNumber, hostnameHash)
-				paramsJson, err := json.Marshal(map[string]interface{}{"tx": hex.EncodeToString(tx)})
+				paramsJSON, err := json.Marshal(map[string]interface{}{"tx": hex.EncodeToString(tx)})
 				if err != nil {
 					fmt.Printf("failed to encode params: %v\n", err)
 					os.Exit(1)
 				}
-				rawParamsJson := json.RawMessage(paramsJson)
+				rawParamsJSON := json.RawMessage(paramsJSON)
 
 				c.SetWriteDeadline(time.Now().Add(sendTimeout))
 				err = c.WriteJSON(rpctypes.RPCRequest{
 					JSONRPC: "2.0",
 					ID:      "tm-bench",
-					Method:  "broadcast_tx_async",
-					Params:  rawParamsJson,
+					Method:  t.BroadcastTxMethod,
+					Params:  rawParamsJSON,
 				})
 				if err != nil {
 					fmt.Printf("%v. Try reducing the connections count and increasing the rate.\n", errors.Wrap(err, "txs send failed"))