Integrate with fasthttp

attack.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"errors"
@@ -17,6 +18,9 @@ import (
 
 	"github.com/tsenart/vegeta/v12/internal/resolver"
 	vegeta "github.com/tsenart/vegeta/v12/lib"
+
+	"github.com/valyala/fasthttp"
+	"golang.org/x/net/http2"
 )
 
 func attackCmd() command {
@@ -38,17 +42,17 @@ func attackCmd() command {
 	fs.StringVar(&opts.certf, "cert", "", "TLS client PEM encoded certificate file")
 	fs.StringVar(&opts.keyf, "key", "", "TLS client PEM encoded private key file")
 	fs.Var(&opts.rootCerts, "root-certs", "TLS root certificate files (comma separated list)")
-	fs.BoolVar(&opts.http2, "http2", true, "Send HTTP/2 requests when supported by the server")
+	fs.BoolVar(&opts.http2, "http2", false, "Send HTTP/2 requests when supported by the server")
 	fs.BoolVar(&opts.h2c, "h2c", false, "Send HTTP/2 requests without TLS encryption")
 	fs.BoolVar(&opts.insecure, "insecure", false, "Ignore invalid server TLS certificates")
 	fs.BoolVar(&opts.lazy, "lazy", false, "Read targets lazily")
 	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.maxConnections, "max-connections", vegeta.DefaultMaxConnections, "Max 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.BoolVar(&opts.skipBody, "skip-body", false, "Skip reading response bodies")
 	fs.Var(&maxBodyFlag{&opts.maxBody}, "max-body", "Maximum number of bytes to capture from response bodies. [-1 = no limit]")
 	fs.Var(&rateFlag{&opts.rate}, "rate", "Number of requests per time unit [0 = infinity]")
 	fs.Var(&opts.headers, "header", "Request header")
@@ -84,12 +88,12 @@ type attackOpts struct {
 	insecure       bool
 	lazy           bool
 	chunked        bool
+	skipBody       bool
 	duration       time.Duration
 	timeout        time.Duration
 	rate           vegeta.Rate
 	workers        uint64
 	maxWorkers     uint64
-	connections    int
 	maxConnections int
 	redirects      int
 	maxBody        int64
@@ -137,10 +141,9 @@ func attack(opts *attackOpts) (err error) {
 	}
 
 	var (
-		tr       vegeta.Targeter
-		src      = files[opts.targetsf]
-		hdr      = opts.headers.Header
-		proxyHdr = opts.proxyHeaders.Header
+		tr  vegeta.Targeter
+		src = files[opts.targetsf]
+		hdr = opts.headers.Header
 	)
 
 	switch opts.format {
@@ -167,36 +170,32 @@ func attack(opts *attackOpts) (err error) {
 	}
 	defer out.Close()
 
-	tlsc, err := tlsConfig(opts.insecure, opts.certf, opts.keyf, opts.rootCerts)
+	hitter, err := newHitter(opts)
 	if err != nil {
 		return err
 	}
 
-	atk := vegeta.NewAttacker(
-		vegeta.Redirects(opts.redirects),
-		vegeta.Timeout(opts.timeout),
-		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.MaxConnections(opts.maxConnections),
-		vegeta.HTTP2(opts.http2),
-		vegeta.H2C(opts.h2c),
-		vegeta.MaxBody(opts.maxBody),
-		vegeta.UnixSocket(opts.unixSocket),
-		vegeta.ProxyHeader(proxyHdr),
-		vegeta.ChunkedBody(opts.chunked),
-	)
+	atk := vegeta.Attack{
+		Hitter:     hitter,
+		Targeter:   tr,
+		Pacer:      opts.rate,
+		Name:       opts.name,
+		Duration:   opts.duration,
+		Workers:    opts.workers,
+		MaxWorkers: opts.maxWorkers,
+	}
+
+	res := make(chan *vegeta.Result)
+	go atk.Run(res)
 
-	res := atk.Attack(tr, opts.rate, opts.duration, opts.name)
 	enc := vegeta.NewEncoder(out)
 	sig := make(chan os.Signal, 1)
 	signal.Notify(sig, os.Interrupt)
 
 	for {
 		select {
+		case <-atk.Done():
+			return nil
 		case <-sig:
 			atk.Stop()
 			return nil
@@ -211,6 +210,110 @@ func attack(opts *attackOpts) (err error) {
 	}
 }
 
+func newHitter(opts *attackOpts) (vegeta.Hitter, error) {
+	tlsc, err := tlsConfig(opts.insecure, opts.certf, opts.keyf, opts.rootCerts)
+	if err != nil {
+		return nil, err
+	}
+
+	needsNetHTTPHitter := opts.http2 || opts.h2c ||
+		os.Getenv("HTTP_PROXY") != "" ||
+		os.Getenv("HTTPS_PROXY") != ""
+
+	if needsNetHTTPHitter {
+		dialer := net.Dialer{
+			LocalAddr: &net.TCPAddr{
+				IP:   opts.laddr.IP,
+				Zone: opts.laddr.Zone,
+			},
+			KeepAlive: 30 * time.Second,
+		}
+
+		tr := &http.Transport{
+			Proxy:              http.ProxyFromEnvironment,
+			DialContext:        dialer.DialContext,
+			TLSClientConfig:    tlsc,
+			MaxConnsPerHost:    opts.maxConnections,
+			DisableKeepAlives:  !opts.keepalive,
+			ProxyConnectHeader: opts.proxyHeaders.Header,
+		}
+
+		client := &http.Client{
+			Timeout:   opts.timeout,
+			Transport: tr,
+			CheckRedirect: func(_ *http.Request, via []*http.Request) error {
+				switch {
+				case opts.redirects == vegeta.NoFollow:
+					return http.ErrUseLastResponse
+				case opts.redirects < len(via):
+					return fmt.Errorf("stopped after %d redirects", opts.redirects)
+				default:
+					return nil
+				}
+			},
+		}
+
+		if opts.unixSocket != "" {
+			tr.DialContext = func(_ context.Context, _, _ string) (net.Conn, error) {
+				return net.Dial("unix", opts.unixSocket)
+			}
+		}
+
+		switch {
+		case opts.http2:
+			if err = http2.ConfigureTransport(tr); err != nil {
+				return nil, err
+			}
+		case opts.h2c:
+			client.Transport = &http2.Transport{
+				AllowHTTP: true,
+				DialTLS: func(network, addr string, cfg *tls.Config) (net.Conn, error) {
+					return tr.DialContext(context.Background(), network, addr)
+				},
+			}
+		}
+
+		return &vegeta.NetHTTPHitter{
+			Chunked:  opts.chunked,
+			MaxBody:  opts.maxBody,
+			SkipBody: opts.skipBody,
+			Client:   client,
+		}, nil
+	}
+
+	dialer := fasthttp.TCPDialer{
+		LocalAddr: &net.TCPAddr{
+			IP:   opts.laddr.IP,
+			Zone: opts.laddr.Zone,
+		},
+	}
+
+	cli := &fasthttp.Client{
+		Name:                          "vegeta " + Version,
+		NoDefaultUserAgentHeader:      true,
+		ReadTimeout:                   opts.timeout,
+		TLSConfig:                     tlsc,
+		MaxConnsPerHost:               opts.maxConnections,
+		MaxResponseBodySize:           int(opts.maxBody),
+		DisableHeaderNamesNormalizing: true,
+		Dial:                          dialer.Dial,
+	}
+
+	if opts.unixSocket != "" {
+		cli.Dial = func(_ string) (net.Conn, error) {
+			return net.Dial("unix", opts.unixSocket)
+		}
+	}
+
+	return &vegeta.FastHTTPHitter{
+		Client:       cli,
+		MaxRedirects: opts.redirects,
+		Chunked:      opts.chunked,
+		KeepAlive:    opts.keepalive,
+		SkipBody:     opts.skipBody,
+	}, nil
+}
+
 // tlsConfig builds a *tls.Config from the given options.
 func tlsConfig(insecure bool, certf, keyf string, rootCerts []string) (*tls.Config, error) {
 	var err error

go.mod

@@ -19,10 +19,16 @@ require (
 	github.com/gonum/stat v0.0.0-20181125101827-41a0da705a5b // indirect
 	github.com/google/go-cmp v0.2.0
 	github.com/influxdata/tdigest v0.0.0-20180711151920-a7d76c6f093a
+	github.com/klauspost/compress v1.10.3 // indirect
 	github.com/mailru/easyjson v0.7.0
 	github.com/miekg/dns v1.1.17
 	github.com/streadway/quantile v0.0.0-20150917103942-b0c588724d25
 	github.com/tsenart/go-tsz v0.0.0-20180814232043-cdeb9e1e981e
-	golang.org/x/net v0.0.0-20190827160401-ba9fcec4b297
+	github.com/valyala/fasthttp v1.9.1-0.20200313102552-38aa88ab52c4
+	golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550 // indirect
+	golang.org/x/net v0.0.0-20200226121028-0de0cce0169b
+	golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e // indirect
 	pgregory.net/rapid v0.3.3
 )
+
+replace github.com/valyala/fasthttp => github.com/tsenart/fasthttp v1.9.1-0.20200315143151-5b99e2de1407

go.sum

@@ -30,23 +30,40 @@ github.com/google/go-cmp v0.2.0 h1:+dTQ8DZQJz0Mb/HjFlkptS1FeQ4cWSnN941F8aEG4SQ=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
 github.com/influxdata/tdigest v0.0.0-20180711151920-a7d76c6f093a h1:vMqgISSVkIqWxCIZs8m1L4096temR7IbYyNdMiBxSPA=
 github.com/influxdata/tdigest v0.0.0-20180711151920-a7d76c6f093a/go.mod h1:9GkyshztGufsdPQWjH+ifgnIr3xNUL5syI70g2dzU1o=
+github.com/klauspost/compress v1.8.2 h1:Bx0qjetmNjdFXASH02NSAREKpiaDwkO1DRZ3dV2KCcs=
+github.com/klauspost/compress v1.8.2/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A=
+github.com/klauspost/compress v1.10.3 h1:OP96hzwJVBIHYU52pVTI6CczrxPvrGfgqF9N5eTO0Q8=
+github.com/klauspost/compress v1.10.3/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
+github.com/klauspost/cpuid v1.2.1 h1:vJi+O/nMdFt0vqm8NZBI6wzALWdA2X+egi0ogNyrC/w=
+github.com/klauspost/cpuid v1.2.1/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
 github.com/mailru/easyjson v0.7.0 h1:aizVhC/NAAcKWb+5QsU1iNOZb4Yws5UO2I+aIprQITM=
 github.com/mailru/easyjson v0.7.0/go.mod h1:KAzv3t3aY1NaHWoQz1+4F1ccyAH66Jk7yos7ldAVICs=
 github.com/miekg/dns v1.1.17 h1:BhJxdA7bH51vKFZSY8Sn9pR7++LREvg0eYFzHA452ew=
 github.com/miekg/dns v1.1.17/go.mod h1:WgzbA6oji13JREwiNsRDNfl7jYdPnmz+VEuLrA+/48M=
 github.com/streadway/quantile v0.0.0-20150917103942-b0c588724d25 h1:7z3LSn867ex6VSaahyKadf4WtSsJIgne6A1WLOAGM8A=
 github.com/streadway/quantile v0.0.0-20150917103942-b0c588724d25/go.mod h1:lbP8tGiBjZ5YWIc2fzuRpTaz0b/53vT6PEs3QuAWzuU=
+github.com/tsenart/fasthttp v1.9.1-0.20200315143151-5b99e2de1407 h1:IHBlcvLiQCL2rzfvfzRqxtJ8A2MFvosGvH6jjvc8/1Q=
+github.com/tsenart/fasthttp v1.9.1-0.20200315143151-5b99e2de1407/go.mod h1:FstJa9V+Pj9vQ7OJie2qMHdwemEDaDiSdBnvPM1Su9w=
 github.com/tsenart/go-tsz v0.0.0-20180814232043-cdeb9e1e981e h1:bB5SXzQmSUsJCmjPDN9fKYx3SSDER5diSjlN6TefTCc=
 github.com/tsenart/go-tsz v0.0.0-20180814232043-cdeb9e1e981e/go.mod h1:SWZznP1z5Ki7hDT2ioqiFKEse8K9tU2OUvaRI0NeGQo=
+github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
+github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
+github.com/valyala/tcplisten v0.0.0-20161114210144-ceec8f93295a/go.mod h1:v3UYOV9WzVtRmSR+PDvWpU/qWl4Wa5LApYYX4ZtKbio=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20190829043050-9756ffdc2472 h1:Gv7RPwsi3eZ2Fgewe3CBsuOebPwO27PoXzRpJPsvSSM=
 golang.org/x/crypto v0.0.0-20190829043050-9756ffdc2472/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
+golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550 h1:ObdrDkeb4kJdCP557AjRjq69pTHfNouLtWZG7j9rPN8=
+golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20190827160401-ba9fcec4b297 h1:k7pJ2yAPLPgbskkFdhRCsA77k2fySZ1zf2zCjvQCiIM=
 golang.org/x/net v0.0.0-20190827160401-ba9fcec4b297/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
+golang.org/x/net v0.0.0-20200226121028-0de0cce0169b h1:0mm1VjtFUOIlE1SbDlwjYaDxZVDP2S5ou6y0gSgXHu8=
+golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58 h1:8gQV6CLnAEikrhgkHFbMAEhagSSnXWGV915qUMm9mrU=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e h1:vcxGaoTs7kV8m5Np9uUNQin4BrLOthgV7252N8V+FwY=
+golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20190904154756-749cb33beabd h1:DBH9mDw0zluJT/R+nGuV3jWFWLFaHyYZWD4tOT+cjn0=
@@ -55,6 +72,7 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.2 h1:tW2bmiBqwgJj/UpqtC8EpXEZVYOwU0yG4iWbprSVAcs=
 golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+golang.org/x/tools v0.0.0-20190907020128-2ca718005c18 h1:xFbv3LvlvQAmbNJFCBKRv1Ccvnh9FVsW0FX2kTWWowE=
 golang.org/x/tools v0.0.0-20190907020128-2ca718005c18/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 pgregory.net/rapid v0.3.3 h1:jCjBsY4ln4Atz78QoBWxUEvAHaFyNDQg9+WU62aCn1U=