Merge pull request #68 from rzkmak/feat/histogram-buckets

Add bucket histogram option

README.md

@@ -72,6 +72,7 @@ Usage:
 Flags:
   -b, --body string         A request body to be sent.
   -B, --body-file string    The path to file whose content will be set as the http request body.
+      --buckets             Histogram buckets in comma-separated value (example: "10ms, 100ms, 500ms".
   -c, --connections int     Amount of maximum open idle connections per target host (default 10000)
       --debug               Run in debug mode.
   -d, --duration duration   The amount of time to issue requests to the targets. Give 0s for an infinite attack. (default 10s)

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()

main.go

@@ -45,6 +45,7 @@ type cli struct {
 	noHTTP2      bool
 	localAddress string
 	noKeepAlive  bool
+	buckets      string
 
 	debug   bool
 	version bool
@@ -81,6 +82,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; comma-separated list.")
 	flagSet.Usage = c.usage
 	if err := flagSet.Parse(os.Args[1:]); err != nil {
 		if !errors.Is(err, flag.ErrHelp) {
@@ -177,6 +179,12 @@ func (c *cli) makeOptions() (*attacker.Options, error) {
 
 	localAddr := net.IPAddr{IP: net.ParseIP(c.localAddress)}
 
+	parsedBuckets, err := parseBucketOptions(c.buckets)
+
+	if err != nil {
+		return nil, fmt.Errorf("wrong buckets format %w", err)
+	}
+
 	return &attacker.Options{
 		Rate:        c.rate,
 		Duration:    c.duration,
@@ -191,6 +199,7 @@ func (c *cli) makeOptions() (*attacker.Options, error) {
 		Connections: c.connections,
 		HTTP2:       !c.noHTTP2,
 		LocalAddr:   localAddr,
+		Buckets:     parsedBuckets,
 	}, nil
 }
 
@@ -202,6 +211,26 @@ func validateMethod(method string) bool {
 	return false
 }
 
+func parseBucketOptions(rawBuckets string) ([]time.Duration, error) {
+	if rawBuckets == "" {
+		return []time.Duration{}, nil
+	}
+
+	stringBuckets := strings.Split(rawBuckets, ",")
+	result := make([]time.Duration, len(stringBuckets))
+
+	for _, bucket := range stringBuckets {
+		trimmedBucket := strings.TrimSpace(bucket)
+		d, err := time.ParseDuration(trimmedBucket)
+		if err != nil {
+			return nil, err
+		}
+		result = append(result, d)
+	}
+
+	return result, nil
+}
+
 // Makes a new file under the working directory only when debug use.
 func setDebug(w io.Writer, debug bool) {
 	if !debug {

main_test.go

@@ -203,6 +203,7 @@ func TestMakeOptions(t *testing.T) {
 				MaxBody:    0,
 				HTTP2:      true,
 				KeepAlive:  true,
+				Buckets:    []time.Duration{},
 			},
 			wantErr: false,
 		},
@@ -229,6 +230,7 @@ func TestMakeOptions(t *testing.T) {
 				MaxBody:    0,
 				HTTP2:      true,
 				KeepAlive:  true,
+				Buckets:    []time.Duration{},
 			},
 			wantErr: false,
 		},
@@ -267,6 +269,7 @@ func TestMakeOptions(t *testing.T) {
 				MaxBody:    0,
 				HTTP2:      false,
 				KeepAlive:  true,
+				Buckets:    []time.Duration{},
 			},
 			wantErr: false,
 		},
@@ -294,6 +297,7 @@ func TestMakeOptions(t *testing.T) {
 				MaxBody:    0,
 				HTTP2:      true,
 				KeepAlive:  false,
+				Buckets:    []time.Duration{},
 			},
 			wantErr: false,
 		},