diff --git a/README.md b/README.md index b30e939..ee87093 100644 --- a/README.md +++ b/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) 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..117c4a4 100644 --- a/main.go +++ b/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 { diff --git a/main_test.go b/main_test.go index 820ded7..61c1ba7 100644 --- a/main_test.go +++ b/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, },