Merge pull request #48 from fsmiamoto/master

Add flag for max number of idle connections

attacker/attacker.go

@@ -10,13 +10,14 @@ import (
 )
 
 const (
-	DefaultRate       = 50
-	DefaultDuration   = 10 * time.Second
-	DefaultTimeout    = 30 * time.Second
-	DefaultMethod     = http.MethodGet
-	DefaultWorkers    = 10
-	DefaultMaxWorkers = math.MaxUint64
-	DefaultMaxBody    = int64(-1)
+	DefaultRate        = 50
+	DefaultDuration    = 10 * time.Second
+	DefaultTimeout     = 30 * time.Second
+	DefaultMethod      = http.MethodGet
+	DefaultWorkers     = 10
+	DefaultMaxWorkers  = math.MaxUint64
+	DefaultMaxBody     = int64(-1)
+	DefaultConnections = 10000
 )
 
 type Attacker interface {
@@ -26,16 +27,17 @@ type Attacker interface {
 
 // Options provides optional settings to attack.
 type Options struct {
-	Rate       int
-	Duration   time.Duration
-	Timeout    time.Duration
-	Method     string
-	Body       []byte
-	Header     http.Header
-	Workers    uint64
-	MaxWorkers uint64
-	MaxBody    int64
-	KeepAlive  bool
+	Rate        int
+	Duration    time.Duration
+	Timeout     time.Duration
+	Method      string
+	Body        []byte
+	MaxBody     int64
+	Header      http.Header
+	Workers     uint64
+	MaxWorkers  uint64
+	KeepAlive   bool
+	Connections int
 
 	Attacker Attacker
 }
@@ -66,12 +68,16 @@ func Attack(ctx context.Context, target string, resCh chan *Result, metricsCh ch
 	if opts.MaxBody == 0 {
 		opts.MaxBody = DefaultMaxBody
 	}
+	if opts.Connections == 0 {
+		opts.Connections = DefaultConnections
+	}
 	if opts.Attacker == nil {
 		opts.Attacker = vegeta.NewAttacker(
 			vegeta.Timeout(opts.Timeout),
 			vegeta.Workers(opts.Workers),
 			vegeta.MaxWorkers(opts.MaxWorkers),
 			vegeta.MaxBody(opts.MaxBody),
+			vegeta.Connections(opts.Connections),
 			vegeta.KeepAlive(opts.KeepAlive),
 		)
 	}

main.go

@@ -30,16 +30,17 @@ var (
 )
 
 type cli struct {
-	rate       int
-	duration   time.Duration
-	timeout    time.Duration
-	method     string
-	headers    []string
-	body       string
-	bodyFile   string
-	maxBody    int64
-	workers    uint64
-	maxWorkers uint64
+	rate        int
+	duration    time.Duration
+	timeout     time.Duration
+	method      string
+	headers     []string
+	body        string
+	bodyFile    string
+	maxBody     int64
+	workers     uint64
+	maxWorkers  uint64
+	connections int
 
 	debug     bool
 	version   bool
@@ -74,6 +75,7 @@ func parseFlags(stdout, stderr io.Writer) (*cli, error) {
 	flagSet.BoolVarP(&c.keepAlive, "keepalive", "k", true, "Use HTTP persistent connection.")
 	flagSet.Uint64VarP(&c.workers, "workers", "w", attacker.DefaultWorkers, "Amount of initial workers to spawn.")
 	flagSet.Uint64VarP(&c.maxWorkers, "max-workers", "W", attacker.DefaultMaxWorkers, "Amount of maximum workers to spawn.")
+	flagSet.IntVarP(&c.connections, "connections", "c", attacker.DefaultConnections, "Amount of maximum open idle connections per target host")
 	flagSet.Usage = c.usage
 	if err := flagSet.Parse(os.Args[1:]); err != nil {
 		if !errors.Is(err, flag.ErrHelp) {

main_test.go

@@ -46,17 +46,18 @@ func TestParseFlags(t *testing.T) {
 		{
 			name: "with default options",
 			want: &cli{
-				rate:       50,
-				duration:   time.Second * 10,
-				timeout:    time.Second * 30,
-				method:     "GET",
-				headers:    []string{},
-				maxBody:    -1,
-				keepAlive:  true,
-				workers:    10,
-				maxWorkers: math.MaxUint64,
-				stdout:     new(bytes.Buffer),
-				stderr:     new(bytes.Buffer),
+				rate:        50,
+				duration:    time.Second * 10,
+				timeout:     time.Second * 30,
+				method:      "GET",
+				headers:     []string{},
+				maxBody:     -1,
+				keepAlive:   true,
+				workers:     10,
+				maxWorkers:  math.MaxUint64,
+				connections: 10000,
+				stdout:      new(bytes.Buffer),
+				stderr:      new(bytes.Buffer),
 			},
 			wantErr: false,
 		},