Skip to content

Commit

Permalink
attack: Add -max-workers flag (tsenart#421)
Browse files Browse the repository at this point in the history 
* attack: Add -max-workers flag

This commit introduces the `-max-workers` flag and associated function
options. This allows controlling the concurrency level that vegeta uses,
but can sacrifice reaching the requested `-rate`, in case the specified
max number of workers is below the what is needed to sustain such rate.

* Updated README

* Update README.md
  • Loading branch information
@tsenart
tsenart authored Jul 19, 2019
1 parent f807d60 commit fb2d8b9
Show file tree
Hide file tree
Showing 3 changed files with 102 additions and 59 deletions.
83 changes: 45 additions & 38 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,85 +51,87 @@ Usage: vegeta [global flags] <command> [command flags]

global flags:
-cpus int
Number of CPUs to use (defaults to the number of CPUs you have)
Number of CPUs to use (defaults to the number of CPUs you have)
-profile string
Enable profiling of [cpu, heap]
Enable profiling of [cpu, heap]
-version
Print version and exit
Print version and exit

attack command:
-body string
Requests body file
Requests body file
-cert string
TLS client PEM encoded certificate file
TLS client PEM encoded certificate file
-connections int
Max open idle connections per target host (default 10000)
Max open idle connections per target host (default 10000)
-duration duration
Duration of the test [0 = forever]
Duration of the test [0 = forever]
-format string
Targets format [http, json] (default "http")
Targets format [http, json] (default "http")
-h2c
Send HTTP/2 requests without TLS encryption
Send HTTP/2 requests without TLS encryption
-header value
Request header
Request header
-http2
Send HTTP/2 requests when supported by the server (default true)
Send HTTP/2 requests when supported by the server (default true)
-insecure
Ignore invalid server TLS certificates
Ignore invalid server TLS certificates
-keepalive
Use persistent connections (default true)
Use persistent connections (default true)
-key string
TLS client PEM encoded private key file
TLS client PEM encoded private key file
-laddr value
Local IP address (default 0.0.0.0)
Local IP address (default 0.0.0.0)
-lazy
Read targets lazily
Read targets lazily
-max-body value
Maximum number of bytes to capture from response bodies. [-1 = no limit] (default -1)
Maximum number of bytes to capture from response bodies. [-1 = no limit] (default -1)
-max-workers uint
Maximum number of workers (default 18446744073709551615)
-name string
Attack name
Attack name
-output string
Output file (default "stdout")
Output file (default "stdout")
-rate value
Number of requests per time unit (default 50/1s)
Number of requests per time unit (default 50/1s)
-redirects int
Number of redirects to follow. -1 will not follow but marks as success (default 10)
Number of redirects to follow. -1 will not follow but marks as success (default 10)
-resolvers value
List of addresses (ip:port) to use for DNS resolution. Disables use of local system DNS. (comma separated list)
List of addresses (ip:port) to use for DNS resolution. Disables use of local system DNS. (comma separated list)
-root-certs value
TLS root certificate files (comma separated list)
TLS root certificate files (comma separated list)
-targets string
Targets file (default "stdin")
Targets file (default "stdin")
-timeout duration
Requests timeout (default 30s)
Requests timeout (default 30s)
-unix-socket string
Connect over a unix socket. This overrides the host address in target URLs
Connect over a unix socket. This overrides the host address in target URLs
-workers uint
Initial number of workers (default 10)


Initial number of workers (default 10)

encode command:
-output string
Output file (default "stdout")
Output file (default "stdout")
-to string
Output encoding [csv, gob, json] (default "json")
Output encoding [csv, gob, json] (default "json")

plot command:
-output string
Output file (default "stdout")
Output file (default "stdout")
-threshold int
Threshold of data points above which series are downsampled. (default 4000)
Threshold of data points above which series are downsampled. (default 4000)
-title string
Title and header of the resulting HTML page (default "Vegeta Plot")
Title and header of the resulting HTML page (default "Vegeta Plot")

report command:
-buckets string
Histogram buckets, e.g.: "[0,1ms,10ms]"
-every duration
Report interval
Report interval
-output string
Output file (default "stdout")
Output file (default "stdout")
-type string
Report type to generate [text, json, hist[buckets]] (default "text")
Report type to generate [text, json, hist[buckets], hdrplot] (default "text")

examples:
echo "GET http://localhost/" | vegeta attack -duration=5s | tee results.bin | vegeta report
Expand Down Expand Up @@ -346,7 +348,12 @@ timeouts.

Specifies the initial number of workers used in the attack. The actual
number of workers will increase if necessary in order to sustain the
requested rate.
requested rate, unless it'd go beyond `-max-workers`.

#### `-max-workers`

Specifies the maximum number of workers used in the attack. It can be used to
control the concurrency level used by an attack.

### `report` command

Expand Down
3 changes: 3 additions & 0 deletions attack.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ func attackCmd() command {
fs.DurationVar(&opts.duration, "duration", 0, "Duration of the test [0 = forever]")
fs.DurationVar(&opts.timeout, "timeout", vegeta.DefaultTimeout, "Requests timeout")
fs.Uint64Var(&opts.workers, "workers", vegeta.DefaultWorkers, "Initial number of workers")
fs.Uint64Var(&opts.maxWorkers, "max-workers", vegeta.DefaultMaxWorkers, "Maximum number of workers")
fs.IntVar(&opts.connections, "connections", vegeta.DefaultConnections, "Max open idle connections per target host")
fs.IntVar(&opts.redirects, "redirects", vegeta.DefaultRedirects, "Number of redirects to follow. -1 will not follow but marks as success")
fs.Var(&maxBodyFlag{&opts.maxBody}, "max-body", "Maximum number of bytes to capture from response bodies. [-1 = no limit]")
Expand Down Expand Up @@ -82,6 +83,7 @@ type attackOpts struct {
timeout time.Duration
rate vegeta.Rate
workers uint64
maxWorkers uint64
connections int
redirects int
maxBody int64
Expand Down Expand Up @@ -168,6 +170,7 @@ func attack(opts *attackOpts) (err error) {
vegeta.LocalAddr(*opts.laddr.IPAddr),
vegeta.TLSConfig(tlsc),
vegeta.Workers(opts.workers),
vegeta.MaxWorkers(opts.maxWorkers),
vegeta.KeepAlive(opts.keepalive),
vegeta.Connections(opts.connections),
vegeta.HTTP2(opts.http2),
Expand Down
75 changes: 54 additions & 21 deletions lib/attack.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"fmt"
"io"
"io/ioutil"
"math"
"net"
"net/http"
"net/url"
Expand All @@ -17,15 +18,16 @@ import (

// Attacker is an attack executor which wraps an http.Client
type Attacker struct {
dialer *net.Dialer
client http.Client
stopch chan struct{}
workers uint64
maxBody int64
redirects int
seqmu sync.Mutex
seq uint64
began time.Time
dialer *net.Dialer
client http.Client
stopch chan struct{}
workers uint64
maxWorkers uint64
maxBody int64
redirects int
seqmu sync.Mutex
seq uint64
began time.Time
}

const (
Expand All @@ -40,6 +42,8 @@ const (
DefaultConnections = 10000
// DefaultWorkers is the default initial number of workers used to carry an attack.
DefaultWorkers = 10
// DefaultMaxWorkers is the default maximum number of workers used to carry an attack.
DefaultMaxWorkers = math.MaxUint64
// DefaultMaxBody is the default max number of bytes to be read from response bodies.
// Defaults to no limit.
DefaultMaxBody = int64(-1)
Expand All @@ -58,10 +62,11 @@ var (
// by the optionally provided opts.
func NewAttacker(opts ...func(*Attacker)) *Attacker {
a := &Attacker{
stopch: make(chan struct{}),
workers: DefaultWorkers,
maxBody: DefaultMaxBody,
began: time.Now(),
stopch: make(chan struct{}),
workers: DefaultWorkers,
maxWorkers: DefaultMaxWorkers,
maxBody: DefaultMaxBody,
began: time.Now(),
}

a.dialer = &net.Dialer{
Expand Down Expand Up @@ -93,6 +98,12 @@ func Workers(n uint64) func(*Attacker) {
return func(a *Attacker) { a.workers = n }
}

// MaxWorkers returns a functional option which sets the maximum number of workers
// an Attacker can use to hit its targets.
func MaxWorkers(n uint64) func(*Attacker) {
return func(a *Attacker) { a.maxWorkers = n }
}

// Connections returns a functional option which sets the number of maximum idle
// open connections per target host.
func Connections(n int) func(*Attacker) {
Expand Down Expand Up @@ -223,37 +234,59 @@ func Client(c *http.Client) func(*Attacker) {
// runs until Stop is called. Results are sent to the returned channel as soon
// as they arrive and will have their Attack field set to the given name.
func (a *Attacker) Attack(tr Targeter, p Pacer, du time.Duration, name string) <-chan *Result {
var workers sync.WaitGroup
var wg sync.WaitGroup

workers := a.workers
if workers > a.maxWorkers {
workers = a.maxWorkers
}

results := make(chan *Result)
ticks := make(chan struct{})
for i := uint64(0); i < a.workers; i++ {
workers.Add(1)
go a.attack(tr, name, &workers, ticks, results)
for i := uint64(0); i < workers; i++ {
wg.Add(1)
go a.attack(tr, name, &wg, ticks, results)
}

go func() {
defer close(results)
defer workers.Wait()
defer wg.Wait()
defer close(ticks)

began, count := time.Now(), uint64(0)
for {
elapsed := time.Since(began)
if du > 0 && elapsed > du {
return
}

wait, stop := p.Pace(elapsed, count)
if stop {
return
}

time.Sleep(wait)

if workers < a.maxWorkers {
select {
case ticks <- struct{}{}:
count++
continue
case <-a.stopch:
return
default:
// all workers are blocked. start one more and try again
workers++
wg.Add(1)
go a.attack(tr, name, &wg, ticks, results)
}
}

select {
case ticks <- struct{}{}:
count++
case <-a.stopch:
return
default: // all workers are blocked. start one more and try again
workers.Add(1)
go a.attack(tr, name, &workers, ticks, results)
}
}
}()
Expand Down

0 comments on commit fb2d8b9

Please sign in to comment.