diff --git a/rego/rego_bench_test.go b/rego/rego_bench_test.go new file mode 100644 index 0000000000..b64e02f9b7 --- /dev/null +++ b/rego/rego_bench_test.go @@ -0,0 +1,70 @@ +package rego + +import ( + "context" + "fmt" + "testing" + + "github.com/open-policy-agent/opa/ast" + "github.com/open-policy-agent/opa/internal/runtime" + inmem "github.com/open-policy-agent/opa/storage/inmem/test" + "github.com/open-policy-agent/opa/util/test" +) + +func BenchmarkPartialObjectRuleCrossModule(b *testing.B) { + ctx := context.Background() + sizes := []int{10, 100, 1000} + + for _, n := range sizes { + b.Run(fmt.Sprint(n), func(b *testing.B) { + store := inmem.NewFromObject(map[string]interface{}{}) + mods := test.PartialObjectBenchmarkCrossModule(n) + query := "data.test.foo" + + input := make(map[string]interface{}) + for idx := 0; idx <= 3; idx++ { + input[fmt.Sprintf("test_input_%d", idx)] = "test_input_10" + } + inputAST, err := ast.InterfaceToValue(input) + if err != nil { + b.Fatal(err) + } + + compiler := ast.MustCompileModules(map[string]string{ + "test/foo.rego": mods[0], + "test/bar.rego": mods[1], + "test/baz.rego": mods[2], + }) + info, err := runtime.Term(runtime.Params{}) + if err != nil { + b.Fatal(err) + } + + pq, err := New( + Query(query), + Compiler(compiler), + Store(store), + Runtime(info), + ).PrepareForEval(ctx) + + if err != nil { + b.Fatal(err) + } + + b.ResetTimer() + + for i := 0; i < b.N; i++ { + _, err = pq.Eval( + ctx, + EvalParsedInput(inputAST), + EvalRuleIndexing(true), + EvalEarlyExit(true), + ) + + if err != nil { + b.Fatal(err) + } + } + }) + } +} diff --git a/topdown/cache.go b/topdown/cache.go index 1b7c455eec..710efee475 100644 --- a/topdown/cache.go +++ b/topdown/cache.go @@ -14,8 +14,9 @@ type virtualCache struct { } type virtualCacheElem struct { - value *ast.Term - children *util.HashMap + value *ast.Term + children *util.HashMap + undefined bool } func newVirtualCache() *virtualCache { @@ -32,18 +33,31 @@ func (c *virtualCache) Pop() { c.stack = c.stack[:len(c.stack)-1] } -func (c *virtualCache) Get(ref ast.Ref) *ast.Term { +// Returns the resolved value of the AST term and a flag indicating if the value +// should be interpretted as undefined: +// +// nil, true indicates the ref is undefined +// ast.Term, false indicates the ref is defined +// nil, false indicates the ref has not been cached +// ast.Term, true is impossible +func (c *virtualCache) Get(ref ast.Ref) (*ast.Term, bool) { node := c.stack[len(c.stack)-1] for i := 0; i < len(ref); i++ { x, ok := node.children.Get(ref[i]) if !ok { - return nil + return nil, false } node = x.(*virtualCacheElem) } - return node.value + if node.undefined { + return nil, true + } + + return node.value, false } +// If value is a nil pointer, set the 'undefined' flag on the cache element to +// indicate that the Ref has resolved to undefined. func (c *virtualCache) Put(ref ast.Ref, value *ast.Term) { node := c.stack[len(c.stack)-1] for i := 0; i < len(ref); i++ { @@ -56,7 +70,11 @@ func (c *virtualCache) Put(ref ast.Ref, value *ast.Term) { node = next } } - node.value = value + if value != nil { + node.value = value + } else { + node.undefined = true + } } func newVirtualCacheElem() *virtualCacheElem { diff --git a/topdown/cache_bench_test.go b/topdown/cache_bench_test.go index ea29298f37..2136490257 100644 --- a/topdown/cache_bench_test.go +++ b/topdown/cache_bench_test.go @@ -39,7 +39,7 @@ func BenchmarkVirtualCache(b *testing.B) { for i := 0; i < b.N; i++ { idx := i % max cache.Put(keys[idx], values[idx]) - result := cache.Get(keys[idx]) + result, _ := cache.Get(keys[idx]) if !result.Equal(values[idx]) { b.Fatal("expected equal") } diff --git a/topdown/cache_test.go b/topdown/cache_test.go index 7e06f2c687..965c60f77c 100644 --- a/topdown/cache_test.go +++ b/topdown/cache_test.go @@ -14,7 +14,7 @@ func TestVirtualCacheCompositeKey(t *testing.T) { cache := newVirtualCache() ref := ast.MustParseRef("data.x.y[[1]].z") cache.Put(ref, ast.BooleanTerm(true)) - result := cache.Get(ref) + result, _ := cache.Get(ref) if !result.Equal(ast.BooleanTerm(true)) { t.Fatalf("Expected true but got %v", result) } @@ -25,12 +25,24 @@ func TestVirtualCacheInvalidate(t *testing.T) { cache.Push() cache.Put(ast.MustParseRef("data.x.p"), ast.BooleanTerm(true)) cache.Pop() - result := cache.Get(ast.MustParseRef("data.x.p")) + result, _ := cache.Get(ast.MustParseRef("data.x.p")) if result != nil { t.Fatal("Expected nil result but got:", result) } } +func TestSetAndRetriveUndefined(t *testing.T) { + cache := newVirtualCache() + cache.Put(ast.MustParseRef("data.foo.bar"), nil) + result, undefined := cache.Get(ast.MustParseRef("data.foo.bar")) + if result != nil { + t.Fatal("Expected nil result but got:", result) + } + if !undefined { + t.Fatal("Expected 'undefined' flag to be false got true") + } +} + func TestBaseCacheGetExactMatch(t *testing.T) { cache := newBaseCache() cache.Put(ast.MustParseRef("data.x.foo"), ast.StringTerm("bar").Value) diff --git a/topdown/eval.go b/topdown/eval.go index 56bbb2c7d4..1072c95c06 100644 --- a/topdown/eval.go +++ b/topdown/eval.go @@ -330,7 +330,6 @@ func (e *eval) evalExpr(iter evalIterator) error { } return nil } - expr := e.query[e.index] e.traceEval(expr) @@ -1875,7 +1874,7 @@ func (e evalFunc) evalCache(argCount int, iter unifyIterator) (ast.Ref, bool, er cacheKey[i] = e.e.bindings.Plug(e.terms[i]) } - cached := e.e.virtualCache.Get(cacheKey) + cached, _ := e.e.virtualCache.Get(cacheKey) if cached != nil { e.e.instr.counterIncr(evalOpVirtualCacheHit) if argCount == len(e.terms)-1 { // f(x) @@ -2391,7 +2390,7 @@ func (e evalVirtualPartial) evalEachRule(iter unifyIterator, unknown bool) error func (e evalVirtualPartial) evalAllRules(iter unifyIterator, rules []*ast.Rule) error { cacheKey := e.plugged[:e.pos+1] - result := e.e.virtualCache.Get(cacheKey) + result, _ := e.e.virtualCache.Get(cacheKey) if result != nil { e.e.instr.counterIncr(evalOpVirtualCacheHit) return e.e.biunify(result, e.rterm, e.bindings, e.rbindings, iter) @@ -2653,7 +2652,7 @@ func (e evalVirtualPartial) evalCache(iter unifyIterator) (evalVirtualPartialCac return hint, nil } - if cached := e.e.virtualCache.Get(e.plugged[:e.pos+1]); cached != nil { // have full extent cached + if cached, _ := e.e.virtualCache.Get(e.plugged[:e.pos+1]); cached != nil { // have full extent cached e.e.instr.counterIncr(evalOpVirtualCacheHit) hint.hit = true return hint, e.evalTerm(iter, e.pos+1, cached, e.bindings) @@ -2664,7 +2663,7 @@ func (e evalVirtualPartial) evalCache(iter unifyIterator) (evalVirtualPartialCac if plugged.IsGround() { hint.key = append(e.plugged[:e.pos+1], plugged) - if cached := e.e.virtualCache.Get(hint.key); cached != nil { + if cached, _ := e.e.virtualCache.Get(hint.key); cached != nil { e.e.instr.counterIncr(evalOpVirtualCacheHit) hint.hit = true return hint, e.evalTerm(iter, e.pos+2, cached, e.bindings) @@ -2752,7 +2751,12 @@ func (e evalVirtualComplete) eval(iter unifyIterator) error { } func (e evalVirtualComplete) evalValue(iter unifyIterator, findOne bool) error { - cached := e.e.virtualCache.Get(e.plugged[:e.pos+1]) + cached, undefined := e.e.virtualCache.Get(e.plugged[:e.pos+1]) + if undefined { + e.e.instr.counterIncr(evalOpVirtualCacheHit) + return nil + } + if cached != nil { e.e.instr.counterIncr(evalOpVirtualCacheHit) return e.evalTerm(iter, cached, e.bindings) @@ -2788,6 +2792,10 @@ func (e evalVirtualComplete) evalValue(iter unifyIterator, findOne bool) error { return err } + if prev == nil { + e.e.virtualCache.Put(e.plugged[:e.pos+1], nil) + } + return nil } diff --git a/util/test/benchmark.go b/util/test/benchmark.go index c232f7562f..1af3eaa2d0 100644 --- a/util/test/benchmark.go +++ b/util/test/benchmark.go @@ -11,6 +11,73 @@ import ( "text/template" ) +// PartialObjectBenchmarkCrossModule returns a module with n "bench_test_" prefixed rules +// that each refer to another "cond_bench_" prefixed rule +func PartialObjectBenchmarkCrossModule(n int) []string { + fooMod := `package test.foo + import data.test.bar + import data.test.baz + + output[key] := value { + value := bar[key] + startswith("bench_test_", key) + }` + barMod := "package test.bar\n" + barMod += ` + cond_bench_0 { + contains(lower(input.test_input_0), lower("input_01")) + } + cond_bench_1 { + contains(lower(input.test_input_1), lower("input")) + } + cond_bench_2 { + contains(lower(input.test_input_2), lower("input_10")) + } + bench_test_out_result := load_tests(test_collector) + + load_tests(in) := out { + out := in + } + ` + + bazMod := "package test.baz\nimport data.test.bar\n" + ruleBuilder := "" + + for idx := 1; idx <= n; idx++ { + barMod += fmt.Sprintf(` + bench_test_%[1]d := result { + input.bench_test_collector_mambo_number_%[3]d + result := input.bench_test_collector_mambo_number_%[3]d + } else := result { + is_null(bench_test_out_result.mambo_number_%[3]d.error) + result := bench_test_out_result.mambo_number_%[3]d.result + } + + test_collector["mambo_number_%[3]d"] := result { + cond_bench_%[2]d + not %[3]d == 2 + not %[3]d == 3 + not input.bench_test_collector_mambo_number_%[3]d + result := { "result": %[3]d, "error": null } + } + `, idx, idx%3, idx%5) + ruleBuilder += fmt.Sprintf(" bar.bench_test_%[1]d == %[1]d\n", idx) + if idx%10 == 0 { + bazMod += fmt.Sprintf(`rule_%d { + %s + }`, idx, ruleBuilder) + fooMod += fmt.Sprintf(` + final_decision = "allow" { + baz.rule_%d + } + `, idx) + ruleBuilder = "" + } + } + + return []string{fooMod, barMod, bazMod} +} + // ArrayIterationBenchmarkModule returns a module that iterates an array // with `n` elements func ArrayIterationBenchmarkModule(n int) string {