From 4e66158fb73102d9c2411ddf39c247e9506c6cf2 Mon Sep 17 00:00:00 2001 From: Edward Paget Date: Fri, 6 Jan 2023 04:13:10 -0600 Subject: [PATCH] topdown: cache undefined rule evaluations (#5523) With this change, `undefined` outcomes of complete rule evaluations are now also cached. Previously, only defined results had been cached, and empty partial sets/objects. In the case of partial rules with string keys, the introduction of ref heads changed how they had been evaluated: Before, they had been evaluated as partial sets, and thus got cached when empty. After, they had been evaluated as complete rules (with ref heads), and if they were undefined, they had _not_ been cached. This caused a performance regression. Fixes #593. Signed-off-by: Edward Paget --- rego/rego_bench_test.go | 70 +++++++++++++++++++++++++++++++++++++ topdown/cache.go | 30 ++++++++++++---- topdown/cache_bench_test.go | 2 +- topdown/cache_test.go | 16 +++++++-- topdown/eval.go | 20 +++++++---- util/test/benchmark.go | 67 +++++++++++++++++++++++++++++++++++ 6 files changed, 190 insertions(+), 15 deletions(-) create mode 100644 rego/rego_bench_test.go 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 {