Merge pull request #27 from kffl/feat/library

Enable usage of gocannon as a library/dependency

args.go

@@ -19,50 +19,51 @@ func parseRequestHeaders(s kingpin.Settings) *common.RequestHeaders {
 	return r
 }
 
-var app = kingpin.New("gocannon", "Performance-focused HTTP load testing tool.")
+func parseArgs() (common.Config, error) {
 
-var config = common.Config{
-	Duration: app.Flag("duration", "Load test duration.").
-		Short('d').
-		Default("10s").
-		Duration(),
-	Connections: app.Flag("connections", "Maximum number of concurrent connections.").
-		Short('c').
-		Default("50").
-		Int(),
-	Timeout: app.Flag("timeout", "HTTP client timeout.").
-		Short('t').
-		Default("200ms").
-		Duration(),
-	Mode: app.Flag("mode", "Statistics collection mode: reqlog (logs each request) or hist (stores histogram of completed requests latencies).").
-		Default("reqlog").
-		Short('m').
-		String(),
-	OutputFile: app.Flag("output", "File to save the request log in CSV format (reqlog mode) or a text file with raw histogram data (hist mode).").
-		PlaceHolder("file.csv").
-		Short('o').
-		String(),
-	Interval: app.Flag("interval", "Interval for statistics calculation (reqlog mode).").
-		Default("250ms").
-		Short('i').
-		Duration(),
-	Preallocate: app.Flag("preallocate", "Number of requests in req log to preallocate memory for per connection (reqlog mode).").
-		Default("1000").
-		Int(),
-	Method:   app.Flag("method", "The HTTP request method (GET, POST, PUT, PATCH or DELETE).").Default("GET").Enum("GET", "POST", "PUT", "PATCH", "DELETE"),
-	Body:     parseRequestBody(app.Flag("body", "HTTP request body.").Short('b').PlaceHolder("\"{data...\"")),
-	Headers:  parseRequestHeaders(kingpin.Flag("header", "HTTP request header(s). You can set more than one header by repeating this flag.").Short('h').PlaceHolder("\"k:v\"")),
-	TrustAll: app.Flag("trust-all", "Omit SSL certificate validation.").Bool(),
-	Format: app.Flag("format", "Load test report format. Either 'default' (verbose), 'json' or 'yaml'. When json or yaml is specified, apart from the load test results, no additional info will be written to std out.").
-		Short('f').
-		Default("default").
-		Enum("default", "json", "yaml"),
-	Plugin: app.Flag("plugin", "Plugin to run Gocannon with (path to .so file).").PlaceHolder("/to/p.so").ExistingFile(),
-	Target: app.Arg("target", "HTTP target URL with port (i.e. http://localhost:80/test or https://host:443/x)").Required().String(),
-}
+	var app = kingpin.New("gocannon", "Performance-focused HTTP load testing tool.")
+
+	var config = common.Config{
+		Duration: app.Flag("duration", "Load test duration.").
+			Short('d').
+			Default("10s").
+			Duration(),
+		Connections: app.Flag("connections", "Maximum number of concurrent connections.").
+			Short('c').
+			Default("50").
+			Int(),
+		Timeout: app.Flag("timeout", "HTTP client timeout.").
+			Short('t').
+			Default("200ms").
+			Duration(),
+		Mode: app.Flag("mode", "Statistics collection mode: reqlog (logs each request) or hist (stores histogram of completed requests latencies).").
+			Default("reqlog").
+			Short('m').
+			String(),
+		OutputFile: app.Flag("output", "File to save the request log in CSV format (reqlog mode) or a text file with raw histogram data (hist mode).").
+			PlaceHolder("file.csv").
+			Short('o').
+			String(),
+		Interval: app.Flag("interval", "Interval for statistics calculation (reqlog mode).").
+			Default("250ms").
+			Short('i').
+			Duration(),
+		Preallocate: app.Flag("preallocate", "Number of requests in req log to preallocate memory for per connection (reqlog mode).").
+			Default("1000").
+			Int(),
+		Method:   app.Flag("method", "The HTTP request method (GET, POST, PUT, PATCH or DELETE).").Default("GET").Enum("GET", "POST", "PUT", "PATCH", "DELETE"),
+		Body:     parseRequestBody(app.Flag("body", "HTTP request body.").Short('b').PlaceHolder("\"{data...\"")),
+		Headers:  parseRequestHeaders(kingpin.Flag("header", "HTTP request header(s). You can set more than one header by repeating this flag.").Short('h').PlaceHolder("\"k:v\"")),
+		TrustAll: app.Flag("trust-all", "Omit SSL certificate validation.").Bool(),
+		Format: app.Flag("format", "Load test report format. Either 'default' (verbose), 'json' or 'yaml'. When json or yaml is specified, apart from the load test results, no additional info will be written to std out.").
+			Short('f').
+			Default("default").
+			Enum("default", "json", "yaml"),
+		Plugin: app.Flag("plugin", "Plugin to run Gocannon with (path to .so file).").PlaceHolder("/to/p.so").ExistingFile(),
+		Target: app.Arg("target", "HTTP target URL with port (i.e. http://localhost:80/test or https://host:443/x)").Required().String(),
+	}
 
-func parseArgs() error {
 	app.Version("1.1.0")
 	_, err := app.Parse(os.Args[1:])
-	return err
+	return config, err
 }

gocannon.go

@@ -14,46 +14,59 @@ func exitWithError(err error) {
 	os.Exit(1)
 }
 
-func runGocannon(cfg common.Config) error {
-	var gocannonPlugin common.GocannonPlugin
+// Gocannon represents a single gocannon instance with a config defined upon its creation.
+type Gocannon struct {
+	cfg    common.Config
+	client *fasthttp.HostClient
+	stats  statsCollector
+	plugin common.GocannonPlugin
+}
+
+// NewGocannon creates a new gocannon instance using a provided config.
+func NewGocannon(cfg common.Config) (Gocannon, error) {
 	var err error
 
-	if *config.Format == "default" {
-		printHeader(config)
-	}
+	gocannon := Gocannon{cfg: cfg}
 
 	if *cfg.Plugin != "" {
-		gocannonPlugin, err = loadPlugin(*cfg.Plugin, *cfg.Format != "default")
+		gocannonPlugin, err := loadPlugin(*cfg.Plugin, *cfg.Format != "default")
 		if err != nil {
-			return err
+			return gocannon, err
 		}
+		gocannon.plugin = gocannonPlugin
 		gocannonPlugin.Startup(cfg)
 	}
 
 	c, err := newHTTPClient(*cfg.Target, *cfg.Timeout, *cfg.Connections, *cfg.TrustAll, true)
 
 	if err != nil {
-		return err
+		return gocannon, err
 	}
 
-	n := *cfg.Connections
+	gocannon.client = c
 
-	stats, scErr := newStatsCollector(*cfg.Mode, n, *cfg.Preallocate, *cfg.Timeout)
+	stats, scErr := newStatsCollector(*cfg.Mode, *cfg.Connections, *cfg.Preallocate, *cfg.Timeout)
+
+	gocannon.stats = stats
 
 	if scErr != nil {
-		return scErr
+		return gocannon, scErr
 	}
 
+	return gocannon, nil
+}
+
+// Run performs the load test.
+func (g Gocannon) Run() (TestResults, error) {
+
+	n := *g.cfg.Connections
+
 	var wg sync.WaitGroup
 
 	wg.Add(n)
 
 	start := makeTimestamp()
-	stop := start + cfg.Duration.Nanoseconds()
-
-	if *cfg.Format == "default" {
-		fmt.Printf("gocannon goes brr...\n")
-	}
+	stop := start + g.cfg.Duration.Nanoseconds()
 
 	for connectionID := 0; connectionID < n; connectionID++ {
 		go func(c *fasthttp.HostClient, cid int, p common.GocannonPlugin) {
@@ -65,43 +78,25 @@ func runGocannon(cfg common.Config) error {
 					plugTarget, plugMethod, plugBody, plugHeaders := p.BeforeRequest(cid)
 					code, start, end = performRequest(c, plugTarget, plugMethod, plugBody, plugHeaders)
 				} else {
-					code, start, end = performRequest(c, *cfg.Target, *cfg.Method, *cfg.Body, *cfg.Headers)
+					code, start, end = performRequest(c, *g.cfg.Target, *g.cfg.Method, *g.cfg.Body, *g.cfg.Headers)
 				}
 				if end >= stop {
 					break
 				}
 
-				stats.RecordResponse(cid, code, start, end)
+				g.stats.RecordResponse(cid, code, start, end)
 			}
 			wg.Done()
-		}(c, connectionID, gocannonPlugin)
+		}(g.client, connectionID, g.plugin)
 	}
 
 	wg.Wait()
 
-	err = stats.CalculateStats(start, stop, *cfg.Interval, *cfg.OutputFile)
+	err := g.stats.CalculateStats(start, stop, *g.cfg.Interval, *g.cfg.OutputFile)
 
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	if *cfg.Format == "default" {
-		printSummary(stats)
-	}
-	stats.PrintReport(*cfg.Format)
-
-	return nil
-}
-
-func main() {
-	err := parseArgs()
-	if err != nil {
-		exitWithError(err)
-	}
-
-	err = runGocannon(config)
-
-	if err != nil {
-		exitWithError(err)
-	}
+	return g.stats, err
 }

hist/histogram.go

@@ -141,3 +141,11 @@ func (h *requestHist) GetReqPerSec() float64 {
 func (h *requestHist) GetLatencyAvg() float64 {
 	return h.results.LatencyAvg * 1000.0
 }
+
+func (h *requestHist) GetLatencyPercentiles() []int64 {
+	asNanoseconds := make([]int64, len(h.results.LatencyPercentiles))
+	for i, p := range h.results.LatencyPercentiles {
+		asNanoseconds[i] = p * 1000
+	}
+	return asNanoseconds
+}

integration_test.go

@@ -126,7 +126,15 @@ func TestGocannonDefaultValues(t *testing.T) {
 		Target:      &target,
 	}
 
-	assert.Nil(t, runGocannon(cfg), "the load test should be completed without errors")
+	g, creationErr := NewGocannon(cfg)
+
+	assert.Nil(t, creationErr, "gocannon instance should be created without errors")
+
+	if creationErr == nil {
+		results, execErr := g.Run()
+		assert.Nil(t, execErr, "the load test should be completed without errors")
+		assert.Greater(t, results.GetReqPerSec(), 100.0, "a throughput of at least 100 req/s should be achieved")
+	}
 }
 
 func TestGocanonWithPlugin(t *testing.T) {
@@ -167,5 +175,16 @@ func TestGocanonWithPlugin(t *testing.T) {
 		Target:      &target,
 	}
 
-	assert.Nil(t, runGocannon(cfg), "the load test should be completed without errors")
+	g, creationErr := NewGocannon(cfg)
+
+	assert.Nil(t, creationErr, "gocannon instance with a plugin should be created without errors")
+
+	if creationErr == nil {
+		results, execErr := g.Run()
+
+		assert.Nil(t, execErr, "the load test should be completed without errors")
+
+		assert.Greater(t, results.GetReqPerSec(), 100.0, "a throughput of at least 100 req/s should be achieved")
+	}
+
 }

main.go

@@ -0,0 +1,35 @@
+package main
+
+import "fmt"
+
+func main() {
+	config, err := parseArgs()
+	if err != nil {
+		exitWithError(err)
+	}
+
+	if *config.Format == "default" {
+		printHeader(config)
+	}
+
+	g, err := NewGocannon(config)
+
+	if err != nil {
+		exitWithError(err)
+	}
+
+	if *config.Format == "default" {
+		fmt.Printf("gocannon goes brr...\n")
+	}
+
+	results, err := g.Run()
+
+	if *config.Format == "default" {
+		printSummary(results)
+	}
+	results.PrintReport(*config.Format)
+
+	if err != nil {
+		exitWithError(err)
+	}
+}

print.go

@@ -10,7 +10,7 @@ func printHeader(cfg common.Config) {
 	fmt.Printf("Attacking %s with %d connections over %s\n", *cfg.Target, *cfg.Connections, *cfg.Duration)
 }
 
-func printSummary(s statsCollector) {
+func printSummary(s TestResults) {
 	fmt.Printf("Total Req:  %8d\n", s.GetReqCount())
 	fmt.Printf("Req/s:      %11.2f\n", s.GetReqPerSec())
 }

reqlog/request.go

@@ -122,6 +122,10 @@ func (r *requestLogCollector) GetLatencyAvg() float64 {
 	return r.results.Summary.LatencyAVG
 }
 
+func (r *requestLogCollector) GetLatencyPercentiles() []int64 {
+	return r.results.Summary.LatencyPercentiles
+}
+
 func (r *requestLogCollector) saveResCodes() {
 	for _, connLog := range *r.reqLog {
 		for _, req := range connLog {

stats_collector.go

@@ -8,13 +8,19 @@ import (
 	"github.com/kffl/gocannon/reqlog"
 )
 
-type statsCollector interface {
-	RecordResponse(conn int, code int, start int64, end int64)
-	CalculateStats(start int64, stop int64, interval time.Duration, fileName string) error
-	PrintReport(format string)
+// TestResults allows for accessing the load test results.
+type TestResults interface {
 	GetReqCount() int64
 	GetReqPerSec() float64
 	GetLatencyAvg() float64
+	GetLatencyPercentiles() []int64
+	PrintReport(format string)
+}
+
+type statsCollector interface {
+	RecordResponse(conn int, code int, start int64, end int64)
+	CalculateStats(start int64, stop int64, interval time.Duration, fileName string) error
+	TestResults
 }
 
 func newStatsCollector(