Add e2e mode for opa bench

This commit adds a new flag to the opa bench command which
allows users to run benchmarks against a running OPA server.
This mode can be used to evaluate the additional overhead the
server is going to introduce.

Co-authored-by: Anders Eknert anders@eknert.com
Signed-off-by: Ashutosh Narkar <anarkar4387@gmail.com>

cmd/bench.go

@@ -5,14 +5,26 @@
 package cmd
 
 import (
+	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"math"
+	"net/http"
 	"os"
 	"sort"
+	"strings"
 	"testing"
 
+	"github.com/open-policy-agent/opa/ast"
+
+	"github.com/open-policy-agent/opa/server/types"
+
+	"github.com/open-policy-agent/opa/logging"
+	"github.com/open-policy-agent/opa/runtime"
+
 	"github.com/olekukonko/tablewriter"
 	"github.com/spf13/cobra"
 
@@ -28,8 +40,11 @@ import (
 // ones compatible with running a benchmark.
 type benchmarkCommandParams struct {
 	evalCommandParams
-	benchMem bool
-	count    int
+	benchMem               bool
+	count                  int
+	e2e                    bool
+	gracefulShutdownPeriod int
+	shutdownWaitPeriod     int
 }
 
 const (
@@ -47,6 +62,7 @@ func newBenchmarkEvalParams() benchmarkCommandParams {
 			target: util.NewEnumFlag(compile.TargetRego, []string{compile.TargetRego, compile.TargetWasm}),
 			schema: &schemaFlags{},
 		},
+		gracefulShutdownPeriod: 10,
 	}
 }
 
@@ -67,6 +83,8 @@ Example with bundle and input data:
 
 To enable more detailed analysis use the --metrics and --benchmem flags.
 
+To run benchmarks against a running OPA server to evaluate server overhead use the --e2e flag.
+
 The optional "gobench" output format conforms to the Go Benchmark Data Format.
 `,
 
@@ -105,6 +123,11 @@ The optional "gobench" output format conforms to the Go Benchmark Data Format.
 	addCountFlag(benchCommand.Flags(), &params.count, "benchmark")
 	addBenchmemFlag(benchCommand.Flags(), &params.benchMem, true)
 
+	addE2EFlag(benchCommand.Flags(), &params.e2e, false)
+
+	benchCommand.Flags().IntVar(&params.gracefulShutdownPeriod, "shutdown-grace-period", 10, "set the time (in seconds) that the server will wait to gracefully shut down. This flag is valid in 'e2e' mode only.")
+	benchCommand.Flags().IntVar(&params.shutdownWaitPeriod, "shutdown-wait-period", 0, "set the time (in seconds) that the server will wait before initiating shutdown. This flag is valid in 'e2e' mode only.")
+
 	RootCommand.AddCommand(benchCommand)
 }
 
@@ -113,13 +136,24 @@ type benchRunner interface {
 }
 
 func benchMain(args []string, params benchmarkCommandParams, w io.Writer, r benchRunner) (int, error) {
+
+	ctx := context.Background()
+
+	if params.e2e {
+		err := benchE2E(ctx, args, params, w)
+		if err != nil {
+			errRender := renderBenchmarkError(params, err, w)
+			return 1, errRender
+		}
+		return 0, nil
+	}
+
 	ectx, err := setupEval(args, params.evalCommandParams)
 	if err != nil {
 		errRender := renderBenchmarkError(params, err, w)
 		return 1, errRender
 	}
 
-	ctx := context.Background()
 	var benchFunc func(context.Context, ...rego.EvalOption) error
 	rg := rego.New(ectx.regoArgs...)
 
@@ -177,13 +211,9 @@ type goBenchRunner struct {
 
 func (r *goBenchRunner) run(ctx context.Context, ectx *evalContext, params benchmarkCommandParams, f func(context.Context, ...rego.EvalOption) error) (testing.BenchmarkResult, error) {
 
-	var hist metrics.Metrics
+	var hist, m metrics.Metrics
 	if params.metrics {
 		hist = metrics.New()
-	}
-
-	var m metrics.Metrics
-	if params.metrics {
 		m = metrics.New()
 	}
 
@@ -221,7 +251,7 @@ func (r *goBenchRunner) run(ctx context.Context, ectx *evalContext, params bench
 			if params.metrics {
 				for name, metric := range m.All() {
 					// Note: We only support int64 metrics right now, this should cover pretty
-					// much all of the ones we would care about (timers and counters).
+					// much all the ones we would care about (timers and counters).
 					switch v := metric.(type) {
 					case int64:
 						hist.Histogram(name).Update(v)
@@ -232,32 +262,286 @@ func (r *goBenchRunner) run(ctx context.Context, ectx *evalContext, params bench
 		}
 
 		if params.metrics {
-			// For each histogram add their values to the benchmark results.
-			// Note: If there are many metrics this gets super verbose.
-			for histName, metric := range hist.All() {
-				histValues, ok := metric.(map[string]interface{})
-				if !ok {
-					continue
-				}
-				for metricName, rawValue := range histValues {
-					unit := fmt.Sprintf("%s_%s", histName, metricName)
-
-					// Only support histogram metrics that are a float64 or int64,
-					// this covers the stock implementation in metrics.Metrics
-					switch v := rawValue.(type) {
-					case int64:
-						b.ReportMetric(float64(v), unit)
-					case float64:
-						b.ReportMetric(v, unit)
-					}
-				}
-			}
+			reportMetrics(b, hist.All())
 		}
 	})
 
 	return br, benchErr
 }
 
+func benchE2E(ctx context.Context, args []string, params benchmarkCommandParams, w io.Writer) error {
+	host := "localhost"
+	port := "18181"
+
+	logger := logging.New()
+	logger.SetLevel(logging.Error)
+
+	paths := params.dataPaths.v
+	if len(params.bundlePaths.v) > 0 {
+		paths = append(paths, params.bundlePaths.v...)
+	}
+
+	rtParams := runtime.Params{
+		Addrs:                  &[]string{fmt.Sprintf("%s:%s", host, port)},
+		Paths:                  paths,
+		Logger:                 logger,
+		BundleMode:             params.bundlePaths.isFlagSet(),
+		SkipBundleVerification: true,
+		EnableVersionCheck:     false,
+		GracefulShutdownPeriod: params.gracefulShutdownPeriod,
+		ShutdownWaitPeriod:     params.shutdownWaitPeriod,
+	}
+
+	rt, err := runtime.NewRuntime(ctx, rtParams)
+	if err != nil {
+		return err
+	}
+
+	cctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	initChannel := rt.Manager.ServerInitializedChannel()
+
+	done := make(chan error)
+	go func() {
+		done <- rt.Serve(cctx)
+	}()
+
+	select {
+	case err := <-done:
+		if err != nil {
+			return err
+		}
+	case <-initChannel:
+		break
+	}
+
+	query, err := readQuery(params, args)
+	if err != nil {
+		return err
+	}
+
+	input, err := readInputBytes(params.evalCommandParams)
+	if err != nil {
+		return err
+	}
+
+	// Wrap input in "input" attribute
+	inp := make(map[string]interface{})
+
+	if input != nil {
+		if err = util.Unmarshal(input, &inp); err != nil {
+			return err
+		}
+	}
+
+	body := map[string]interface{}{"input": inp}
+
+	var path string
+	if params.partial {
+		path = "compile"
+		body["query"] = query
+		if len(params.unknowns) > 0 {
+			body["unknowns"] = params.unknowns
+		}
+	} else {
+		_, err := ast.ParseBody(query)
+		if err != nil {
+			return fmt.Errorf("error occurred while parsing query")
+		}
+
+		if strings.HasPrefix(query, "data.") {
+			path = strings.ReplaceAll(query, ".", "/")
+		} else {
+			path = "query"
+			body["query"] = query
+		}
+	}
+
+	url := fmt.Sprintf("http://%s:%s/v1/%v", host, port, path)
+	if params.metrics {
+		url += "?metrics=true"
+	}
+
+	for i := 0; i < params.count; i++ {
+		br, err := runE2E(params, url, body)
+		if err != nil {
+			return err
+		}
+		renderBenchmarkResult(params, br, w)
+	}
+	return nil
+}
+
+func runE2E(params benchmarkCommandParams, url string, input map[string]interface{}) (testing.BenchmarkResult, error) {
+	hist := metrics.New()
+
+	var benchErr error
+
+	br := testing.Benchmark(func(b *testing.B) {
+
+		// Track memory allocations, if enabled
+		if params.benchMem {
+			b.ReportAllocs()
+		}
+
+		// Reset the histogram for each invocation of the bench function
+		hist.Clear()
+
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+
+			// Start the timer
+			b.StartTimer()
+
+			// Execute the query API call
+			m, err := e2eQuery(params, url, input)
+
+			// Stop the timer while processing the results
+			b.StopTimer()
+			if err != nil {
+				benchErr = err
+				b.FailNow()
+			}
+
+			// Add metrics for that evaluation into the top level histogram
+			if params.metrics {
+				for name, metric := range m {
+					switch v := metric.(type) {
+					case json.Number:
+						num, err := v.Int64()
+						if err != nil {
+							benchErr = err
+							b.FailNow()
+						}
+						hist.Histogram(name).Update(num)
+
+					}
+				}
+			}
+		}
+
+		if params.metrics {
+			reportMetrics(b, hist.All())
+		}
+	})
+
+	return br, benchErr
+}
+
+func e2eQuery(params benchmarkCommandParams, url string, input map[string]interface{}) (types.MetricsV1, error) {
+
+	reqBody, err := json.Marshal(input)
+	if err != nil {
+		return nil, err
+	}
+
+	req, err := http.NewRequest("POST", url, bytes.NewReader(reqBody))
+	if err != nil {
+		return nil, fmt.Errorf("unexpected error: %w", err)
+	}
+
+	resp, err := http.DefaultClient.Do(req)
+	if err != nil {
+		return nil, err
+	}
+
+	defer resp.Body.Close()
+
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	if resp.StatusCode != 200 {
+		var e map[string]interface{}
+		if err = util.Unmarshal(body, &e); err != nil {
+			return nil, err
+		}
+
+		if _, ok := e["errors"]; !ok {
+			return nil, fmt.Errorf("request failed, OPA server replied with HTTP %v: %v", resp.StatusCode, e["message"])
+		}
+
+		bs, err := json.Marshal(e["errors"])
+		if err != nil {
+			return nil, err
+		}
+
+		var astErrs ast.Errors
+		if err = util.Unmarshal(bs, &astErrs); err != nil {
+			// ignore err
+			return nil, fmt.Errorf("request failed, OPA server replied with HTTP %v: %v", resp.StatusCode, e["message"])
+		}
+
+		return nil, astErrs
+	}
+
+	if !params.partial {
+		var result types.DataResponseV1
+		if err = util.Unmarshal(body, &result); err != nil {
+			return nil, err
+		}
+
+		if result.Result == nil && params.fail {
+			return nil, fmt.Errorf("undefined result")
+		}
+
+		return result.Metrics, nil
+	}
+
+	var result types.CompileResponseV1
+	if err = util.Unmarshal(body, &result); err != nil {
+		return nil, err
+	}
+
+	if params.fail {
+		if result.Result == nil {
+			return nil, fmt.Errorf("undefined result")
+		}
+
+		i := *result.Result
+
+		peResult, ok := i.(map[string]interface{})
+		if !ok {
+			return nil, fmt.Errorf("invalid result for compile response")
+		}
+
+		if len(peResult) == 0 {
+			return nil, fmt.Errorf("undefined result")
+		}
+
+		if val, ok := peResult["queries"]; ok {
+			queries, ok := val.([]interface{})
+			if !ok {
+				return nil, fmt.Errorf("invalid result for output of partial evaluation")
+			}
+
+			if len(queries) == 0 {
+				return nil, fmt.Errorf("undefined result")
+			}
+		} else {
+			return nil, fmt.Errorf("invalid result for output of partial evaluation")
+		}
+	}
+
+	return result.Metrics, nil
+}
+
+func readQuery(params benchmarkCommandParams, args []string) (string, error) {
+	var query string
+	if params.stdin {
+		bs, err := ioutil.ReadAll(os.Stdin)
+		if err != nil {
+			return "", err
+		}
+		query = string(bs)
+	} else {
+		query = args[0]
+	}
+	return query, nil
+}
+
 func renderBenchmarkResult(params benchmarkCommandParams, br testing.BenchmarkResult, w io.Writer) {
 	switch params.outputFormat.String() {
 	case evalJSONOutput:
@@ -333,3 +617,26 @@ func prettyFormatFloat(x float64) string {
 	}
 	return fmt.Sprintf(format, x)
 }
+
+func reportMetrics(b *testing.B, m map[string]interface{}) {
+	// For each histogram add their values to the benchmark results.
+	// Note: If there are many metrics this gets super verbose.
+	for histName, metric := range m {
+		histValues, ok := metric.(map[string]interface{})
+		if !ok {
+			continue
+		}
+		for metricName, rawValue := range histValues {
+			unit := fmt.Sprintf("%s_%s", histName, metricName)
+
+			// Only support histogram metrics that are a float64 or int64,
+			// this covers the stock implementation in metrics.Metrics
+			switch v := rawValue.(type) {
+			case int64:
+				b.ReportMetric(float64(v), unit)
+			case float64:
+				b.ReportMetric(v, unit)
+			}
+		}
+	}
+}

cmd/bench_test.go

@@ -58,6 +58,158 @@ func TestRunBenchmark(t *testing.T) {
 	}
 }
 
+func TestRunBenchmarkE2E(t *testing.T) {
+	params := testBenchParams()
+	params.e2e = true
+
+	args := []string{"1 + 1"}
+	var buf bytes.Buffer
+
+	rc, err := benchMain(args, params, &buf, &goBenchRunner{})
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+
+	if rc != 0 {
+		t.Fatalf("Unexpected return code %d, expected 0", rc)
+	}
+
+	// Expect a json serialized benchmark result with histogram fields
+	var br testing.BenchmarkResult
+	err = util.UnmarshalJSON(buf.Bytes(), &br)
+	if err != nil {
+		t.Fatalf("Unexpected error unmarshalling output: %s", err)
+	}
+
+	if br.N == 0 || br.T == 0 || br.MemAllocs == 0 || br.MemBytes == 0 {
+		t.Fatalf("Expected benchmark results to be non-zero, got: %+v", br)
+	}
+
+	if _, ok := br.Extra["histogram_timer_rego_query_eval_ns_count"]; !ok {
+		t.Fatalf("Expected benchmark results to contain 'histogram_timer_rego_query_eval_ns_count', got: %+v", br)
+	}
+
+	if float64(br.N) != br.Extra["histogram_timer_rego_query_eval_ns_count"] {
+		t.Fatalf("Expected 'histogram_timer_rego_query_eval_ns_count' to be equal to N")
+	}
+
+	if _, ok := br.Extra["histogram_timer_server_handler_ns_count"]; !ok {
+		t.Fatalf("Expected benchmark results to contain 'histogram_timer_server_handler_ns_count', got: %+v", br)
+	}
+
+	if float64(br.N) != br.Extra["histogram_timer_server_handler_ns_count"] {
+		t.Fatalf("Expected 'histogram_timer_server_handler_ns_count' to be equal to N")
+	}
+}
+
+func TestRunBenchmarkFailFastE2E(t *testing.T) {
+	params := testBenchParams()
+	params.fail = true // configured to fail on undefined results
+	params.e2e = true
+
+	args := []string{"a := 1; a > 2"}
+	var buf bytes.Buffer
+
+	rc, err := benchMain(args, params, &buf, &goBenchRunner{})
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+
+	if rc != 1 {
+		t.Fatalf("Unexpected return code %d, expected 1", rc)
+	}
+
+	// Expect a json serialized benchmark result with histogram fields
+	var pr presentation.Output
+	err = util.UnmarshalJSON(buf.Bytes(), &pr)
+	if err != nil {
+		t.Fatalf("Unexpected error unmarshalling output: %s", err)
+	}
+
+	if len(pr.Errors) != 1 {
+		t.Fatalf("Expected 1 error in result, got:\n\n%s\n", buf.String())
+	}
+}
+
+func TestBenchPartialE2E(t *testing.T) {
+	params := testBenchParams()
+	params.partial = true
+	params.fail = true
+	params.e2e = true
+	params.unknowns = []string{"input"}
+	args := []string{"input.x > 0"}
+	var buf bytes.Buffer
+
+	rc, err := benchMain(args, params, &buf, &goBenchRunner{})
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+
+	if rc != 0 {
+		t.Fatalf("Unexpected return code %d, expected 0", rc)
+	}
+
+	var br testing.BenchmarkResult
+	err = util.UnmarshalJSON(buf.Bytes(), &br)
+	if err != nil {
+		t.Fatalf("Unexpected error unmarshalling output: %s", err)
+	}
+
+	if br.N == 0 || br.T == 0 || br.MemAllocs == 0 || br.MemBytes == 0 {
+		t.Fatalf("Expected benchmark results to be non-zero, got: %+v", br)
+	}
+
+	if _, ok := br.Extra["histogram_timer_rego_partial_eval_ns_count"]; !ok {
+		t.Fatalf("Expected benchmark results to contain 'histogram_timer_rego_partial_eval_ns_count', got: %+v", br)
+	}
+
+	if float64(br.N) != br.Extra["histogram_timer_rego_partial_eval_ns_count"] {
+		t.Fatalf("Expected 'histogram_timer_rego_partial_eval_ns_count' to be equal to N")
+	}
+
+	if _, ok := br.Extra["histogram_timer_server_handler_ns_count"]; !ok {
+		t.Fatalf("Expected benchmark results to contain 'histogram_timer_server_handler_ns_count', got: %+v", br)
+	}
+
+	if float64(br.N) != br.Extra["histogram_timer_server_handler_ns_count"] {
+		t.Fatalf("Expected 'histogram_timer_server_handler_ns_count' to be equal to N")
+	}
+}
+
+func TestRunBenchmarkPartialFailFastE2E(t *testing.T) {
+	params := testBenchParams()
+	params.partial = true
+	params.unknowns = []string{}
+	params.fail = true
+	params.e2e = true
+	args := []string{"1 == 2"}
+	var buf bytes.Buffer
+
+	rc, err := benchMain(args, params, &buf, &goBenchRunner{})
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+
+	if rc != 1 {
+		t.Fatalf("Unexpected return code %d, expected 1", rc)
+	}
+
+	actual := buf.String()
+	expected := `{
+  "errors": [
+    {
+      "message": "undefined result"
+    }
+  ]
+}
+`
+
+	if actual != expected {
+		t.Fatalf("\nExpected:\n%s\n\nGot:\n%s\n", expected, actual)
+	}
+
+}
+
 func TestRunBenchmarkFailFast(t *testing.T) {
 	params := testBenchParams()
 	params.fail = true // configured to fail on undefined results
@@ -290,33 +442,79 @@ func TestBenchMainInvalidInputFile(t *testing.T) {
 	})
 }
 
+func TestBenchMainWithJSONInputFileE2E(t *testing.T) {
+	params := testBenchParams()
+	params.e2e = true
+	files := map[string]string{
+		"/input.json": `{"x": 42}`,
+	}
+	args := []string{"input.x == 42"}
+	test.WithTempFS(files, func(path string) {
+		params.inputPath = filepath.Join(path, "input.json")
+
+		var buf bytes.Buffer
+
+		rc, err := benchMain(args, params, &buf, &goBenchRunner{})
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		if rc != 0 {
+			t.Fatalf("Unexpected return code %d, expected 0", rc)
+		}
+	})
+}
+
+func TestBenchMainWithYAMLInputFileE2E(t *testing.T) {
+	params := testBenchParams()
+	params.e2e = true
+	files := map[string]string{
+		"/input.yaml": `x: 42`,
+	}
+	args := []string{"input.x == 42"}
+	test.WithTempFS(files, func(path string) {
+		params.inputPath = filepath.Join(path, "input.yaml")
+
+		var buf bytes.Buffer
+
+		rc, err := benchMain(args, params, &buf, &goBenchRunner{})
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		if rc != 0 {
+			t.Fatalf("Unexpected return code %d, expected 0", rc)
+		}
+	})
+}
+
+func TestBenchMainInvalidInputFileE2E(t *testing.T) {
+	params := testBenchParams()
+	params.e2e = true
+	files := map[string]string{
+		"/input.yaml": `x: 42`,
+	}
+	args := []string{"1+1"}
+	test.WithTempFS(files, func(path string) {
+		params.inputPath = filepath.Join(path, "definitely/not/input.yaml")
+
+		var buf bytes.Buffer
+
+		rc, err := benchMain(args, params, &buf, &goBenchRunner{})
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		if rc != 1 {
+			t.Fatalf("Unexpected return code %d, expected 1", rc)
+		}
+	})
+}
+
 func TestBenchMainWithBundleData(t *testing.T) {
 	params := testBenchParams()
 
-	mod := `package a.b
-
-	x {
-	   data.a.b.c == 42
-	}
-	`
-
-	b := bundle.Bundle{
-		Manifest: bundle.Manifest{},
-		Data: map[string]interface{}{
-			"a": map[string]interface{}{
-				"b": map[string]interface{}{
-					"c": 42,
-				},
-			},
-		},
-		Modules: []bundle.ModuleFile{
-			{
-				Path:   "/a/b/policy.rego",
-				Raw:    []byte(mod),
-				Parsed: ast.MustParseModule(mod),
-			},
-		},
-	}
+	b := testBundle()
 
 	files := map[string]string{
 		"bundle.tar.gz": "",
@@ -343,7 +541,101 @@ func TestBenchMainWithBundleData(t *testing.T) {
 
 		validateBenchMainPrep(t, args, params)
 	})
+}
 
+func TestBenchMainWithBundleDataE2E(t *testing.T) {
+	params := testBenchParams()
+	params.e2e = true
+
+	b := testBundle()
+
+	files := map[string]string{
+		"bundle.tar.gz": "",
+	}
+
+	test.WithTempFS(files, func(path string) {
+		bundlePath := filepath.Join(path, "bundle.tar.gz")
+		f, err := os.OpenFile(bundlePath, os.O_WRONLY, os.ModePerm)
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		err = bundle.Write(f, b)
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		err = params.bundlePaths.Set(bundlePath)
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		args := []string{"data.a.b.x"}
+
+		var buf bytes.Buffer
+
+		rc, err := benchMain(args, params, &buf, &goBenchRunner{})
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		if rc != 0 {
+			t.Fatalf("Unexpected return code %d, expected 0", rc)
+		}
+	})
+}
+
+func TestBenchMainWithDataE2E(t *testing.T) {
+	params := testBenchParams()
+	params.e2e = true
+
+	mod := `package a.b
+
+	x {
+	   data.a.b.c == 42
+	}
+	`
+
+	files := map[string]string{
+		"p.rego": mod,
+	}
+
+	test.WithTempFS(files, func(path string) {
+		err := params.dataPaths.Set(filepath.Join(path, "p.rego"))
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		args := []string{"data.a.b.x"}
+
+		var buf bytes.Buffer
+
+		rc, err := benchMain(args, params, &buf, &goBenchRunner{})
+		if err != nil {
+			t.Fatalf("Unexpected error: %s", err)
+		}
+
+		if rc != 0 {
+			t.Fatalf("Unexpected return code %d, expected 0", rc)
+		}
+	})
+}
+
+func TestBenchMainBadQueryE2E(t *testing.T) {
+	params := testBenchParams()
+	params.e2e = true
+	args := []string{"foo.bar"}
+
+	var buf bytes.Buffer
+
+	rc, err := benchMain(args, params, &buf, &goBenchRunner{})
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+
+	if rc != 1 {
+		t.Fatalf("Unexpected return code %d, expected 1", rc)
+	}
 }
 
 func TestRenderBenchmarkResultJSONOutput(t *testing.T) {
@@ -600,3 +892,30 @@ func fakeBenchResults() testing.BenchmarkResult {
 		},
 	}
 }
+
+func testBundle() bundle.Bundle {
+	mod := `package a.b
+
+	x {
+	   data.a.b.c == 42
+	}
+	`
+
+	return bundle.Bundle{
+		Manifest: bundle.Manifest{},
+		Data: map[string]interface{}{
+			"a": map[string]interface{}{
+				"b": map[string]interface{}{
+					"c": 42,
+				},
+			},
+		},
+		Modules: []bundle.ModuleFile{
+			{
+				Path:   "/a/b/policy.rego",
+				Raw:    []byte(mod),
+				Parsed: ast.MustParseModule(mod),
+			},
+		},
+	}
+}

cmd/flags.go

@@ -153,6 +153,10 @@ func addStrictFlag(fs *pflag.FlagSet, strict *bool, value bool) {
 	fs.BoolVarP(strict, "strict", "S", value, "enable compiler strict mode")
 }
 
+func addE2EFlag(fs *pflag.FlagSet, e2e *bool, value bool) {
+	fs.BoolVar(e2e, "e2e", value, "run benchmarks against a running OPA server")
+}
+
 const (
 	explainModeOff   = "off"
 	explainModeFull  = "full"

server/server.go

@@ -2156,6 +2156,7 @@ func (s *Server) v1QueryGet(w http.ResponseWriter, r *http.Request) {
 
 func (s *Server) v1QueryPost(w http.ResponseWriter, r *http.Request) {
 	m := metrics.New()
+	m.Timer(metrics.ServerHandler).Start()
 
 	decisionID := s.generateDecisionID()
 	ctx := r.Context()
@@ -2221,6 +2222,12 @@ func (s *Server) v1QueryPost(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	m.Timer(metrics.ServerHandler).Stop()
+
+	if includeMetrics || includeInstrumentation {
+		results.Metrics = m.All()
+	}
+
 	writer.JSON(w, http.StatusOK, results, pretty)
 }