Vegeta as a library

* Slicker API
* Privatises a bunch of stuff
* More tests
* More documentation

README.md

@@ -2,6 +2,7 @@
 
 Vegeta is a versatile HTTP load testing tool built out of need to drill
 HTTP services with a constant request rate.
+It can be used both as a command line utility and a library.
 
 ![Vegeta](https://dl.dropboxusercontent.com/u/83217940/vegeta.png)
 
@@ -12,7 +13,7 @@ command:
 $ go install github.com/tsenart/vegeta
 ```
 
-## Usage
+## Usage (CLI)
 ```shell
 $ vegeta -h
 Usage of vegeta:
@@ -59,6 +60,28 @@ HEAD http://goku:9090/path/to/success
 ...
 ```
 
+## Usage (Library)
+```go
+package main
+
+import (
+  vegeta "github.com/tsenart/vegeta/lib"
+  "time"
+  "os"
+)
+
+func main() {
+  targets := vegeta.NewTargets([]string{"GET http://localhost:9100/"})
+  rate := uint64(100) // per second
+  duration := 4 * time.Second
+  reporter := vegeta.NewTextReporter()
+
+  vegeta.Attack(targets, rate, duration, reporter)
+
+  reporter.Report(os.Stdout)
+}
+```
+
 #### Limitations
 There will be an upper bound of the supported `rate` which varies on the
 machine being used.
@@ -79,7 +102,6 @@ Just pass a new number as the argument to change it.
 * Add timeout options to the requests
 * Graphical reporters
 * Cluster mode (to overcome single machine limits)
-* More tests
 * HTTPS
 
 ## Licence

client.go

@@ -1,52 +0,0 @@
-package main
-
-import (
-	"errors"
-	"io/ioutil"
-	"net/http"
-	"time"
-)
-
-// Client is an http.Client with rate limiting
-// TODO: Add timeouts
-type Client struct{ http.Client }
-
-// Response represents the metrics we want out of an http.Response
-type Response struct {
-	code      uint64
-	timestamp time.Time
-	timing    time.Duration
-	bytesOut  uint64
-	bytesIn   uint64
-	err       error
-}
-
-// Drill loops over the passed reqs channel and executes each request.
-// It is throttled to the rate specified
-func (c *Client) Drill(rate uint64, reqs chan *http.Request, res chan *Response) {
-	throttle := time.Tick(time.Duration(1e9 / rate))
-	for req := range reqs {
-		<-throttle
-		go c.Do(req, res)
-	}
-}
-
-// Do executes the passed http.Request and puts a generated *Response into res.
-func (c *Client) Do(req *http.Request, res chan *Response) {
-	began := time.Now()
-	r, err := c.Client.Do(req)
-	resp := &Response{
-		timestamp: began,
-		timing:    time.Since(began),
-		bytesOut:  uint64(req.ContentLength),
-		err:       err,
-	}
-	if err == nil {
-		resp.bytesIn, resp.code = uint64(r.ContentLength), uint64(r.StatusCode)
-		if body, err := ioutil.ReadAll(r.Body); err != nil && resp.code < 200 || resp.code >= 300 {
-			resp.err = errors.New(string(body))
-		}
-	}
-
-	res <- resp
-}

lib/attack.go

@@ -0,0 +1,68 @@
+package vegeta
+
+import (
+	"errors"
+	"io/ioutil"
+	"net/http"
+	"time"
+)
+
+// Attack hits the passed Targets (http.Requests) at the rate specified for
+// duration time and then waits for all the requests to come back.
+// The results of the attack are put into the rep Reporter.
+func Attack(targets Targets, rate uint64, duration time.Duration, rep Reporter) {
+	hits := make(chan *http.Request, rate*uint64((duration).Seconds()))
+	defer close(hits)
+	responses := make(chan *result, cap(hits))
+	defer close(responses)
+	go drill(rate, hits, responses) // Attack!
+	for i := 0; i < cap(hits); i++ {
+		hits <- targets[i%len(targets)]
+	}
+	// Wait for all requests to finish
+	for i := 0; i < cap(responses); i++ {
+		rep.add(<-responses)
+	}
+}
+
+// result represents the metrics we want out of an http.Response
+type result struct {
+	code      uint64
+	timestamp time.Time
+	timing    time.Duration
+	bytesOut  uint64
+	bytesIn   uint64
+	err       error
+}
+
+// drill loops over the passed reqs channel and executes each request.
+// It is throttled to the rate specified.
+func drill(rate uint64, reqs chan *http.Request, res chan *result) {
+	throttle := time.Tick(time.Duration(1e9 / rate))
+	for req := range reqs {
+		<-throttle
+		go hit(req, res)
+	}
+}
+
+// hit executes the passed http.Request and puts a generated *result into res.
+// Both transport errors and unsucessfull requests (non {2xx,3xx}) are
+// considered errors which are set in the Response.
+func hit(req *http.Request, res chan *result) {
+	began := time.Now()
+	r, err := http.DefaultClient.Do(req)
+	result := &result{
+		timestamp: began,
+		timing:    time.Since(began),
+		bytesOut:  uint64(req.ContentLength),
+		err:       err,
+	}
+	if err == nil {
+		result.bytesIn, result.code = uint64(r.ContentLength), uint64(r.StatusCode)
+		if body, err := ioutil.ReadAll(r.Body); err != nil && result.code < 200 || result.code >= 300 {
+			result.err = errors.New(string(body))
+		}
+	}
+
+	res <- result
+}

lib/attack_test.go

@@ -1,7 +1,6 @@
-package main
+package vegeta
 
 import (
-	"bytes"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -17,13 +16,10 @@ func TestAttackRate(t *testing.T) {
 			atomic.AddUint64(&hitCount, 1)
 		}),
 	)
-	targets, err := NewTargets(bytes.NewBufferString("GET " + server.URL + "\n"))
-	if err != nil {
-		t.Fatal(err)
-	}
+	request, _ := http.NewRequest("GET", server.URL, nil)
 	rate := uint64(5000)
 	rep := NewTextReporter()
-	attack(targets, "random", rate, 1*time.Second, rep)
+	Attack(Targets{request}, rate, 1*time.Second, rep)
 	if hits := atomic.LoadUint64(&hitCount); hits != rate {
 		rep.Report(os.Stdout)
 		t.Fatalf("Wrong number of hits: want %d, got %d\n", rate, hits)

lib/reporters.go

@@ -1,4 +1,4 @@
-package main
+package vegeta
 
 import (
 	"fmt"
@@ -9,23 +9,20 @@ import (
 
 // Reporter represents any reporter of the results of the test
 type Reporter interface {
-	Add(res *Response)
 	Report(io.Writer) error
+	add(res *result)
 }
 
+// TextReporter prints the test results as text
+// Metrics incude avg time per request, success ratio,
+// total number of request, avg bytes in and avg bytes out
 type TextReporter struct {
-	responses []*Response
+	responses []*result
 }
 
 // NewTextReporter initializes a TextReporter with n responses
 func NewTextReporter() *TextReporter {
-	return &TextReporter{responses: make([]*Response, 0)}
-}
-
-// Add adds a response to be used in the report
-// Order of arrival is not relevant for this reporter
-func (r *TextReporter) Add(res *Response) {
-	r.responses = append(r.responses, res)
+	return &TextReporter{responses: make([]*result, 0)}
 }
 
 // Report computes and writes the report to out.
@@ -77,3 +74,9 @@ func (r *TextReporter) Report(out io.Writer) error {
 
 	return w.Flush()
 }
+
+// add adds a response to be used in the report
+// Order of arrival is not relevant for this reporter
+func (r *TextReporter) add(res *result) {
+	r.responses = append(r.responses, res)
+}

lib/targets.go

@@ -1,4 +1,4 @@
-package main
+package vegeta
 
 import (
 	"bufio"
@@ -10,26 +10,41 @@ import (
 	"strings"
 )
 
+// Targets represents the http.Requests which will be issued during the test
 type Targets []*http.Request
 
+// NewTargetsFromFile reads and parses targets from a text file
 func NewTargetsFromFile(filename string) (Targets, error) {
 	file, err := os.Open(filename)
 	if err != nil {
 		return Targets{}, err
 	}
 	defer file.Close()
-	return NewTargets(file)
+	return readTargets(file)
 }
 
-func NewTargets(source io.Reader) (Targets, error) {
-	targets := make([]*http.Request, 0)
+// readTargets reads targets out of a line separated source skipping empty lines
+func readTargets(source io.Reader) (Targets, error) {
 	scanner := bufio.NewScanner(source)
-
+	lines := make([]string, 0)
 	for scanner.Scan() {
 		line := scanner.Text()
 		if line = strings.TrimSpace(line); line == "" { // Empty line
 			continue
 		}
+		lines = append(lines, line)
+	}
+	if err := scanner.Err(); err != nil {
+		return Targets{}, err
+	}
+
+	return NewTargets(lines)
+}
+
+// NewTargets instantiates Targets from a slice of strings
+func NewTargets(lines []string) (Targets, error) {
+	targets := make([]*http.Request, 0)
+	for _, line := range lines {
 		parts := strings.SplitN(line, " ", 2)
 		if len(parts) != 2 {
 			return targets, fmt.Errorf("Invalid request format: `%s`", line)
@@ -41,12 +56,10 @@ func NewTargets(source io.Reader) (Targets, error) {
 		}
 		targets = append(targets, req)
 	}
-	if err := scanner.Err(); err != nil {
-		return targets, err
-	}
 	return targets, nil
 }
 
+// Shuffle randomly alters the order of Targets with the provided seed
 func (t Targets) Shuffle(seed int64) {
 	rand.Seed(seed)
 	for i, rnd := range rand.Perm(len(t)) {

lib/targets_test.go

@@ -1,4 +1,4 @@
-package main
+package vegeta
 
 import (
 	"bytes"
@@ -6,9 +6,24 @@ import (
 	"testing"
 )
 
+func TestReadTargets(t *testing.T) {
+	lines := bytes.NewBufferString("GET http://lolcathost:9999/\n\nHEAD http://lolcathost:9999/\n")
+	targets, err := readTargets(lines)
+	if err != nil {
+		t.Fatalf("Couldn't parse valid source: %s", err)
+	}
+	for i, method := range []string{"GET", "HEAD"} {
+		if targets[i].Method != method ||
+			targets[i].URL.String() != "http://lolcathost:9999/" {
+			t.Fatalf("Request was parsed incorrectly. Got: %s %s",
+				targets[i].Method, targets[i].URL.String())
+		}
+	}
+}
+
 func TestNewTargets(t *testing.T) {
-	text := "GET http://lolcathost:9999/\n\nHEAD http://lolcathost:9999/\n"
-	targets, err := NewTargets(bytes.NewBufferString(text))
+	lines := []string{"GET http://lolcathost:9999/", "HEAD http://lolcathost:9999/"}
+	targets, err := NewTargets(lines)
 	if err != nil {
 		t.Fatalf("Couldn't parse valid source: %s", err)
 	}

main.go

@@ -2,9 +2,9 @@ package main
 
 import (
 	"flag"
+	vegeta "github.com/tsenart/vegeta/lib"
 	"io"
 	"log"
-	"net/http"
 	"os"
 	"runtime"
 	"time"
@@ -34,7 +34,7 @@ func main() {
 		log.Fatal("rate can't be zero")
 	}
 
-	targets, err := NewTargetsFromFile(*targetsf)
+	targets, err := vegeta.NewTargetsFromFile(*targetsf)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -52,13 +52,13 @@ func main() {
 		log.Fatal("Duration provided is invalid")
 	}
 
-	var rep Reporter
+	var rep vegeta.Reporter
 	switch *reporter {
 	case "text":
-		rep = NewTextReporter()
+		rep = vegeta.NewTextReporter()
 	default:
 		log.Println("reporter provided is not supported. using text")
-		rep = NewTextReporter()
+		rep = vegeta.NewTextReporter()
 	}
 
 	var out io.Writer
@@ -75,7 +75,7 @@ func main() {
 	}
 
 	log.Printf("Vegeta is attacking %d targets in %s order for %s...\n", len(targets), *ordering, *duration)
-	attack(targets, *ordering, *rate, *duration, rep)
+	vegeta.Attack(targets, *rate, *duration, rep)
 	log.Println("Done!")
 
 	log.Printf("Writing report to '%s'...", *output)
@@ -83,19 +83,3 @@ func main() {
 		log.Println("Failed to report!")
 	}
 }
-
-func attack(targets Targets, ordering string, rate uint64, duration time.Duration, rep Reporter) {
-	hits := make(chan *http.Request, rate*uint64((duration).Seconds()))
-	defer close(hits)
-	responses := make(chan *Response, cap(hits))
-	defer close(responses)
-	client := Client{}
-	go client.Drill(rate, hits, responses) // Attack!
-	for i := 0; i < cap(hits); i++ {
-		hits <- targets[i%len(targets)]
-	}
-	// Wait for all requests to finish
-	for i := 0; i < cap(responses); i++ {
-		rep.Add(<-responses)
-	}
-}