Merge pull request #18 from kffl/feat/plugins

Support for custom plugins

CITATION.cff

@@ -6,4 +6,4 @@ license: "Apache-2.0"
 message: "Please use the following citation metadata."
 repository-code: "https://github.com/kffl/gocannon"
 title: "Gocannon - Performance-focused HTTP benchmarking tool"
-version: "0.2.1"
+version: "1.0.0"

README.md

@@ -27,6 +27,7 @@ Flags:
   -h, --header="k:v" ...  HTTP request header(s). You can set more than one header by
                           repeating this flag.
       --trust-all         Omit SSL certificate validation
+      --plugin=/to/p.so   Plugin to run Gocannon with
       --version           Show application version.
 
 Args:
@@ -112,6 +113,17 @@ Similarly to saving CSV output in request log mode, you can write the histogram
 45
 [...]
 ```
+### Custom plugins
+
+Gocannon supports user-provided plugins, which can customize the requests sent during the load test. A custom plugin has to satisfy the `GocannonPlugin` interface defined in the `common` package, which also contains types required for plugin development. An example implementation with additional comments is provided in `_example_plugin` folder.
+
+In order to build a plugin, use the following command inside its directory:
+
+```
+go build -buildmode=plugin -o plugin.so plugin.go
+```
+
+Once you obtain a shared object (`.so`) file, you can provide a path to it via `--plugin` flag. Bare in mind that a custom plugin `.so` file and the Gocannon binary using it must both be compiled using the same Go version.
 
 ## Load testing recommendations
 

_example_plugin/plugin.go

@@ -0,0 +1,56 @@
+package main
+
+import (
+	"fmt"
+	"sync/atomic"
+
+	"github.com/kffl/gocannon/common"
+)
+
+type plugin string
+
+var config common.Config
+
+// you can use global variables within the plugin to persist its state between invocations of BeforeRequest
+var reqCounter int64 = 0
+
+func (p plugin) Startup(cfg common.Config) {
+	// saving the config for later
+	// make sure not to mutate the config contents
+	config = cfg
+}
+
+func (p plugin) BeforeRequest(cid int) (target string, method string, body common.RawRequestBody, headers common.RequestHeaders) {
+	// there can be multiple invocations of BeforeRequest (with different connections ids) happening in parallel
+	// therefore it's necessary to ensure thread-safe usage of plugin global variables
+
+	// as an example, we are going to use atomic add operation
+	// in order to track how many requests were sent so far
+	reqNum := atomic.AddInt64(&reqCounter, 1)
+
+	headers = *config.Headers
+
+	// every 10th request, we want to add a special header
+	if reqNum%10 == 0 {
+		headers = append(*config.Headers, common.RequestHeader{"X-Special-Header", "gocannon"})
+	}
+
+	// appending connectionID to the target (i.e. http://target:123/?connection=5)
+	target = fmt.Sprintf("%s?connection=%d", *config.Target, cid)
+
+	// we leave the HTTP method unchanged (from config)
+	method = *config.Method
+
+	// and the same for body
+	// the body shall not be mutated after being passed as a return value, as gocannon uses fasthttp's SetBodyRaw
+	body = *config.Body
+
+	return
+}
+
+func (p plugin) GetName() string {
+	return string(p)
+}
+
+// GocannonPlugin is an exported instance of the plugin (the name "GocannonPlugin" is mandatory)
+var GocannonPlugin plugin = "Sample Plugin"

args.go

@@ -1,100 +1,64 @@
 package main
 
 import (
-	"fmt"
-	"strings"
+	"os"
 
+	"github.com/kffl/gocannon/common"
 	"gopkg.in/alecthomas/kingpin.v2"
 )
 
-type rawRequestBody []byte
-
-func (b *rawRequestBody) Set(value string) error {
-	(*b) = []byte(value)
-	return nil
-}
-
-func (b *rawRequestBody) String() string {
-	return fmt.Sprint(*b)
-}
-
-func (b *rawRequestBody) IsCumulative() bool {
-	return false
-}
-
-func parseRequestBody(s kingpin.Settings) *rawRequestBody {
-	r := &rawRequestBody{}
-	s.SetValue((*rawRequestBody)(r))
+func parseRequestBody(s kingpin.Settings) *common.RawRequestBody {
+	r := &common.RawRequestBody{}
+	s.SetValue((*common.RawRequestBody)(r))
 	return r
 }
 
-type requestHeader struct {
-	key   string
-	value string
-}
-
-type requestHeaders []requestHeader
-
-func (r *requestHeaders) Set(value string) error {
-	tokenized := strings.Split(value, ":")
-	if len(tokenized) != 2 {
-		return fmt.Errorf("Header '%s' doesn't match 'Key:Value' format (i.e. 'Content-Type:application/json')", value)
-	}
-	h := requestHeader{tokenized[0], tokenized[1]}
-	(*r) = append(*r, h)
-	return nil
-}
-
-func (r *requestHeaders) String() string {
-	return fmt.Sprint(*r)
-}
-
-func (r *requestHeaders) IsCumulative() bool {
-	return true
-}
-
-func parseRequestHeaders(s kingpin.Settings) *requestHeaders {
-	r := &requestHeaders{}
-	s.SetValue((*requestHeaders)(r))
+func parseRequestHeaders(s kingpin.Settings) *common.RequestHeaders {
+	r := &common.RequestHeaders{}
+	s.SetValue((*common.RequestHeaders)(r))
 	return r
 }
 
-var (
-	duration = kingpin.Flag("duration", "Load test duration").
-			Short('d').
-			Default("10s").
-			Duration()
-	connections = kingpin.Flag("connections", "Maximum number of concurrent connections").
-			Short('c').
-			Default("50").
-			Int()
-	timeout = kingpin.Flag("timeout", "HTTP client timeout").
+var app = kingpin.New("gocannon", "Performance-focused HTTP benchmarking 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 = kingpin.Flag("mode", "Statistics collection mode: reqlog (logs each request) or hist (stores histogram of completed requests latencies)").
+		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 = kingpin.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 = kingpin.Flag("interval", "Interval for statistics calculation (reqlog mode)").
-			Default("250ms").
-			Short('i').
-			Duration()
-	preallocate = kingpin.Flag("preallocate", "Number of requests in req log to preallocate memory for per connection (reqlog mode)").
-			Default("1000").
-			Int()
-	method   = kingpin.Flag("method", "The HTTP request method (GET, POST, PUT, PATCH or DELETE)").Default("GET").Enum("GET", "POST", "PUT", "PATCH", "DELETE")
-	body     = parseRequestBody(kingpin.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 = kingpin.Flag("trust-all", "Omit SSL certificate validation").Bool()
-	target   = kingpin.Arg("target", "HTTP target URL with port (i.e. http://localhost:80/test or https://host:443/x)").Required().String()
-)
+		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(),
+	Plugin:   app.Flag("plugin", "Plugin to run Gocannon with").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() {
-	kingpin.Version("0.2.1")
-	kingpin.Parse()
+func parseArgs() error {
+	app.Version("1.0.0")
+	_, err := app.Parse(os.Args[1:])
+	return err
 }

common/config.go

@@ -0,0 +1,20 @@
+package common
+
+import "time"
+
+// Config is a struct containing all of the parsed CLI flags and arguments
+type Config struct {
+	Duration    *time.Duration
+	Connections *int
+	Timeout     *time.Duration
+	Mode        *string
+	OutputFile  *string
+	Interval    *time.Duration
+	Preallocate *int
+	Method      *string
+	Body        *RawRequestBody
+	Headers     *RequestHeaders
+	TrustAll    *bool
+	Plugin      *string
+	Target      *string
+}

common/headers.go

@@ -0,0 +1,33 @@
+package common
+
+import (
+	"fmt"
+	"strings"
+)
+
+// RequestHeader represents a single HTTP request header (a key: value pair)
+type RequestHeader struct {
+	Key   string
+	Value string
+}
+
+// RequestHeaders is a slice of request headers that will be added to the request
+type RequestHeaders []RequestHeader
+
+func (r *RequestHeaders) Set(value string) error {
+	tokenized := strings.Split(value, ":")
+	if len(tokenized) != 2 {
+		return fmt.Errorf("Header '%s' doesn't match 'Key:Value' format (i.e. 'Content-Type:application/json')", value)
+	}
+	h := RequestHeader{tokenized[0], tokenized[1]}
+	(*r) = append(*r, h)
+	return nil
+}
+
+func (r *RequestHeaders) String() string {
+	return fmt.Sprint(*r)
+}
+
+func (r *RequestHeaders) IsCumulative() bool {
+	return true
+}

args_test.go → common/headers_test.go

@@ -1,4 +1,4 @@
-package main
+package common
 
 import (
 	"testing"
@@ -7,7 +7,7 @@ import (
 )
 
 func TestSetRequestHeaders(t *testing.T) {
-	r := requestHeaders{}
+	r := RequestHeaders{}
 
 	errHeaderOk := r.Set("Content-Type:application/json")
 	errHeaderWrong := r.Set("WrongHeader")

common/plugin_interface.go

@@ -0,0 +1,11 @@
+package common
+
+// GocannonPlugin is an interface that has to be satisfied by a custom gocannnon plugin
+type GocannonPlugin interface {
+	// function called on gocannon startup with a config passed to it
+	Startup(cfg Config)
+	// function called before each request is sent
+	BeforeRequest(cid int) (target string, method string, body RawRequestBody, headers RequestHeaders)
+	// function that returns the plugin's name
+	GetName() string
+}

common/request_body.go

@@ -0,0 +1,18 @@
+package common
+
+import "fmt"
+
+type RawRequestBody []byte
+
+func (b *RawRequestBody) Set(value string) error {
+	(*b) = []byte(value)
+	return nil
+}
+
+func (b *RawRequestBody) String() string {
+	return fmt.Sprint(*b)
+}
+
+func (b *RawRequestBody) IsCumulative() bool {
+	return false
+}