From 28707121bb7840461a563aabcb510a9636e5c5a1 Mon Sep 17 00:00:00 2001
From: Felix Kaiser <felix.kaiser@fxkr.net>
Date: Mon, 24 Jun 2019 10:22:47 -0700
Subject: [PATCH] Include histogram in JSON output (#395)

* report: Add -buckets option for -type=json

Closes #394

* report: Add support for -buckets option for -type=hist too
---
 README.md        | 10 ++++++++++
 lib/histogram.go | 20 ++++++++++++++++++++
 lib/metrics.go   |  6 ++++++
 report.go        | 22 ++++++++++++++++------
 4 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index d1e5aabc..fe9146a4 100644
--- a/README.md
+++ b/README.md
@@ -363,6 +363,8 @@ Options:
   --type    Which report type to generate (text | json | hist[buckets]).
             [default: text]
 
+  --buckets Histogram buckets, e.g.: '[0,1ms,10ms]'
+
   --every   Write the report to --output at every given interval (e.g 100ms)
             The default of 0 means the report will only be written after
             all results have been processed. [default: 0]
@@ -434,6 +436,7 @@ The `Error Set` shows a unique set of errors returned by all issued requests. Th
     "99th": 3530000,
     "max": 3660505
   },
+  "buckets": {"0":9952,"1000000":40,"2000000":6,"3000000":0,"4000000":0,"5000000":2},
   "bytes_in": {
     "total": 606700,
     "mean": 6067
@@ -457,6 +460,13 @@ The `Error Set` shows a unique set of errors returned by all issued requests. Th
 }
 ```
 
+In the `buckets` field, each key is a nanosecond value representing the lower bound of a bucket.
+The upper bound is implied by the next higher bucket.
+Upper bounds are non-inclusive.
+The highest bucket is the overflow bucket; it has no upper bound.
+The values are counts of how many requests fell into that particular bucket.
+If the `-buckets` parameter is not present, the `buckets` field is omitted.
+
 #### `report -type=hist`
 
 Computes and prints a text based histogram for the given buckets.
diff --git a/lib/histogram.go b/lib/histogram.go
index dd16ba17..f8fd8446 100644
--- a/lib/histogram.go
+++ b/lib/histogram.go
@@ -1,6 +1,7 @@
 package vegeta
 
 import (
+	"bytes"
 	"fmt"
 	"strings"
 	"time"
@@ -35,6 +36,25 @@ func (h *Histogram) Add(r *Result) {
 	h.Counts[i]++
 }
 
+// MarshalJSON returns a JSON encoding of the buckets and their counts.
+func (h *Histogram) MarshalJSON() ([]byte, error) {
+	var buf bytes.Buffer
+
+	// Custom marshalling to guarantee order.
+	buf.WriteString("{")
+	for i := range h.Buckets {
+		if i > 0 {
+			buf.WriteString(", ")
+		}
+		if _, err := fmt.Fprintf(&buf, "\"%d\": %d", h.Buckets[i], h.Counts[i]); err != nil {
+			return nil, err
+		}
+	}
+	buf.WriteString("}")
+
+	return buf.Bytes(), nil
+}
+
 // Nth returns the nth bucket represented as a string.
 func (bs Buckets) Nth(i int) (left, right string) {
 	if i >= len(bs)-1 {
diff --git a/lib/metrics.go b/lib/metrics.go
index 20b03857..051117e5 100644
--- a/lib/metrics.go
+++ b/lib/metrics.go
@@ -12,6 +12,8 @@ import (
 type Metrics struct {
 	// Latencies holds computed request latency metrics.
 	Latencies LatencyMetrics `json:"latencies"`
+	// Histogram, only if requested
+	Histogram *Histogram `json:"buckets,omitempty"`
 	// BytesIn holds computed incoming byte metrics.
 	BytesIn ByteMetrics `json:"bytes_in"`
 	// BytesOut holds computed outgoing byte metrics.
@@ -75,6 +77,10 @@ func (m *Metrics) Add(r *Result) {
 			m.Errors = append(m.Errors, r.Error)
 		}
 	}
+
+	if m.Histogram != nil {
+		m.Histogram.Add(r)
+	}
 }
 
 // Close implements the Close method of the Report interface by computing
diff --git a/report.go b/report.go
index abec15a4..e1d4bca7 100644
--- a/report.go
+++ b/report.go
@@ -40,6 +40,7 @@ func reportCmd() command {
 	typ := fs.String("type", "text", "Report type to generate [text, json, hist[buckets]]")
 	every := fs.Duration("every", 0, "Report interval")
 	output := fs.String("output", "stdout", "Output file")
+	buckets := fs.String("buckets", "", "Histogram buckets, e.g.: \"[0,1ms,10ms]\"")
 
 	fs.Usage = func() {
 		fmt.Fprintln(os.Stderr, reportUsage)
@@ -51,11 +52,11 @@ func reportCmd() command {
 		if len(files) == 0 {
 			files = append(files, "stdin")
 		}
-		return report(files, *typ, *output, *every)
+		return report(files, *typ, *output, *every, *buckets)
 	}}
 }
 
-func report(files []string, typ, output string, every time.Duration) error {
+func report(files []string, typ, output string, every time.Duration, bucketsStr string) error {
 	if len(typ) < 4 {
 		return fmt.Errorf("invalid report type: %s", typ)
 	}
@@ -85,13 +86,22 @@ func report(files []string, typ, output string, every time.Duration) error {
 		rep, report = vegeta.NewTextReporter(&m), &m
 	case "json":
 		var m vegeta.Metrics
+		if bucketsStr != "" {
+			m.Histogram = &vegeta.Histogram{}
+			if err := m.Histogram.Buckets.UnmarshalText([]byte(bucketsStr)); err != nil {
+				return err
+			}
+		}
 		rep, report = vegeta.NewJSONReporter(&m), &m
 	case "hist":
-		if len(typ) < 6 {
-			return fmt.Errorf("bad buckets: '%s'", typ[4:])
-		}
 		var hist vegeta.Histogram
-		if err := hist.Buckets.UnmarshalText([]byte(typ[4:])); err != nil {
+		if bucketsStr == "" { // Old way
+			if len(typ) < 6 {
+				return fmt.Errorf("bad buckets: '%s'", typ[4:])
+			}
+			bucketsStr = typ[4:]
+		}
+		if err := hist.Buckets.UnmarshalText([]byte(bucketsStr)); err != nil {
 			return err
 		}
 		rep, report = vegeta.NewHistogramReporter(&hist), &hist