Repository: yarpc/yab
Pull request #301 (2020-06-09 18:55:17)

Add JSON Format Option to Benchmark Output
These changes allow a user to specify a 'format' option when running yab. Currently, the two formats supported are 'json' and 'plaintext' with 'plaintext' being the default output that is currently outputted by yab. When a user specifies the 'json' output, fields such as benchmark parameters, latencies, and benchmarking summaries will be printed in JSON format. This output format has been tested in benchmark_test.go.

Example JSON Output
```
{
  "body": {
    "result": {
      "count": 2,
      "message": "hello"
    }
  },
  "headers": {
    "Content-Length": "38",
    "Content-Type": "application/vnd.apache.thrift.binary",
    "Date": "Thu, 18 Jun 2020 23:32:07 GMT",
    "Rpc-Service": "hello",
    "Rpc-Status": "success"
  },
  "statusCode": 200
}

{
  "benchmarkParameters": {
    "cpus": 8,
    "connections": 16,
    "concurrency": 1,
    "maxRequests": 0,
    "maxDuration": "1s",
    "maxRPS": 0
  },
  "latencies": {
    "0.5000": "224.56µs",
    "0.9000": "541.276µs",
    "0.9500": "780.953µs",
    "0.9900": "1.409557ms",
    "0.9990": "4.644993ms",
    "0.9995": "5.223061ms",
    "1.0000": "7.468056ms"
  },
  "summary": {
    "elapsedTimeSeconds": 1,
    "totalRequests": 51891,
    "rps": 51851.72
  }
}
```

CLAassistant: [![CLA assistant check](https://cla-assistant.io/pull/badge/not_signed)](https://cla-assistant.io/yarpc/yab?pullRequest=301) <br/>Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our [Contributor License Agreement](https://cla-assistant.io/yarpc/yab?pullRequest=301) before we can accept your contribution.<br/><sub>You have signed the CLA already but the status is still pending? Let us [recheck](https://cla-assistant.io/check/yarpc/yab?pullRequest=301) it.</sub>
codecov[bot]: # [Codecov](https://codecov.io/gh/yarpc/yab/pull/301?src=pr&el=h1) Report
> Merging [#301](https://codecov.io/gh/yarpc/yab/pull/301?src=pr&el=desc) into [dev](https://codecov.io/gh/yarpc/yab/commit/968dcc9358bd84b243b5389af32db82a7f5724f9&el=desc) will **increase** coverage by `0.09%`.
> The diff coverage is `94.59%`.

[![Impacted file tree graph](https://codecov.io/gh/yarpc/yab/pull/301/graphs/tree.svg?width=650&height=150&src=pr&token=yEgrqUTCuh)](https://codecov.io/gh/yarpc/yab/pull/301?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##              dev     #301      +/-   ##
==========================================
+ Coverage   80.84%   80.93%   +0.09%     
==========================================
  Files          59       59              
  Lines        3190     3216      +26     
==========================================
+ Hits         2579     2603      +24     
- Misses        513      514       +1     
- Partials       98       99       +1     
```


| [Impacted Files](https://codecov.io/gh/yarpc/yab/pull/301?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [options.go](https://codecov.io/gh/yarpc/yab/pull/301/diff?src=pr&el=tree#diff-b3B0aW9ucy5nbw==) | `94.73% <ø> (ø)` | |
| [benchmark.go](https://codecov.io/gh/yarpc/yab/pull/301/diff?src=pr&el=tree#diff-YmVuY2htYXJrLmdv) | `93.10% <93.33%> (-0.72%)` | :arrow_down: |
| [bench\_state.go](https://codecov.io/gh/yarpc/yab/pull/301/diff?src=pr&el=tree#diff-YmVuY2hfc3RhdGUuZ28=) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/yarpc/yab/pull/301?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/yarpc/yab/pull/301?src=pr&el=footer). Last update [968dcc9...2dcd649](https://codecov.io/gh/yarpc/yab/pull/301?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

jaricftw: LGTM, @prashantv can you also take another look

Patch:
diff --git a/bench_state.go b/bench_state.go
index 415403c..88a7317 100644
--- a/bench_state.go
+++ b/bench_state.go
@@ -79,13 +79,14 @@ func (s *benchmarkState) recordLatency(d time.Duration) {
 	s.statter.Timing("latency", d)
 }
 
-func (s *benchmarkState) printLatencies(out output) {
-	// TODO JSON output?
+// Returns a mapping of quantiles to latency values
+func (s *benchmarkState) getLatencies() map[float64]time.Duration {
 	sort.Sort(byDuration(s.latencies))
-	out.Printf("Latencies:\n")
-	for _, quantile := range []float64{0.5, 0.9, 0.95, 0.99, 0.999, 0.9995, 1.0} {
-		out.Printf("  %.4f: %v\n", quantile, s.getQuantile(quantile))
+	latencyValues := make(map[float64]time.Duration, len(_quantiles))
+	for _, quantile := range _quantiles {
+		latencyValues[quantile] = s.getQuantile(quantile)
 	}
+	return latencyValues
 }
 
 func (s *benchmarkState) printErrors(out output) {
diff --git a/benchmark.go b/benchmark.go
index 4220998..c49481b 100644
--- a/benchmark.go
+++ b/benchmark.go
@@ -21,11 +21,14 @@
 package main
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
+	"math"
 	"os"
 	"os/signal"
 	"runtime"
+	"strings"
 	"sync"
 	"time"
 
@@ -39,8 +42,36 @@ import (
 var (
 	errNegativeDuration = errors.New("duration cannot be negative")
 	errNegativeMaxReqs  = errors.New("max requests cannot be negative")
+
+	// using a global _quantiles slice mainly for ease of testing, and not passing
+	// the same array around to multiple functions
+	_quantiles = []float64{0.5000, 0.9000, 0.9500, 0.9900, 0.9990, 0.9995, 1.0000}
 )
 
+// Parameters holds values of all benchmark parameters
+type Parameters struct {
+	CPUs        int    `json:"cpus"`
+	Connections int    `json:"connections"`
+	Concurrency int    `json:"concurrency"`
+	MaxRequests int    `json:"maxRequests"`
+	MaxDuration string `json:"maxDuration"`
+	MaxRPS      int    `json:"maxRPS"`
+}
+
+// Summary stores the benchmarking summary
+type Summary struct {
+	ElapsedTimeSeconds float64 `json:"elapsedTimeSeconds"`
+	TotalRequests      int     `json:"totalRequests"`
+	RPS                float64 `json:"rps"`
+}
+
+// BenchmarkOutput stores benchmark settings and results for JSON output
+type BenchmarkOutput struct {
+	Parameters Parameters        `json:"benchmarkParameters"`
+	Latencies  map[string]string `json:"latencies"`
+	Summary    Summary           `json:"summary"`
+}
+
 // setGoMaxProcs sets runtime.GOMAXPROCS if the option is set
 // and returns the number of GOMAXPROCS configured.
 func (o BenchmarkOptions) setGoMaxProcs() int {
@@ -112,13 +143,27 @@ func runBenchmark(out output, logger *zap.Logger, allOpts Options, resolved reso
 
 	goMaxProcs := opts.setGoMaxProcs()
 	numConns := opts.getNumConnections(goMaxProcs)
-	out.Printf("Benchmark parameters:\n")
-	out.Printf("  CPUs:            %v\n", goMaxProcs)
-	out.Printf("  Connections:     %v\n", numConns)
-	out.Printf("  Concurrency:     %v\n", opts.Concurrency)
-	out.Printf("  Max requests:    %v\n", opts.MaxRequests)
-	out.Printf("  Max duration:    %v\n", opts.MaxDuration)
-	out.Printf("  Max RPS:         %v\n", opts.RPS)
+
+	parameters := Parameters{
+		CPUs:        goMaxProcs,
+		Connections: numConns,
+		Concurrency: opts.Concurrency,
+		MaxRequests: opts.MaxRequests,
+		MaxDuration: opts.MaxDuration.String(),
+		MaxRPS:      opts.RPS,
+	}
+
+	// If format is JSON, benchmark parameters are printed after benchmark is run to maintain a single JSON blob
+	formatAsJSON := false
+	switch format := strings.ToLower(opts.Format); format {
+	case "text", "":
+		printParameters(out, parameters)
+	case "json":
+		formatAsJSON = true
+	default:
+		out.Warnf("Unrecognized format option %q, please specify 'json' for JSON output. Printing plaintext output as default.\n\n", opts.Format)
+		printParameters(out, parameters)
+	}
 
 	// Warm up number of connections.
 	logger.Debug("Warming up connections.", zap.Int("numConns", numConns))
@@ -170,9 +215,6 @@ func runBenchmark(out output, logger *zap.Logger, allOpts Options, resolved reso
 		}
 	}
 
-	// TODO: Support streaming updates.
-	// TOOD: Aggregate per-backend state for JSON output in future.
-
 	// Wait for all the worker goroutines to end.
 	wg.Wait()
 	total := time.Since(start)
@@ -188,12 +230,73 @@ func runBenchmark(out output, logger *zap.Logger, allOpts Options, resolved reso
 		zap.Time("startTime", start),
 	)
 
+	// Print out errors
+	// TODO: allow errors to be printed in JSON format for output consistency
 	overall.printErrors(out)
-	overall.printLatencies(out)
 
-	out.Printf("Elapsed time:      %v\n", (total / time.Millisecond * time.Millisecond))
-	out.Printf("Total requests:    %v\n", overall.totalRequests)
-	out.Printf("RPS:               %.2f\n", float64(overall.totalRequests)/total.Seconds())
+	latencyValues := overall.getLatencies()
+
+	// Rounding RPS value to the hundredths place
+	rps := float64(overall.totalRequests) / total.Seconds()
+	rps = (math.Round(rps * 100)) / 100
+
+	summary := Summary{
+		ElapsedTimeSeconds: (total / time.Millisecond * time.Millisecond).Seconds(),
+		TotalRequests:      overall.totalRequests,
+		RPS:                rps,
+	}
+
+	if formatAsJSON {
+		outputJSON(out, parameters, latencyValues, summary)
+	} else {
+		outputPlaintext(out, latencyValues, summary)
+	}
+}
+
+func outputJSON(out output, parameters Parameters, latencyValues map[float64]time.Duration, summary Summary) {
+	latencies := make(map[string]string, len(_quantiles))
+	for _, quantile := range _quantiles {
+		latencies[fmt.Sprintf("%.4f", quantile)] = latencyValues[quantile].String()
+	}
+
+	benchmarkOutput := BenchmarkOutput{
+		Parameters: parameters,
+		Latencies:  latencies,
+		Summary:    summary,
+	}
+
+	jsonOutput, err := json.MarshalIndent(&benchmarkOutput, "" /* prefix */, "  " /* indent */)
+	if err != nil {
+		out.Fatalf("Failed to marshal benchmark output: %v\n", err)
+	}
+	out.Printf("%s\n", jsonOutput)
+}
+
+func outputPlaintext(out output, latencyValues map[float64]time.Duration, summary Summary) {
+	// Print out latencies
+	printLatencies(out, latencyValues)
+
+	// Print out summary
+	out.Printf("Elapsed time (seconds):   %.2f\n", summary.ElapsedTimeSeconds)
+	out.Printf("Total requests:           %v\n", summary.TotalRequests)
+	out.Printf("RPS:                      %.2f\n", summary.RPS)
+}
+
+func printParameters(out output, parameters Parameters) {
+	out.Printf("Benchmark parameters:\n")
+	out.Printf("  CPUs:            %v\n", parameters.CPUs)
+	out.Printf("  Connections:     %v\n", parameters.Connections)
+	out.Printf("  Concurrency:     %v\n", parameters.Concurrency)
+	out.Printf("  Max requests:    %v\n", parameters.MaxRequests)
+	out.Printf("  Max duration:    %v\n", parameters.MaxDuration)
+	out.Printf("  Max RPS:         %v\n", parameters.MaxRPS)
+}
+
+func printLatencies(out output, latencyValues map[float64]time.Duration) {
+	out.Printf("Latencies:\n")
+	for _, quantile := range _quantiles {
+		out.Printf("  %.4f: %v\n", quantile, latencyValues[quantile])
+	}
 }
 
 // stopOnInterrupt sets up a signal that will trigger the run to stop.
diff --git a/options.go b/options.go
index f11657c..9d9eea2 100644
--- a/options.go
+++ b/options.go
@@ -109,6 +109,7 @@ type BenchmarkOptions struct {
 	// Benchmark metrics can optionally be reported via statsd.
 	StatsdHostPort string `long:"statsd" description:"Optional host:port of a StatsD server to report metrics"`
 	PerPeerStats   bool   `long:"per-peer-stats" description:"Whether to emit stats by peer rather than aggregated"`
+	Format         string `long:"format" description:"Prints benchmark output in either text or JSON format. Default is text."`
 }
 
 func newOptions() *Options {


Test patch:
diff --git a/bench_state_test.go b/bench_state_test.go
index 335dded..37ad019 100644
--- a/bench_state_test.go
+++ b/bench_state_test.go
@@ -102,7 +102,7 @@ func TestBenchmarkStateLatencies(t *testing.T) {
 	assert.Equal(t, state.totalSuccess, 10001, "Success count mismatch")
 	assert.Equal(t, state.totalRequests, 10001, "Request count mismatch")
 
-	state.printLatencies(out)
+	printLatencies(out, state.getLatencies())
 
 	expected := []string{
 		"0.5000: 5ms",
@@ -147,7 +147,8 @@ func TestBenchmarkStateMergeLatencies(t *testing.T) {
 	assert.Equal(t, state1.totalRequests, 10001, "Request count mismatch")
 
 	buf, _, out := getOutput(t)
-	state1.printLatencies(out)
+
+	printLatencies(out, state1.getLatencies())
 
 	expected := []string{
 		"0.5000: 5ms",
diff --git a/benchmark_test.go b/benchmark_test.go
index efdbcca..669dc3b 100644
--- a/benchmark_test.go
+++ b/benchmark_test.go
@@ -21,6 +21,7 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
 	"os"
 	"sync"
@@ -31,6 +32,7 @@ import (
 	"github.com/yarpc/yab/transport"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/uber/tchannel-go/testutils"
 	"go.uber.org/atomic"
 )
@@ -218,3 +220,98 @@ func TestBenchmarkStatsPerPeer(t *testing.T) {
 	})
 	assert.Equal(t, want, statsServer.Aggregated(), "unexpected stats")
 }
+
+func TestBenchmarkOutput(t *testing.T) {
+	// Testing text, json, and unrecognized output
+	tests := []struct {
+		name          string
+		format        []string
+		wantJSON      bool
+		wantWarn      string
+		wantOutput    []string
+		notWantOutput []string
+	}{
+		{
+			name:          "json",
+			format:        []string{"json", "JSON", "Json", "jsON", "jsoN"},
+			wantJSON:      true,
+			wantWarn:      "",
+			wantOutput:    []string{"summary", "benchmarkParameters", "maxRPS", "latencies"},
+			notWantOutput: []string{"Errors", "Benchmark parameters", "Max RPS", "Unrecognized format option"},
+		},
+		{
+			name:          "text",
+			format:        []string{"text", ""},
+			wantJSON:      false,
+			wantWarn:      "",
+			wantOutput:    []string{"Benchmark parameters", "Max RPS"},
+			notWantOutput: []string{"Errors", "summary", "maxRPS", "Unrecognized format option"},
+		},
+		// Unimplemented format should print an error and default to plaintext
+		{
+			name:          "unrecognized",
+			format:        []string{"csv", "plaintext", "blob"},
+			wantJSON:      false,
+			wantWarn:      "Unrecognized format option",
+			wantOutput:    []string{"Benchmark parameters", "Max RPS"},
+			notWantOutput: []string{"Errors", "benchmarkParameters", "maxRPS"},
+		},
+	}
+
+	var requests atomic.Int32
+	s := newServer(t)
+	defer s.shutdown()
+	s.register(fooMethod, methods.errorIf(func() bool {
+		requests.Inc()
+		return false
+	}))
+	m := benchmarkMethodForTest(t, fooMethod, transport.TChannel)
+
+	for _, tt := range tests {
+		for _, format := range tt.format {
+			t.Run(tt.name, func(t *testing.T) {
+				requests.Store(0)
+				buf, bufWarn, out := getOutput(t)
+				opts := Options{
+					BOpts: BenchmarkOptions{
+						MaxRequests:    1,
+						MaxDuration:    100 * time.Millisecond,
+						Connections:    1,
+						WarmupRequests: 0,
+						Concurrency:    1,
+						Format:         format,
+					},
+					TOpts: s.transportOpts(),
+				}
+
+				runBenchmark(out, _testLogger, opts, _resolvedTChannelThrift, m)
+				bufStr := buf.String()
+				bufWarnStr := bufWarn.String()
+
+				for _, want := range tt.wantOutput {
+					assert.Contains(t, bufStr, want)
+				}
+				for _, notWant := range tt.notWantOutput {
+					assert.NotContains(t, bufStr, notWant)
+				}
+
+				if tt.wantWarn != "" {
+					assert.Contains(t, bufWarnStr, tt.wantWarn)
+				} else {
+					assert.Empty(t, bufWarnStr)
+				}
+
+				if tt.wantJSON {
+					// Creating struct from JSON output string
+					var benchmarkOutput BenchmarkOutput
+					b := []byte(bufStr)
+					err := json.Unmarshal(b, &benchmarkOutput)
+					require.NoError(t, err)
+					assert.Equal(t, benchmarkOutput.Parameters.MaxRPS, opts.BOpts.RPS)
+					assert.GreaterOrEqual(t, opts.BOpts.MaxRequests, benchmarkOutput.Summary.TotalRequests)
+					assert.Equal(t, len(_quantiles), len(benchmarkOutput.Latencies))
+				}
+			})
+		}
+	}
+}