Vegeta as a library

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

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

@@ -0,0 +1,27 @@
+package vegeta
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestAttackRate(t *testing.T) {
+	hitCount := uint64(0)
+	server := httptest.NewServer(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			atomic.AddUint64(&hitCount, 1)
+		}),
+	)
+	request, _ := http.NewRequest("GET", server.URL, nil)
+	rate := uint64(5000)
+	rep := NewTextReporter()
+	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

@@ -0,0 +1,82 @@
+package vegeta
+
+import (
+	"fmt"
+	"io"
+	"text/tabwriter"
+	"time"
+)
+
+// Reporter represents any reporter of the results of the test
+type Reporter interface {
+	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 []*result
+}
+
+// NewTextReporter initializes a TextReporter with n responses
+func NewTextReporter() *TextReporter {
+	return &TextReporter{responses: make([]*result, 0)}
+}
+
+// Report computes and writes the report to out.
+// It returns an error in case of failure.
+func (r *TextReporter) Report(out io.Writer) error {
+	totalRequests := len(r.responses)
+	totalTime := time.Duration(0)
+	totalBytesOut := uint64(0)
+	totalBytesIn := uint64(0)
+	totalSuccess := uint64(0)
+	histogram := map[uint64]uint64{}
+	errors := map[string]struct{}{}
+
+	for _, res := range r.responses {
+		histogram[res.code]++
+		totalTime += res.timing
+		totalBytesOut += res.bytesOut
+		totalBytesIn += res.bytesIn
+		if res.code >= 200 && res.code < 300 {
+			totalSuccess++
+		}
+		if res.err != nil {
+			errors[res.err.Error()] = struct{}{}
+		}
+	}
+
+	avgTime := time.Duration(float64(totalTime) / float64(totalRequests))
+	avgBytesOut := float64(totalBytesOut) / float64(totalRequests)
+	avgBytesIn := float64(totalBytesIn) / float64(totalRequests)
+	avgSuccess := float64(totalSuccess) / float64(totalRequests)
+
+	w := tabwriter.NewWriter(out, 0, 8, 2, '\t', tabwriter.StripEscape)
+	fmt.Fprintf(w, "Time(avg)\tRequests\tSuccess\tBytes(rx/tx)\n")
+	fmt.Fprintf(w, "%s\t%d\t%.2f%%\t%.2f/%.2f\n", avgTime, totalRequests, avgSuccess*100, avgBytesOut, avgBytesIn)
+
+	fmt.Fprintf(w, "\nCount:\t")
+	for _, count := range histogram {
+		fmt.Fprintf(w, "%d\t", count)
+	}
+	fmt.Fprintf(w, "\nStatus:\t")
+	for code, _ := range histogram {
+		fmt.Fprintf(w, "%d\t", code)
+	}
+
+	fmt.Fprintln(w, "\n\nError Set:")
+	for err, _ := range errors {
+		fmt.Fprintln(w, err)
+	}
+
+	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

@@ -0,0 +1,70 @@
+package vegeta
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"math/rand"
+	"net/http"
+	"os"
+	"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 readTargets(file)
+}
+
+// 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)
+		}
+		// Build request
+		req, err := http.NewRequest(parts[0], parts[1], nil)
+		if err != nil {
+			return targets, fmt.Errorf("Failed to build request: %s", err)
+		}
+		targets = append(targets, req)
+	}
+	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)) {
+		tmp := t[i]
+		t[i] = t[rnd]
+		t[rnd] = tmp
+	}
+}

lib/targets_test.go

@@ -0,0 +1,54 @@
+package vegeta
+
+import (
+	"bytes"
+	"net/http"
+	"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) {
+	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)
+	}
+	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 TestShuffle(t *testing.T) {
+	targets := make(Targets, 50)
+	for i := 0; i < 50; i++ {
+		targets[i], _ = http.NewRequest("GET", "http://lolcathost:9999/", nil)
+	}
+	targetsCopy := make(Targets, 50)
+	copy(targetsCopy, targets)
+
+	targets.Shuffle(0)
+	for i, target := range targets {
+		if targetsCopy[i] != target {
+			return
+		}
+	}
+	t.Fatal("Targets were not shuffled correctly")
+}