Merge pull request #91 from jpfuentes2/jf_keepalive_option

Add -keepalive flag to toggle persistent conns

attack.go

@@ -38,6 +38,7 @@ func attackCmd() command {
 	fs.IntVar(&opts.redirects, "redirects", vegeta.DefaultRedirects, "Number of redirects to follow")
 	fs.Var(&opts.headers, "header", "Request header")
 	fs.Var(&opts.laddr, "laddr", "Local IP address")
+	fs.BoolVar(&opts.keepalive, "keepalive", true, "Use persistent connections")
 
 	return command{fs, func(args []string) error {
 		fs.Parse(args)
@@ -65,6 +66,7 @@ type attackOpts struct {
 	redirects int
 	headers   headers
 	laddr     localAddr
+	keepalive bool
 }
 
 // attack validates the attack arguments, sets up the
@@ -134,6 +136,7 @@ func attack(opts *attackOpts) (err error) {
 		vegeta.LocalAddr(*opts.laddr.IPAddr),
 		vegeta.TLSConfig(&tlsc),
 		vegeta.Workers(opts.workers),
+		vegeta.KeepAlive(opts.keepalive),
 	)
 
 	res := atk.Attack(tr, opts.rate, opts.duration)

lib/attack.go

@@ -99,6 +99,20 @@ func LocalAddr(addr net.IPAddr) func(*Attacker) {
 	}
 }
 
+// KeepAlive returns a functional option which toggles KeepAlive
+// connections on the dialer and transport.
+func KeepAlive(keepalive bool) func(*Attacker) {
+	return func(a *Attacker) {
+		tr := a.client.Transport.(*http.Transport)
+		tr.DisableKeepAlives = !keepalive
+		if !keepalive {
+			a.dialer.KeepAlive = 0
+			tr.Dial = a.dialer.Dial
+		}
+		a.client.Transport = tr
+	}
+}
+
 // TLSConfig returns a functional option which sets the *tls.Config for a
 // Attacker to use with its requests.
 func TLSConfig(c *tls.Config) func(*Attacker) {

lib/attack_test.go

@@ -126,3 +126,18 @@ func TestLocalAddr(t *testing.T) {
 		}
 	}
 }
+
+func TestKeepAlive(t *testing.T) {
+	t.Parallel()
+
+	atk := NewAttacker(KeepAlive(false))
+
+	if atk.dialer.KeepAlive != 0 {
+		t.Fatalf("Dialer KeepAlive is not disabled. Want 0. Got %d", atk.dialer.KeepAlive)
+	}
+
+	disableKeepAlive := atk.client.Transport.(*http.Transport).DisableKeepAlives
+	if disableKeepAlive == false {
+		t.Fatalf("Transport DisableKeepAlives is not enabled. Want true. Got %t", disableKeepAlive)
+	}
+}