From fd2f4eb505e6cbfda9652d13a7dee3cbfad4e647 Mon Sep 17 00:00:00 2001
From: Rizki Maulana Akbar <rizkimaulanaakbar@outlook.com>
Date: Wed, 14 Oct 2020 12:47:08 +0700
Subject: [PATCH] feat: add bucket histogram option

---
 attacker/attacker.go |  9 ++++++++-
 main.go              | 25 +++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/attacker/attacker.go b/attacker/attacker.go
index 9d8dea0..33c39ed 100644
--- a/attacker/attacker.go
+++ b/attacker/attacker.go
@@ -43,6 +43,7 @@ type Options struct {
 	Connections int
 	HTTP2       bool
 	LocalAddr   net.IPAddr
+	Buckets     []time.Duration
 
 	Attacker Attacker
 }
@@ -100,7 +101,13 @@ func Attack(ctx context.Context, target string, resCh chan *Result, metricsCh ch
 		Header: opts.Header,
 	})
 
-	metrics := &vegeta.Metrics{}
+	var metrics *vegeta.Metrics
+	if len(opts.Buckets) > 0 {
+		histogram := &vegeta.Histogram{Buckets: opts.Buckets}
+		metrics = &vegeta.Metrics{Histogram: histogram}
+	} else {
+		metrics = &vegeta.Metrics{}
+	}
 
 	child, cancelChild := context.WithCancel(ctx)
 	defer cancelChild()
diff --git a/main.go b/main.go
index 89fe9b8..c9401df 100644
--- a/main.go
+++ b/main.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"log"
 	"net"
 	"net/http"
 	"net/url"
@@ -45,6 +46,7 @@ type cli struct {
 	noHTTP2      bool
 	localAddress string
 	noKeepAlive  bool
+	buckets      string
 
 	debug   bool
 	version bool
@@ -81,6 +83,7 @@ func parseFlags(stdout, stderr io.Writer) (*cli, error) {
 	flagSet.IntVarP(&c.connections, "connections", "c", attacker.DefaultConnections, "Amount of maximum open idle connections per target host")
 	flagSet.BoolVar(&c.noHTTP2, "no-http2", false, "Don't issue HTTP/2 requests to servers which support it.")
 	flagSet.StringVar(&c.localAddress, "local-addr", "0.0.0.0", "Local IP address.")
+	flagSet.StringVar(&c.buckets, "buckets", "", "Histogram buckets.")
 	flagSet.Usage = c.usage
 	if err := flagSet.Parse(os.Args[1:]); err != nil {
 		if !errors.Is(err, flag.ErrHelp) {
@@ -191,6 +194,7 @@ func (c *cli) makeOptions() (*attacker.Options, error) {
 		Connections: c.connections,
 		HTTP2:       !c.noHTTP2,
 		LocalAddr:   localAddr,
+		Buckets:     parseBucketOptions(c.buckets),
 	}, nil
 }
 
@@ -202,6 +206,27 @@ func validateMethod(method string) bool {
 	return false
 }
 
+func parseBucketOptions(rawBuckets string) []time.Duration {
+	if rawBuckets == "" {
+		return []time.Duration{}
+	}
+
+	var result []time.Duration
+
+	stringBuckets := strings.Split(rawBuckets, ",")
+
+	for _, bucket := range stringBuckets {
+		trimmedBucket := strings.TrimSpace(bucket)
+		d, err := time.ParseDuration(trimmedBucket)
+		if err != nil {
+			log.Fatalf("Invalid bucket time format")
+		}
+		result = append(result, d)
+	}
+
+	return result
+}
+
 // Makes a new file under the working directory only when debug use.
 func setDebug(w io.Writer, debug bool) {
 	if !debug {