ewilliams28/releases
1
0
Fork 0
mirror of https://github.com/open-policy-agent/opa.git synced 2026-08-12 19:32:48 -06:00
Code Issues Packages Projects Releases Wiki Activity

cmd: Add support for benchmarking partial evaluation queries

Browse Source 
This change allows users to get an accurate estimate of partial eval
performance on their policies.

Signed-off-by: Torin Sandall <torinsandall@gmail.com>
This commit is contained in:
Torin Sandall
2021-01-06 16:18:09 -05:00
parent 2c3060466d
commit 740b33bd13
4 changed files with 72 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/bench.go
+41 -17
View File
@@ -6,7 +6,6 @@ package cmd
import (
"context"
"errors"
"fmt"
"io"
"math"
@@ -18,7 +17,6 @@ import (
"github.com/spf13/cobra"
"github.com/open-policy-agent/opa/internal/presentation"
"github.com/open-policy-agent/opa/metrics"
"github.com/open-policy-agent/opa/rego"
"github.com/open-policy-agent/opa/util"
@@ -64,7 +62,6 @@ Example with bundle and input data:
opa bench -b ./policy-bundle -i input.json 'data.authz.allow'
To enable more detailed analysis use the --metrics and --benchmem flags.
The optional "gobench" output format conforms to the Go Benchmark Data Format.
@@ -79,6 +76,8 @@ The optional "gobench" output format conforms to the Go Benchmark Data Format.
}
// Sub-set of the standard `opa eval ..` flags
addPartialFlag(benchCommand.Flags(), &params.partial, false)
addUnknownsFlag(benchCommand.Flags(), &params.unknowns, []string{"input"})
addFailFlag(benchCommand.Flags(), &params.fail, true)
addDataFlag(benchCommand.Flags(), &params.dataPaths)
addBundleFlag(benchCommand.Flags(), &params.bundlePaths)
@@ -99,7 +98,7 @@ The optional "gobench" output format conforms to the Go Benchmark Data Format.
}
type benchRunner interface {
run(ctx context.Context, ectx *evalContext, pq rego.PreparedEvalQuery, params benchmarkCommandParams) (testing.BenchmarkResult, error)
run(ctx context.Context, ectx *evalContext, params benchmarkCommandParams, f func(context.Context, ...rego.EvalOption) error) (testing.BenchmarkResult, error)
}
func benchMain(args []string, params benchmarkCommandParams, w io.Writer, r benchRunner) int {
@@ -110,17 +109,47 @@ func benchMain(args []string, params benchmarkCommandParams, w io.Writer, r benc
}
ctx := context.Background()
var benchFunc func(context.Context, ...rego.EvalOption) error
// Take the eval context and prepare anything else we possible can before benchmarking the evaluation
pq, err := ectx.r.PrepareForEval(ctx)
if err != nil {
renderBenchmarkError(params, err, w)
return 1
if !params.partial {
// Take the eval context and prepare anything else we possible can before benchmarking the evaluation
pq, err := ectx.r.PrepareForEval(ctx)
if err != nil {
renderBenchmarkError(params, err, w)
return 1
}
benchFunc = func(ctx context.Context, opts ...rego.EvalOption) error {
result, err := pq.Eval(ctx, opts...)
if err != nil {
return err
} else if len(result) == 0 && params.fail {
return fmt.Errorf("undefined result")
}
return nil
}
} else {
// As with normal evaluation, prepare as much as possible up front.
pq, err := ectx.r.PrepareForPartial(ctx)
if err != nil {
renderBenchmarkError(params, err, w)
return 1
}
benchFunc = func(ctx context.Context, opts ...rego.EvalOption) error {
result, err := pq.Partial(ctx, opts...)
if err != nil {
return err
} else if len(result.Queries) == 0 && params.fail {
return fmt.Errorf("undefined result")
}
return nil
}
}
// Run the benchmark as many times as specified, re-use the prepared objects for each
for i := 0; i < params.count; i++ {
br, err := r.run(ctx, ectx, pq, params)
br, err := r.run(ctx, ectx, params, benchFunc)
if err != nil {
renderBenchmarkError(params, err, w)
return 1
@@ -134,7 +163,7 @@ func benchMain(args []string, params benchmarkCommandParams, w io.Writer, r benc
type goBenchRunner struct {
}
func (r *goBenchRunner) run(ctx context.Context, ectx *evalContext, pq rego.PreparedEvalQuery, params benchmarkCommandParams) (testing.BenchmarkResult, error) {
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
if params.metrics {
@@ -167,7 +196,7 @@ func (r *goBenchRunner) run(ctx context.Context, ectx *evalContext, pq rego.Prep
b.StartTimer()
// Perform the evaluation
rs, err := pq.Eval(ctx, ectx.evalArgs...)
err := f(ctx, ectx.evalArgs...)
// Stop the timer while processing the results
b.StopTimer()
@@ -176,11 +205,6 @@ func (r *goBenchRunner) run(ctx context.Context, ectx *evalContext, pq rego.Prep
b.FailNow()
}
if len(rs) == 0 && params.fail {
benchErr = errors.New("undefined result")
b.FailNow()
}
// Add metrics for that evaluation into the top level histogram
if params.metrics {
for name, metric := range m.All() {
cmd/bench_test.go
+21 -7
View File
@@ -83,16 +83,30 @@ func TestRunBenchmarkFailFast(t *testing.T) {
// mockBenchRunner lets us test the bench CLI operations without having to wait ~10 seconds
// while the actual benchmark runner does its thing.
type mockBenchRunner struct {
onRun func(ctx context.Context, ectx *evalContext, pq rego.PreparedEvalQuery, params benchmarkCommandParams) (testing.BenchmarkResult, error)
onRun func(ctx context.Context, ectx *evalContext, params benchmarkCommandParams, f func(context.Context, ...rego.EvalOption) error) (testing.BenchmarkResult, error)
}
func (r *mockBenchRunner) run(ctx context.Context, ectx *evalContext, pq rego.PreparedEvalQuery, params benchmarkCommandParams) (testing.BenchmarkResult, error) {
func (r *mockBenchRunner) run(ctx context.Context, ectx *evalContext, params benchmarkCommandParams, f func(context.Context, ...rego.EvalOption) error) (testing.BenchmarkResult, error) {
if r.onRun != nil {
return r.onRun(ctx, ectx, pq, params)
return r.onRun(ctx, ectx, params, f)
}
return testing.BenchmarkResult{}, nil
}
func TestBenchPartial(t *testing.T) {
params := testBenchParams()
params.partial = true
params.fail = true
args := []string{"input=1"}
var buf bytes.Buffer
rc := benchMain(args, params, &buf, &mockBenchRunner{})
if rc != 0 {
t.Fatalf("Unexpected return code %d, expected 0", rc)
}
}
func TestBenchMainErrPreparing(t *testing.T) {
params := testBenchParams()
args := []string{"???"} // query compile error
@@ -111,7 +125,7 @@ func TestBenchMainErrRunningBenchmark(t *testing.T) {
var buf bytes.Buffer
mockRunner := &mockBenchRunner{}
mockRunner.onRun = func(ctx context.Context, ectx *evalContext, pq rego.PreparedEvalQuery, params benchmarkCommandParams) (testing.BenchmarkResult, error) {
mockRunner.onRun = func(ctx context.Context, ectx *evalContext, params benchmarkCommandParams, f func(context.Context, ...rego.EvalOption) error) (testing.BenchmarkResult, error) {
return testing.BenchmarkResult{}, errors.New("error error error")
}
@@ -131,7 +145,7 @@ func TestBenchMainWithCount(t *testing.T) {
params.count = 25
actualCount := 0
mockRunner.onRun = func(ctx context.Context, ectx *evalContext, pq rego.PreparedEvalQuery, params benchmarkCommandParams) (testing.BenchmarkResult, error) {
mockRunner.onRun = func(ctx context.Context, ectx *evalContext, params benchmarkCommandParams, f func(context.Context, ...rego.EvalOption) error) (testing.BenchmarkResult, error) {
actualCount++
return testing.BenchmarkResult{}, nil
}
@@ -156,7 +170,7 @@ func TestBenchMainWithNegativeCount(t *testing.T) {
params.count = -1
actualCount := 0
mockRunner.onRun = func(ctx context.Context, ectx *evalContext, pq rego.PreparedEvalQuery, params benchmarkCommandParams) (testing.BenchmarkResult, error) {
mockRunner.onRun = func(ctx context.Context, ectx *evalContext, params benchmarkCommandParams, f func(context.Context, ...rego.EvalOption) error) (testing.BenchmarkResult, error) {
actualCount++
return testing.BenchmarkResult{}, nil
}
@@ -178,7 +192,7 @@ func validateBenchMainPrep(t *testing.T, args []string, params benchmarkCommandP
mockRunner := &mockBenchRunner{}
mockRunner.onRun = func(ctx context.Context, ectx *evalContext, pq rego.PreparedEvalQuery, params benchmarkCommandParams) (testing.BenchmarkResult, error) {
mockRunner.onRun = func(ctx context.Context, ectx *evalContext, params benchmarkCommandParams, f func(context.Context, ...rego.EvalOption) error) (testing.BenchmarkResult, error) {
// cheat and use the ectx to evalute the query to ensure the input setup on it was valid
pq, err := ectx.r.PrepareForEval(ctx)
cmd/eval.go
+2 -2
View File
@@ -218,8 +218,6 @@ Set the output format with the --format flag.
// Eval specific flags
evalCommand.Flags().BoolVarP(&params.coverage, "coverage", "", false, "report coverage")
evalCommand.Flags().BoolVarP(&params.partial, "partial", "p", false, "perform partial evaluation")
evalCommand.Flags().StringArrayVarP(&params.unknowns, "unknowns", "u", []string{"input"}, "set paths to treat as unknown during partial evaluation")
evalCommand.Flags().StringArrayVarP(&params.disableInlining, "disable-inlining", "", []string{}, "set paths of documents to exclude from inlining")
evalCommand.Flags().BoolVarP(&params.shallowInlining, "shallow-inlining", "", false, "disable inlining of rules that depend on unknowns")
evalCommand.Flags().BoolVar(&params.disableIndexing, "disable-indexing", false, "disable indexing optimizations")
@@ -232,6 +230,8 @@ Set the output format with the --format flag.
evalCommand.Flags().BoolVarP(&params.failDefined, "fail-defined", "", false, "exits with non-zero exit code on defined/non-empty result and errors")
// Shared flags
addPartialFlag(evalCommand.Flags(), &params.partial, false)
addUnknownsFlag(evalCommand.Flags(), &params.unknowns, []string{"input"})
addFailFlag(evalCommand.Flags(), &params.fail, false)
addDataFlag(evalCommand.Flags(), &params.dataPaths)
addBundleFlag(evalCommand.Flags(), &params.bundlePaths)
cmd/flags.go
+8
View File
@@ -122,6 +122,14 @@ func addCapabilitiesFlag(fs *pflag.FlagSet, f *capabilitiesFlag) {
fs.VarP(f, "capabilities", "", "set capabilities.json file path")
}
func addPartialFlag(fs *pflag.FlagSet, partial *bool, value bool) {
fs.BoolVarP(partial, "partial", "p", value, "perform partial evaluation")
}
func addUnknownsFlag(fs *pflag.FlagSet, unknowns *[]string, value []string) {
fs.StringArrayVarP(unknowns, "unknowns", "u", value, "set paths to treat as unknown during partial evaluation")
}
const (
explainModeOff = "off"
explainModeFull = "full"