Merge pull request #14 from tsenart/improve-library

Improve Vegeta as a library

README.md

@@ -97,17 +97,23 @@ import (
   vegeta "github.com/tsenart/vegeta/lib"
   "time"
   "os"
+  "fmt"
 )
 
 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)
+  results := vegeta.Attack(targets, rate, duration)
 
-  reporter.Report(os.Stdout)
+  totalTime := time.Duration(0)
+  for _, result := range results {
+    totalTime += result.Timing
+  }
+  meanTime := time.Duration(float64(totalTime) / float64(len(results)))
+
+  fmt.Printf("Average timing: %s", meanTime)
 }
 ```
 

lib/attack.go

@@ -4,40 +4,55 @@ import (
 	"errors"
 	"io/ioutil"
 	"net/http"
+	"sort"
 	"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.
+// The results of the attack are put into a slice which is returned.
-func Attack(targets Targets, rate uint64, duration time.Duration, rep Reporter) {
+func Attack(targets Targets, rate uint64, duration time.Duration) []Result {
-	hits := make(chan *http.Request, rate*uint64((duration).Seconds()))
+	total := rate * uint64(duration.Seconds())
-	defer close(hits)
+	hits := make(chan *http.Request, total)
-	responses := make(chan *result, cap(hits))
+	res := make(chan Result, total)
-	defer close(responses)
+	results := make(Results, total)
-	go drill(rate, hits, responses) // Attack!
+	// Scatter
+	go drill(rate, hits, res)
 	for i := 0; i < cap(hits); i++ {
 		hits <- targets[i%len(targets)]
 	}
-	// Wait for all requests to finish
+	close(hits)
-	for i := 0; i < cap(responses); i++ {
+	// Gather
-		rep.add(<-responses)
+	for i := 0; i < cap(res); i++ {
+		results[i] = <-res
 	}
+	close(res)
+
+	sort.Sort(results)
+
+	return results
 }
 
-// result represents the metrics we want out of an http.Response
+// Result represents the metrics we want out of an http.Response
-type result struct {
+type Result struct {
-	code      uint64
+	Code      uint16
-	timestamp time.Time
+	Timestamp time.Time
-	timing    time.Duration
+	Timing    time.Duration
-	bytesOut  uint64
+	BytesOut  uint64
-	bytesIn   uint64
+	BytesIn   uint64
-	err       error
+	Error     error
 }
 
+// Results is a slice of Result defined only to be sortable with sort.Interface
+type Results []Result
+
+func (r Results) Len() int           { return len(r) }
+func (r Results) Less(i, j int) bool { return r[i].Timestamp.Before(r[j].Timestamp) }
+func (r Results) Swap(i, j int)      { r[i], r[j] = r[j], r[i] }
+
 // 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) {
+func drill(rate uint64, reqs chan *http.Request, res chan Result) {
 	throttle := time.Tick(time.Duration(1e9 / rate))
 	for req := range reqs {
 		<-throttle
@@ -45,24 +60,27 @@ func drill(rate uint64, reqs chan *http.Request, res chan *result) {
 	}
 }
 
-// hit executes the passed http.Request and puts a generated *result into res.
+// hit executes the passed http.Request and puts the result into results.
 // Both transport errors and unsucessfull requests (non {2xx,3xx}) are
-// considered errors which are set in the Response.
+// considered errors.
-func hit(req *http.Request, res chan *result) {
+func hit(req *http.Request, res chan Result) {
 	began := time.Now()
 	r, err := http.DefaultClient.Do(req)
-	result := &result{
+	result := Result{
-		timestamp: began,
+		Timestamp: began,
-		timing:    time.Since(began),
+		Timing:    time.Since(began),
-		bytesOut:  uint64(req.ContentLength),
+		BytesOut:  uint64(req.ContentLength),
-		err:       err,
+		Error:     err,
 	}
 	if err == nil {
-		result.bytesIn, result.code = uint64(r.ContentLength), uint64(r.StatusCode)
+		result.Code = uint16(r.StatusCode)
-		if body, err := ioutil.ReadAll(r.Body); err != nil && (result.code < 200 || result.code >= 300) {
+		if body, err := ioutil.ReadAll(r.Body); err != nil {
-			result.err = errors.New(string(body))
+			if result.Code < 200 || result.Code >= 300 {
+				result.Error = errors.New(string(body))
+			}
+		} else {
+			result.BytesIn = uint64(len(body))
 		}
 	}
-
 	res <- result
 }

lib/attack_test.go

@@ -3,7 +3,6 @@ package vegeta
 import (
 	"net/http"
 	"net/http/httptest"
-	"os"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -18,10 +17,8 @@ func TestAttackRate(t *testing.T) {
 	)
 	request, _ := http.NewRequest("GET", server.URL, nil)
 	rate := uint64(5000)
-	rep := NewTextReporter()
+	Attack(Targets{request}, rate, 1*time.Second)
-	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/reporter.go

@@ -1,11 +0,0 @@
-package vegeta
-
-import (
-	"io"
-)
-
-// Reporter represents any reporter of the results of the test
-type Reporter interface {
-	Report(io.Writer) error
-	add(res *result)
-}

lib/reporters.go

@@ -0,0 +1,102 @@
+package vegeta
+
+import (
+	"code.google.com/p/plotinum/plot"
+	"code.google.com/p/plotinum/plotter"
+	"code.google.com/p/plotinum/plotutil"
+	"code.google.com/p/plotinum/vg"
+	"code.google.com/p/plotinum/vg/vgsvg"
+	"fmt"
+	"io"
+	"text/tabwriter"
+	"time"
+)
+
+// Reporter represents any function which takes a slice of Results and
+// generates a report, writing it to an io.Writer and returning an error
+// in case of failure
+type Reporter func([]Result, io.Writer) error
+
+// ReportText computes and prints some metrics out of results
+// as formatted text. Metrics include avg time per request, success ratio,
+// total number of request, avg bytes in and avg bytes out.
+func ReportText(results []Result, out io.Writer) error {
+	totalRequests := float64(len(results))
+	totalTime := time.Duration(0)
+	totalBytesOut := uint64(0)
+	totalBytesIn := uint64(0)
+	totalSuccess := uint64(0)
+	histogram := map[uint16]uint64{}
+	errors := map[string]struct{}{}
+
+	for _, res := range results {
+		histogram[res.Code]++
+		totalTime += res.Timing
+		totalBytesOut += res.BytesOut
+		totalBytesIn += res.BytesIn
+		if res.Code >= 200 && res.Code < 300 {
+			totalSuccess++
+		}
+		if res.Error != nil {
+			errors[res.Error.Error()] = struct{}{}
+		}
+	}
+
+	avgTime := time.Duration(float64(totalTime) / totalRequests)
+	avgBytesOut := float64(totalBytesOut) / totalRequests
+	avgBytesIn := float64(totalBytesIn) / totalRequests
+	avgSuccess := float64(totalSuccess) / 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, int(totalRequests), avgSuccess*100, avgBytesIn, avgBytesOut)
+
+	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()
+}
+
+// ReportTimingsPlot builds up a plot of the response times of the requests
+// in SVG format and writes it to out
+func ReportTimingsPlot(results []Result, out io.Writer) error {
+	p, err := plot.New()
+	if err != nil {
+		return err
+	}
+	pts := make(plotter.XYs, len(results))
+	for i := 0; i < len(pts); i++ {
+		pts[i].X = results[i].Timestamp.Sub(results[0].Timestamp).Seconds()
+		pts[i].Y = results[i].Timing.Seconds() * 1000
+	}
+
+	line, err := plotter.NewLine(pts)
+	if err != nil {
+		return err
+	}
+	line.Color = plotutil.Color(1)
+
+	p.Add(line)
+	p.X.Padding = vg.Length(3.0)
+	p.X.Label.Text = "Time elapsed"
+	p.Y.Padding = vg.Length(3.0)
+	p.Y.Label.Text = "Latency (ms)"
+
+	w, h := vg.Millimeters(float64(len(results))), vg.Centimeters(12.0)
+	canvas := vgsvg.New(w, h)
+	p.Draw(plot.MakeDrawArea(canvas))
+
+	_, err = canvas.WriteTo(out)
+	return err
+}

lib/text_reporter.go

@@ -1,76 +0,0 @@
-package vegeta
-
-import (
-	"fmt"
-	"io"
-	"text/tabwriter"
-	"time"
-)
-
-// 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/timings_plot_reporter.go

@@ -1,90 +0,0 @@
-package vegeta
-
-import (
-	"code.google.com/p/plotinum/plot"
-	"code.google.com/p/plotinum/plotter"
-	"code.google.com/p/plotinum/plotutil"
-	"code.google.com/p/plotinum/vg"
-	"code.google.com/p/plotinum/vg/vgsvg"
-	"container/list"
-	"io"
-	"time"
-)
-
-type TimingsPlotReporter struct {
-	responses *list.List
-}
-
-// NewTimingsPlotReporter initializes a TimingsPlotReporter
-func NewTimingsPlotReporter() *TimingsPlotReporter {
-	return &TimingsPlotReporter{responses: list.New()}
-}
-
-// add inserts response to be used in the report, sorted by timestamp.
-func (r *TimingsPlotReporter) add(res *result) {
-	// Empty list
-	if r.responses.Len() == 0 {
-		r.responses.PushFront(res)
-		return
-	}
-	// Happened after all others
-	if last := r.responses.Back().Value.(*result); last.timestamp.Before(res.timestamp) {
-		r.responses.PushBack(res)
-		return
-	}
-	// Happened before all others
-	if first := r.responses.Front().Value.(*result); first.timestamp.After(res.timestamp) {
-		r.responses.PushFront(res)
-		return
-	}
-	// O(n) worst case insertion time
-	for e := r.responses.Front(); e != nil; e = e.Next() {
-		needle := e.Value.(*result)
-		if res.timestamp.Before(needle.timestamp) {
-			r.responses.InsertBefore(res, e)
-			return
-		}
-	}
-}
-
-// Report builds up a plot of the response times of the requests
-// in SVG format and writes it to out
-func (r *TimingsPlotReporter) Report(out io.Writer) error {
-	timestamps := make([]time.Time, 0)
-	timings := make([]time.Duration, 0)
-
-	for e := r.responses.Front(); e != nil; e = e.Next() {
-		r := e.Value.(*result)
-		timestamps = append(timestamps, r.timestamp)
-		timings = append(timings, r.timing)
-	}
-
-	p, err := plot.New()
-	if err != nil {
-		return err
-	}
-	pts := make(plotter.XYs, len(timestamps))
-	for i := 0; i < len(pts); i++ {
-		pts[i].X = timestamps[i].Sub(timestamps[0]).Seconds()
-		pts[i].Y = timings[i].Seconds() * 1000
-	}
-
-	line, err := plotter.NewLine(pts)
-	if err != nil {
-		return err
-	}
-	line.Color = plotutil.Color(1)
-
-	p.Add(line)
-	p.X.Padding = vg.Length(3.0)
-	p.X.Label.Text = "Time elapsed"
-	p.Y.Padding = vg.Length(3.0)
-	p.Y.Label.Text = "Latency (ms)"
-
-	w, h := vg.Millimeters(float64(len(timestamps))), vg.Centimeters(12.0)
-	canvas := vgsvg.New(w, h)
-	p.Draw(plot.MakeDrawArea(canvas))
-
-	_, err = canvas.WriteTo(out)
-	return err
-}

main.go

@@ -71,12 +71,12 @@ func run(rate uint64, duration time.Duration, targetsf, ordering, reporter, outp
 	var rep vegeta.Reporter
 	switch reporter {
 	case "text":
-		rep = vegeta.NewTextReporter()
+		rep = vegeta.ReportText
 	case "plot:timings":
-		rep = vegeta.NewTimingsPlotReporter()
+		rep = vegeta.ReportTimingsPlot
 	default:
 		log.Println("Reporter provided is not supported. Using text")
-		rep = vegeta.NewTextReporter()
+		rep = vegeta.ReportText
 	}
 
 	targets, err := vegeta.NewTargetsFromFile(targetsf)
@@ -94,11 +94,10 @@ func run(rate uint64, duration time.Duration, targetsf, ordering, reporter, outp
 	}
 
 	log.Printf("Vegeta is attacking %d targets in %s order for %s...\n", len(targets), ordering, duration)
-	vegeta.Attack(targets, rate, duration, rep)
+	results := vegeta.Attack(targets, rate, duration)
 	log.Println("Done!")
-
 	log.Printf("Writing report to '%s'...", output)
-	if err = rep.Report(out); err != nil {
+	if err = rep(results, out); err != nil {
 		return fmt.Errorf(errReportingPrefix+"%s", err)
 	}
 	return nil