diff --git a/ast/compile.go b/ast/compile.go index ad1a9208d1..f2c8c35c2c 100644 --- a/ast/compile.go +++ b/ast/compile.go @@ -90,13 +90,14 @@ type Compiler struct { metricName string f func() } - maxErrs int - sorted []string // list of sorted module names - pathExists func([]string) (bool, error) - after map[string][]CompilerStageDefinition - metrics metrics.Metrics - builtins map[string]*Builtin - unsafeBuiltinsMap map[string]struct{} + maxErrs int + sorted []string // list of sorted module names + pathExists func([]string) (bool, error) + after map[string][]CompilerStageDefinition + metrics metrics.Metrics + builtins map[string]*Builtin + unsafeBuiltinsMap map[string]struct{} + comprehensionIndices map[*Term]*ComprehensionIndex } // CompilerStage defines the interface for stages in the compiler. @@ -187,6 +188,10 @@ type QueryCompiler interface { // parsed query. For example, given the query "input := 1" the rewritten // query would be "__local0__ = 1". The mapping would then be {__local0__: input}. RewrittenVars() map[Var]Var + + // ComprehensionIndex returns an index data structure for the given comprehension + // term. If no index is found, returns nil. + ComprehensionIndex(term *Term) *ComprehensionIndex } // QueryCompilerStage defines the interface for stages in the query compiler. @@ -214,9 +219,10 @@ func NewCompiler() *Compiler { }, func(x util.T) int { return x.(Ref).Hash() }), - maxErrs: CompileErrorLimitDefault, - after: map[string][]CompilerStageDefinition{}, - unsafeBuiltinsMap: map[string]struct{}{}, + maxErrs: CompileErrorLimitDefault, + after: map[string][]CompilerStageDefinition{}, + unsafeBuiltinsMap: map[string]struct{}{}, + comprehensionIndices: map[*Term]*ComprehensionIndex{}, } c.ModuleTree = NewModuleTree(nil) @@ -261,6 +267,7 @@ func NewCompiler() *Compiler { {"CheckTypes", "compile_stage_check_types", c.checkTypes}, {"CheckUnsafeBuiltins", "compile_state_check_unsafe_builtins", c.checkUnsafeBuiltins}, {"BuildRuleIndices", "compile_stage_rebuild_indices", c.buildRuleIndices}, + {"BuildComprehensionIndices", "compile_stage_rebuild_comprehension_indices", c.buildComprehensionIndices}, } return c @@ -350,6 +357,13 @@ func (c *Compiler) Failed() bool { return len(c.Errors) > 0 } +// ComprehensionIndex returns a data structure specifying how to index comprehension +// results so that callers do not have to recompute the comprehension more than once. +// If no index is found, returns nil. +func (c *Compiler) ComprehensionIndex(term *Term) *ComprehensionIndex { + return c.comprehensionIndices[term] +} + // GetArity returns the number of args a function referred to by ref takes. If // ref refers to built-in function, the built-in declaration is consulted, // otherwise, the ref is used to perform a ruleset lookup. @@ -610,7 +624,13 @@ func (c *Compiler) WithModuleLoader(f ModuleLoader) *Compiler { return c } -// buildRuleIndices constructs indices for rules. +func (c *Compiler) counterAdd(name string, n uint64) { + if c.metrics == nil { + return + } + c.metrics.Counter(name).Add(n) +} + func (c *Compiler) buildRuleIndices() { c.RuleTree.DepthFirst(func(node *TreeNode) bool { @@ -628,6 +648,18 @@ func (c *Compiler) buildRuleIndices() { } +func (c *Compiler) buildComprehensionIndices() { + for _, name := range c.sorted { + WalkRules(c.Modules[name], func(r *Rule) bool { + candidates := r.Head.Args.Vars() + candidates.Update(ReservedVars) + n := buildComprehensionIndices(c.GetArity, candidates, r.Body, c.comprehensionIndices) + c.counterAdd(compileStageComprehensionIndexBuild, n) + return false + }) + } +} + // checkRecursion ensures that there are no recursive definitions, i.e., there are // no cycles in the Graph. func (c *Compiler) checkRecursion() { @@ -1250,19 +1282,21 @@ func (c *Compiler) setGraph() { } type queryCompiler struct { - compiler *Compiler - qctx *QueryContext - typeEnv *TypeEnv - rewritten map[Var]Var - after map[string][]QueryCompilerStageDefinition - unsafeBuiltins map[string]struct{} + compiler *Compiler + qctx *QueryContext + typeEnv *TypeEnv + rewritten map[Var]Var + after map[string][]QueryCompilerStageDefinition + unsafeBuiltins map[string]struct{} + comprehensionIndices map[*Term]*ComprehensionIndex } func newQueryCompiler(compiler *Compiler) QueryCompiler { qc := &queryCompiler{ - compiler: compiler, - qctx: nil, - after: map[string][]QueryCompilerStageDefinition{}, + compiler: compiler, + qctx: nil, + after: map[string][]QueryCompilerStageDefinition{}, + comprehensionIndices: map[*Term]*ComprehensionIndex{}, } return qc } @@ -1286,6 +1320,15 @@ func (qc *queryCompiler) RewrittenVars() map[Var]Var { return qc.rewritten } +func (qc *queryCompiler) ComprehensionIndex(term *Term) *ComprehensionIndex { + if result, ok := qc.comprehensionIndices[term]; ok { + return result + } else if result, ok := qc.compiler.comprehensionIndices[term]; ok { + return result + } + return nil +} + func (qc *queryCompiler) runStage(metricName string, qctx *QueryContext, query Body, s func(*QueryContext, Body) (Body, error)) (Body, error) { if qc.compiler.metrics != nil { qc.compiler.metrics.Timer(metricName).Start() @@ -1321,6 +1364,7 @@ func (qc *queryCompiler) Compile(query Body) (Body, error) { {"RewriteDynamicTerms", "query_compile_stage_rewrite_dynamic_terms", qc.rewriteDynamicTerms}, {"CheckTypes", "query_compile_stage_check_types", qc.checkTypes}, {"CheckUnsafeBuiltins", "query_compile_stage_check_unsafe_builtins", qc.checkUnsafeBuiltins}, + {"BuildComprehensionIndex", "query_compile_stage_build_comprehension_index", qc.buildComprehensionIndices}, } qctx := qc.qctx.Copy() @@ -1462,6 +1506,189 @@ func (qc *queryCompiler) rewriteWithModifiers(qctx *QueryContext, body Body) (Bo return body, nil } +func (qc *queryCompiler) buildComprehensionIndices(qctx *QueryContext, body Body) (Body, error) { + // NOTE(tsandall): The query compiler does not have a metrics object so we + // cannot record index metrics currently. + _ = buildComprehensionIndices(qc.compiler.GetArity, ReservedVars, body, qc.comprehensionIndices) + return body, nil +} + +// ComprehensionIndex specifies how the comprehension term can be indexed. The keys +// tell the evaluator what variables to use for indexing. In the future, the index +// could be expanded with more information that would allow the evaluator to index +// a larger fragment of comprehensions (e.g., by closing over variables in the outer +// query.) +type ComprehensionIndex struct { + Term *Term + Keys []*Term +} + +func (ci *ComprehensionIndex) String() string { + if ci == nil { + return "" + } + return fmt.Sprintf("", Array(ci.Keys)) +} + +func buildComprehensionIndices(arity func(Ref) int, candidates VarSet, node interface{}, result map[*Term]*ComprehensionIndex) (n uint64) { + WalkBodies(node, func(b Body) bool { + cpy := candidates.Copy() + for _, expr := range b { + if index := getComprehensionIndex(arity, cpy, expr); index != nil { + result[index.Term] = index + n++ + } + // Any variables appearing in the expressions leading up to the comprehension + // are fair-game to be used as index keys. + cpy.Update(expr.Vars(VarVisitorParams{SkipClosures: true, SkipRefCallHead: true})) + } + return false + }) + return n +} + +func getComprehensionIndex(arity func(Ref) int, candidates VarSet, expr *Expr) *ComprehensionIndex { + + // Ignore everything except = expressions. Extract + // the comprehension term from the expression. + if !expr.IsEquality() || expr.Negated || len(expr.With) > 0 { + return nil + } + + var term *Term + + lhs, rhs := expr.Operand(0), expr.Operand(1) + + if _, ok := lhs.Value.(Var); ok && IsComprehension(rhs.Value) { + term = rhs + } else if _, ok := rhs.Value.(Var); ok && IsComprehension(lhs.Value) { + term = lhs + } + + if term == nil { + return nil + } + + // Ignore comprehensions that contain expressions that close over variables + // in the outer body if those variables are not also output variables in the + // comprehension body. In other words, ignore comprehensions that we cannot + // safely evaluate without bindings from the outer body. For example: + // + // x = [1] + // [true | data.y[z] = x] # safe to evaluate w/o outer body + // [true | data.y[z] = x[0]] # NOT safe to evaluate because 'x' would be unsafe. + // + // By identifying output variables in the body we also know what to index on by + // intersecting with candidate variables from the outer query. + // + // For example: + // + // x = data.foo[_] + // _ = [y | data.bar[y] = x] # index on 'x' + // + // This query goes from O(data.foo*data.bar) to O(data.foo+data.bar). + var body Body + + switch x := term.Value.(type) { + case *ArrayComprehension: + body = x.Body + case *SetComprehension: + body = x.Body + case *ObjectComprehension: + body = x.Body + } + + outputs := outputVarsForBody(body, arity, ReservedVars) + unsafe := body.Vars(safetyCheckVarVisitorParams).Diff(outputs).Diff(ReservedVars) + if len(unsafe) > 0 { + return nil + } + + // Similarly, ignore comprehensions that contain references with output variables + // that intersect with the candidates. Indexing these comprehensions could worsen + // performance. + vis := newComprehensionIndexRegressionCheckVisitor(candidates) + vis.Walk(body) + if vis.worse { + return nil + } + + indexVars := candidates.Intersect(outputs) + if len(indexVars) == 0 { + return nil + } + + // Make a sorted set of variable names that will serve as the index key set. + // Sort to ensure deterministic indexing. In future this could be relaxed + // if we can decide that one ordering is better than another. + result := make([]*Term, 0, len(indexVars)) + + for v := range indexVars { + result = append(result, NewTerm(v)) + } + + sort.Slice(result, func(i, j int) bool { + return result[i].Value.Compare(result[j].Value) < 0 + }) + + return &ComprehensionIndex{Term: term, Keys: result} +} + +type comprehensionIndexRegressionCheckVisitor struct { + candidates VarSet + seen VarSet + worse bool +} + +// TOOD(tsandall): Improve this so that users can either supply this list explicitly +// or the information is maintained on the built-in function declaration. What we really +// need to know is whether the built-in function allows callers to push down output +// values or not. It's unlikely that anything outside of OPA does this today so this +// solution is fine for now. +var comprehensionIndexBlacklist = map[string]int{ + WalkBuiltin.Name: len(WalkBuiltin.Decl.Args()), +} + +func newComprehensionIndexRegressionCheckVisitor(candidates VarSet) *comprehensionIndexRegressionCheckVisitor { + return &comprehensionIndexRegressionCheckVisitor{ + candidates: candidates, + seen: NewVarSet(), + } +} + +func (vis *comprehensionIndexRegressionCheckVisitor) Walk(x interface{}) { + NewGenericVisitor(vis.visit).Walk(x) +} + +func (vis *comprehensionIndexRegressionCheckVisitor) visit(x interface{}) bool { + if !vis.worse { + switch x := x.(type) { + case *Expr: + operands := x.Operands() + if pos := comprehensionIndexBlacklist[x.Operator().String()]; pos > 0 && pos < len(operands) { + vis.assertEmptyIntersection(operands[pos].Vars()) + } + case Ref: + vis.assertEmptyIntersection(x.OutputVars()) + case Var: + vis.seen.Add(x) + // Always skip comprehensions. We do not have to visit their bodies here. + case *ArrayComprehension, *SetComprehension, *ObjectComprehension: + return true + } + } + return vis.worse +} + +func (vis *comprehensionIndexRegressionCheckVisitor) assertEmptyIntersection(vs VarSet) { + for v := range vs { + if vis.candidates.Contains(v) && !vis.seen.Contains(v) { + vis.worse = true + return + } + } +} + // ModuleTreeNode represents a node in the module tree. The module // tree is keyed by the package path. type ModuleTreeNode struct { diff --git a/ast/compile_test.go b/ast/compile_test.go index c24b996f68..e0e592e64e 100644 --- a/ast/compile_test.go +++ b/ast/compile_test.go @@ -3377,6 +3377,225 @@ func TestCompilerWithStageAfterWithMetrics(t *testing.T) { } } +func TestCompilerBuildComprehensionIndexKeySet(t *testing.T) { + + tests := []struct { + note string + module string + atRow int + wantTerm string + wantKeys string + }{ + { + note: "example: invert object", + module: ` + package test + + p { + value = input[i] + keys = [j | value = input[j]] + } + `, + atRow: 6, + wantTerm: `[j | value = input[j]]`, + wantKeys: `[value]`, + }, + { + note: "example: multiple keys from body", + module: ` + package test + + p { + v1 = input[i].v1 + v2 = input[i].v2 + keys = [j | v1 = input[j].v1; v2 = input[j].v2] + } + `, + atRow: 7, + wantTerm: `[j | v1 = input[j].v1; v2 = input[j].v2]`, + wantKeys: `[v1, v2]`, + }, + { + note: "example: nested comprehensions are supported", + module: ` + package test + + p = {x: ys | + x = input[i] + ys = {y | x = input[y]} + } + `, + atRow: 6, + wantTerm: `{y | x = input[y]}`, + wantKeys: `[x]`, + }, + { + note: "skip: lone comprehensions", + module: ` + package test + + p { + [v | input[i] = v] # skip because no assignment + }`, + }, + { + note: "skip: due to with modifier", + module: ` + package test + + p { + v = input[i] + ks = [j | input[j] = v] with data.x as 1 # skip because of with modifier + }`, + }, + { + note: "skip: due to negation", + module: ` + package test + + p { + v = input[i] + a = [] + not a = [j | input[j] = v] # skip due to negation + }`, + }, + { + note: "skip: due to lack of comprehension", + module: ` + package test + + p { + v = input[i] + }`, + }, + { + note: "skip: due to unsafe comprehension body", + module: ` + package test + + f(x) { + v = input[i] + ys = [y | y = x[j]] # x is not safe + }`, + }, + { + note: "skip: due to no candidates", + module: ` + package test + + p { + ys = [y | y = input[j]] + }`, + }, + { + note: "skip: avoid increasing runtime (func arg)", + module: ` + package test + + f(x) { + y = input[x] + ys = [y | y = input[x]] + }`, + }, + { + note: "skip: avoid increasing runtime (head key)", + module: ` + package test + + p[x] { + y = input[x] + ys = [y | y = input[x]] + }`, + }, + { + note: "skip: avoid increasing runtime (walk)", + module: ` + package test + + p[x] { + y = input.bar[x] + ys = [y | a = input.foo; walk(a, [x, y])] + }`, + }, + { + note: "bypass: use intermediate var to skip regression check", + module: ` + package test + + p[x] { + y = input[x] + ys = [y | y = input[z]; z = x] + }`, + atRow: 6, + wantTerm: ` [y | y = input[z]; z = x]`, + wantKeys: `[x, y]`, + }, + } + + for _, tc := range tests { + t.Run(tc.note, func(t *testing.T) { + + m := metrics.New() + compiler := NewCompiler().WithMetrics(m) + compiler.Compile(map[string]*Module{"test.rego": MustParseModule(tc.module)}) + if compiler.Failed() { + t.Fatal(compiler.Errors) + } + + n := m.Counter(compileStageComprehensionIndexBuild).Value().(uint64) + + if tc.atRow == 0 { + if n > 0 || len(compiler.comprehensionIndices) > 0 { + t.Fatal("expected no indices to be built. got:", compiler.comprehensionIndices) + } + return + } + + if n != 1 { + t.Fatal("expected counter to be incremented") + } + + var comprehension *Term + + WalkTerms(compiler.Modules["test.rego"], func(x *Term) bool { + if !IsComprehension(x.Value) { + return true + } + if x.Location.Row != tc.atRow { + return false + } else if comprehension != nil { + t.Fatal("expected at most one comprehension per line in test module") + } + comprehension = x + return false + }) + + if comprehension == nil { + t.Fatal("expected comprehension at line:", tc.atRow) + } + + result := compiler.ComprehensionIndex(comprehension) + + if result == nil { + t.Fatal("expected result") + } + + expTerm := MustParseTerm(tc.wantTerm) + + if !result.Term.Equal(expTerm) { + t.Fatalf("expected term to be %v but got: %v", expTerm, result.Term) + } + + expKeys := MustParseTerm(tc.wantKeys).Value.(Array) + + if Array(result.Keys).Compare(expKeys) != 0 { + t.Fatalf("expected keys to be %v but got: %v", expKeys, result.Keys) + } + + }) + } +} + func TestQueryCompiler(t *testing.T) { tests := []struct { note string diff --git a/ast/compilemetrics.go b/ast/compilemetrics.go new file mode 100644 index 0000000000..5d952258da --- /dev/null +++ b/ast/compilemetrics.go @@ -0,0 +1,9 @@ +// Copyright 2020 The OPA Authors. All rights reserved. +// Use of this source code is governed by an Apache2 +// license that can be found in the LICENSE file. + +package ast + +const ( + compileStageComprehensionIndexBuild = "compile_stage_comprehension_index_build" +) diff --git a/docs/content/policy-performance.md b/docs/content/policy-performance.md index 8d5824286a..ea256eb587 100644 --- a/docs/content/policy-performance.md +++ b/docs/content/policy-performance.md @@ -3,6 +3,7 @@ title: Policy Performance kind: documentation weight: 3 --- + ## High Performance Policy Decisions For low-latency/high-performance use-cases, e.g. microservice API authorization, policy evaluation has a budget on the order of 1 millisecond. Not all use cases require that kind of performance, and OPA is powerful enough that you can write expressive policies that take longer than 1 millisecond to evaluate. But for high-performance use cases, there is a fragment of the policy language that has been engineered to evaluate quickly. Even as the size of the policies grow, the performance for this fragment can be nearly constant-time. @@ -131,6 +132,140 @@ For `glob.match(pattern, delimiter, match)` statements to be indexed the pattern | `glob.match("foo:**:bar", [":"], input.x)` | no | pattern contains `**` | | `glob.match("foo:*:bar", [":"], input.x[i])` | no | match contains variable(s) | +### Comprehension Indexing + +Rego does not support mutation. As a result, certain operations like "group by" require +use of comprehensions to aggregate values. To avoid O(n^2) runtime complexity in +queries/rules that perform group-by, OPA may compute and memoize the entire collection +produced by comprehensions at once. This ensures that runtime complexity is O(n) where +n is the size of the collection that group-by/aggregation is being performed on. + +For example, suppose the policy must check if the number of ports exposed on an interface +exceeds some threshold (e.g., any interface may expose up to 100 ports.) The policy is given +the port->interface mapping as a JSON array under `input`: + +```json +{ + "exposed": [ + { + "interface": "eth0", + "port": 8080, + }, + { + "interface": "eth0", + "port": 8081, + }, + { + "interface": "eth1", + "port": 443, + }, + { + "interface": "lo1", + "port": 5000, + } + ] +} +``` + +In this case, the policy must count the number of ports exposed on each interface. To do this, +the policy must first aggregate/group the ports by the interface name. Conceptually, +the policy should generate a document like this: + +```json +{ + "exposed_ports_by_interface": { + "eth0": [8080, 8081], + "eth1": [443], + "lo1": [5000] + } +} +``` + +Since multiple ports could be exposed on a single interface, the policy must use a comprehension to +aggregate the port values by the interface names. To implement this logic in Rego, we would write: + +```rego +some i +intf := input.exposed[i].interface +ports := [port | some j; input.exposed[j].interface == intf; port := input.exposed[j].port] +``` + +Without comprehension indexing, this query would be O(n^2) where n is the size of `input.exposed`. +However, with comprehension indexing, the query remains O(n) because OPA only computes the comprehension +_once_. In this case, the comprehension is evaluated and all possible values of `ports` are computed +at once. These values are indexed by the assignments of `intf`. + +To implement the policy above we could write: + +```rego +deny[msg] { + some i + count(exposed_ports_by_interface[i]) > 100 + msg := sprintf("interface '%v' exposes too many ports", [i]) +} + +exposed_ports_by_interface := {intf: ports | + some i + intf := input.exposed[i].interface + ports := [port | + some j + input.exposed[j].interface == intf + port := input.exposed[j].port + ] +} +``` + +Indices can be built for comprehensions (nested or not) that generate collections (i.e., arrays, sets, or objects) +based on variables in an outer query. In the example above: + +* `intf` is the variable in the outer query. +* `[port | some j; input.exposed[j].interface == intf; port := input.exposed[j].port]` is the comprehension. +* `ports` is the variable the collection is assigned to. + +In order to be indexed, comprehensions must meet the following conditions: + +1. The comprehension appears in an assignment or unification statement. +1. The expression containing the comprehension does not include a `with` statement. +1. The expression containing the comprehension is not negated. +1. The comprehension body is safe when considered independent from the outer query. +1. The comprehension body closes over at least one variable in the outer query and none of these variables appear as outputs in references or `walk()` calls. + +The following examples show cases that are NOT indexed: + +```rego +not_indexed_because_missing_assignment { + x := input[_] + [y | some y; x == input[y]] +} + +not_indexed_because_includes_with { + x := input[_] + ys := [y | some y; x := input[y]] with input as {} +} + +not_indexed_because_negated { + x := input[_] + not data.arr = [y | some y; x := input[y]] +} + +not_indexed_because_safety { + obj := input.foo.bar + x := obj[_] + ys := [y | some y; x == obj[y]] +} + +not_indexed_because_no_closure { + ys := [y | x := input[y]] +} + +not_indexed_because_reference_operand_closure { + x := input[y].x + ys := [y | x == input[y].z[_]] +} +``` + +> The 4th and 5th restrictions may be relaxed in the future. + ### Profiling You can also _profile_ your policies using `opa eval`. The profiler is useful if you need to understand diff --git a/topdown/cache.go b/topdown/cache.go index 33a8b16f26..96c250e13c 100644 --- a/topdown/cache.go +++ b/topdown/cache.go @@ -164,3 +164,74 @@ func (s *refStack) Prefixed(ref ast.Ref) bool { } return false } + +type comprehensionCache struct { + stack []map[*ast.Term]*comprehensionCacheElem +} + +type comprehensionCacheElem struct { + value *ast.Term + children *util.HashMap +} + +func newComprehensionCache() *comprehensionCache { + cache := &comprehensionCache{} + cache.Push() + return cache +} + +func (c *comprehensionCache) Push() { + c.stack = append(c.stack, map[*ast.Term]*comprehensionCacheElem{}) +} + +func (c *comprehensionCache) Pop() { + c.stack = c.stack[:len(c.stack)-1] +} + +func (c *comprehensionCache) Elem(t *ast.Term) (*comprehensionCacheElem, bool) { + elem, ok := c.stack[len(c.stack)-1][t] + return elem, ok +} + +func (c *comprehensionCache) Set(t *ast.Term, elem *comprehensionCacheElem) { + c.stack[len(c.stack)-1][t] = elem +} + +func newComprehensionCacheElem() *comprehensionCacheElem { + return &comprehensionCacheElem{children: newComprehensionCacheHashMap()} +} + +func (c *comprehensionCacheElem) Get(key []*ast.Term) *ast.Term { + node := c + for i := 0; i < len(key); i++ { + x, ok := node.children.Get(key[i]) + if !ok { + return nil + } + node = x.(*comprehensionCacheElem) + } + return node.value +} + +func (c *comprehensionCacheElem) Put(key []*ast.Term, value *ast.Term) { + node := c + for i := 0; i < len(key); i++ { + x, ok := node.children.Get(key[i]) + if ok { + node = x.(*comprehensionCacheElem) + } else { + next := newComprehensionCacheElem() + node.children.Put(key[i], next) + node = next + } + } + node.value = value +} + +func newComprehensionCacheHashMap() *util.HashMap { + return util.NewHashMap(func(a, b util.T) bool { + return a.(*ast.Term).Equal(b.(*ast.Term)) + }, func(x util.T) int { + return x.(*ast.Term).Hash() + }) +} diff --git a/topdown/errors.go b/topdown/errors.go index d0457f1a83..c3e21bc1a4 100644 --- a/topdown/errors.go +++ b/topdown/errors.go @@ -117,3 +117,11 @@ func mergeConflictErr(loc *ast.Location) error { Message: "real and replacement data could not be merged", } } + +func internalErr(loc *ast.Location, msg string) error { + return &Error{ + Code: InternalErr, + Location: loc, + Message: msg, + } +} diff --git a/topdown/eval.go b/topdown/eval.go index 4399184b51..0ec38ad02e 100644 --- a/topdown/eval.go +++ b/topdown/eval.go @@ -30,37 +30,38 @@ func (f *queryIDFactory) Next() uint64 { } type eval struct { - ctx context.Context - seed io.Reader - queryID uint64 - queryIDFact *queryIDFactory - parent *eval - caller *eval - cancel Cancel - query ast.Body - queryCompiler ast.QueryCompiler - index int - indexing bool - bindings *bindings - store storage.Store - baseCache *baseCache - txn storage.Transaction - compiler *ast.Compiler - input *ast.Term - data *ast.Term - targetStack *refStack - tracers []Tracer - instr *Instrumentation - builtins map[string]*Builtin - builtinCache builtins.Cache - virtualCache *virtualCache - saveSet *saveSet - saveStack *saveStack - saveSupport *saveSupport - saveNamespace *ast.Term - disableInlining [][]ast.Ref - genvarprefix string - runtime *ast.Term + ctx context.Context + seed io.Reader + queryID uint64 + queryIDFact *queryIDFactory + parent *eval + caller *eval + cancel Cancel + query ast.Body + queryCompiler ast.QueryCompiler + index int + indexing bool + bindings *bindings + store storage.Store + baseCache *baseCache + txn storage.Transaction + compiler *ast.Compiler + input *ast.Term + data *ast.Term + targetStack *refStack + tracers []Tracer + instr *Instrumentation + builtins map[string]*Builtin + builtinCache builtins.Cache + virtualCache *virtualCache + comprehensionCache *comprehensionCache + saveSet *saveSet + saveStack *saveStack + saveSupport *saveSupport + saveNamespace *ast.Term + disableInlining [][]ast.Ref + genvarprefix string + runtime *ast.Term } func (e *eval) Run(iter evalIterator) error { @@ -384,7 +385,6 @@ func (e *eval) evalWith(iter evalIterator) error { } input, err := mergeTermWithValues(e.input, pairsInput) - if err != nil { return &Error{ Code: ConflictErr, @@ -432,6 +432,7 @@ func (e *eval) evalWithPush(input *ast.Term, data *ast.Term, targets []ast.Ref, e.data = data } + e.comprehensionCache.Push() e.virtualCache.Push() e.targetStack.Push(targets) e.disableInlining = append(e.disableInlining, disable) @@ -443,6 +444,7 @@ func (e *eval) evalWithPop(input *ast.Term, data *ast.Term) { e.disableInlining = e.disableInlining[:len(e.disableInlining)-1] e.targetStack.Pop() e.virtualCache.Pop() + e.comprehensionCache.Pop() e.data = data e.input = input } @@ -850,6 +852,14 @@ func (e *eval) biunifyComprehension(a, b *ast.Term, b1, b2 *bindings, swap bool, return e.biunifyComprehensionPartial(a, b, b1, b2, swap, iter) } + value, err := e.buildComprehensionCache(a) + + if err != nil { + return err + } else if value != nil { + return e.biunify(value, b, b1, b2, iter) + } + switch a := a.Value.(type) { case *ast.ArrayComprehension: return e.biunifyComprehensionArray(a, b, b1, b2, iter) @@ -859,7 +869,106 @@ func (e *eval) biunifyComprehension(a, b *ast.Term, b1, b2 *bindings, swap bool, return e.biunifyComprehensionObject(a, b, b1, b2, iter) } - return fmt.Errorf("illegal comprehension %T", a) + return internalErr(e.query[e.index].Location, "illegal comprehension type") +} + +func (e *eval) buildComprehensionCache(a *ast.Term) (*ast.Term, error) { + + index := e.comprehensionIndex(a) + if index == nil { + e.instr.counterIncr(evalOpComprehensionCacheSkip) + return nil, nil + } + + cache, ok := e.comprehensionCache.Elem(a) + if !ok { + var err error + switch x := a.Value.(type) { + case *ast.ArrayComprehension: + cache, err = e.buildComprehensionCacheArray(x, index.Keys) + case *ast.SetComprehension: + cache, err = e.buildComprehensionCacheSet(x, index.Keys) + case *ast.ObjectComprehension: + cache, err = e.buildComprehensionCacheObject(x, index.Keys) + default: + err = internalErr(e.query[e.index].Location, "illegal comprehension type") + } + if err != nil { + return nil, err + } + e.comprehensionCache.Set(a, cache) + e.instr.counterIncr(evalOpComprehensionCacheBuild) + } else { + e.instr.counterIncr(evalOpComprehensionCacheHit) + } + + values := make([]*ast.Term, len(index.Keys)) + + for i := range index.Keys { + values[i] = e.bindings.Plug(index.Keys[i]) + } + + return cache.Get(values), nil +} + +func (e *eval) buildComprehensionCacheArray(x *ast.ArrayComprehension, keys []*ast.Term) (*comprehensionCacheElem, error) { + child := e.child(x.Body) + node := newComprehensionCacheElem() + return node, child.Run(func(child *eval) error { + values := make([]*ast.Term, len(keys)) + for i := range keys { + values[i] = child.bindings.Plug(keys[i]) + } + head := child.bindings.Plug(x.Term) + cached := node.Get(values) + if cached != nil { + cached.Value = append(cached.Value.(ast.Array), head) + } else { + node.Put(values, ast.ArrayTerm(head)) + } + return nil + }) +} + +func (e *eval) buildComprehensionCacheSet(x *ast.SetComprehension, keys []*ast.Term) (*comprehensionCacheElem, error) { + child := e.closure(x.Body) + node := newComprehensionCacheElem() + return node, child.Run(func(child *eval) error { + values := make([]*ast.Term, len(keys)) + for i := range keys { + values[i] = child.bindings.Plug(keys[i]) + } + head := child.bindings.Plug(x.Term) + cached := node.Get(values) + if cached != nil { + set := cached.Value.(ast.Set) + set.Add(head) + } else { + node.Put(values, ast.SetTerm(head)) + } + return nil + }) +} + +func (e *eval) buildComprehensionCacheObject(x *ast.ObjectComprehension, keys []*ast.Term) (*comprehensionCacheElem, error) { + child := e.closure(x.Body) + node := newComprehensionCacheElem() + return node, child.Run(func(child *eval) error { + values := make([]*ast.Term, len(keys)) + for i := range keys { + values[i] = child.bindings.Plug(keys[i]) + } + headKey := child.bindings.Plug(x.Key) + headValue := child.bindings.Plug(x.Value) + cached := node.Get(values) + if cached != nil { + obj := cached.Value.(ast.Object) + obj.Insert(headKey, headValue) + } else { + node.Put(values, ast.ObjectTerm(ast.Item(headKey, headValue))) + } + return nil + }) } func (e *eval) biunifyComprehensionPartial(a, b *ast.Term, b1, b2 *bindings, swap bool, iter unifyIterator) error { @@ -2227,6 +2336,13 @@ func (e evalTerm) save(iter unifyIterator) error { return e.e.biunify(ast.NewTerm(ref), e.rterm, e.termbindings, e.rbindings, iter) } +func (e *eval) comprehensionIndex(term *ast.Term) *ast.ComprehensionIndex { + if e.queryCompiler != nil { + return e.queryCompiler.ComprehensionIndex(term) + } + return e.compiler.ComprehensionIndex(term) +} + func applyCopyPropagation(p *copypropagation.CopyPropagator, instr *Instrumentation, body ast.Body) ast.Body { instr.startTimer(partialOpCopyPropagation) result := p.Apply(body) diff --git a/topdown/instrumentation.go b/topdown/instrumentation.go index fa871cf908..7c391bb8e2 100644 --- a/topdown/instrumentation.go +++ b/topdown/instrumentation.go @@ -7,18 +7,21 @@ package topdown import "github.com/open-policy-agent/opa/metrics" const ( - evalOpPlug = "eval_op_plug" - evalOpResolve = "eval_op_resolve" - evalOpRuleIndex = "eval_op_rule_index" - evalOpBuiltinCall = "eval_op_builtin_call" - evalOpVirtualCacheHit = "eval_op_virtual_cache_hit" - evalOpVirtualCacheMiss = "eval_op_virtual_cache_miss" - evalOpBaseCacheHit = "eval_op_base_cache_hit" - evalOpBaseCacheMiss = "eval_op_base_cache_miss" - partialOpSaveUnify = "partial_op_save_unify" - partialOpSaveSetContains = "partial_op_save_set_contains" - partialOpSaveSetContainsRec = "partial_op_save_set_contains_rec" - partialOpCopyPropagation = "partial_op_copy_propagation" + evalOpPlug = "eval_op_plug" + evalOpResolve = "eval_op_resolve" + evalOpRuleIndex = "eval_op_rule_index" + evalOpBuiltinCall = "eval_op_builtin_call" + evalOpVirtualCacheHit = "eval_op_virtual_cache_hit" + evalOpVirtualCacheMiss = "eval_op_virtual_cache_miss" + evalOpBaseCacheHit = "eval_op_base_cache_hit" + evalOpBaseCacheMiss = "eval_op_base_cache_miss" + evalOpComprehensionCacheSkip = "eval_op_comprehension_cache_skip" + evalOpComprehensionCacheBuild = "eval_op_comprehension_cache_build" + evalOpComprehensionCacheHit = "eval_op_comprehension_cache_hit" + partialOpSaveUnify = "partial_op_save_unify" + partialOpSaveSetContains = "partial_op_save_set_contains" + partialOpSaveSetContainsRec = "partial_op_save_set_contains_rec" + partialOpCopyPropagation = "partial_op_copy_propagation" ) // Instrumentation implements helper functions to instrument query evaluation diff --git a/topdown/query.go b/topdown/query.go index ced1ef77e0..1ce6a9a7be 100644 --- a/topdown/query.go +++ b/topdown/query.go @@ -185,32 +185,33 @@ func (q *Query) PartialRun(ctx context.Context) (partials []ast.Body, support [] f := &queryIDFactory{} b := newBindings(0, q.instr) e := &eval{ - ctx: ctx, - seed: q.seed, - cancel: q.cancel, - query: q.query, - queryCompiler: q.queryCompiler, - queryIDFact: f, - queryID: f.Next(), - bindings: b, - compiler: q.compiler, - store: q.store, - baseCache: newBaseCache(), - targetStack: newRefStack(), - txn: q.txn, - input: q.input, - tracers: q.tracers, - instr: q.instr, - builtins: q.builtins, - builtinCache: builtins.Cache{}, - virtualCache: newVirtualCache(), - saveSet: newSaveSet(q.unknowns, b, q.instr), - saveStack: newSaveStack(), - saveSupport: newSaveSupport(), - saveNamespace: ast.StringTerm(q.partialNamespace), - genvarprefix: q.genvarprefix, - runtime: q.runtime, - indexing: q.indexing, + ctx: ctx, + seed: q.seed, + cancel: q.cancel, + query: q.query, + queryCompiler: q.queryCompiler, + queryIDFact: f, + queryID: f.Next(), + bindings: b, + compiler: q.compiler, + store: q.store, + baseCache: newBaseCache(), + targetStack: newRefStack(), + txn: q.txn, + input: q.input, + tracers: q.tracers, + instr: q.instr, + builtins: q.builtins, + builtinCache: builtins.Cache{}, + virtualCache: newVirtualCache(), + comprehensionCache: newComprehensionCache(), + saveSet: newSaveSet(q.unknowns, b, q.instr), + saveStack: newSaveStack(), + saveSupport: newSaveSupport(), + saveNamespace: ast.StringTerm(q.partialNamespace), + genvarprefix: q.genvarprefix, + runtime: q.runtime, + indexing: q.indexing, } if len(q.disableInlining) > 0 { @@ -285,28 +286,29 @@ func (q *Query) Iter(ctx context.Context, iter func(QueryResult) error) error { } f := &queryIDFactory{} e := &eval{ - ctx: ctx, - seed: q.seed, - cancel: q.cancel, - query: q.query, - queryCompiler: q.queryCompiler, - queryIDFact: f, - queryID: f.Next(), - bindings: newBindings(0, q.instr), - compiler: q.compiler, - store: q.store, - baseCache: newBaseCache(), - targetStack: newRefStack(), - txn: q.txn, - input: q.input, - tracers: q.tracers, - instr: q.instr, - builtins: q.builtins, - builtinCache: builtins.Cache{}, - virtualCache: newVirtualCache(), - genvarprefix: q.genvarprefix, - runtime: q.runtime, - indexing: q.indexing, + ctx: ctx, + seed: q.seed, + cancel: q.cancel, + query: q.query, + queryCompiler: q.queryCompiler, + queryIDFact: f, + queryID: f.Next(), + bindings: newBindings(0, q.instr), + compiler: q.compiler, + store: q.store, + baseCache: newBaseCache(), + targetStack: newRefStack(), + txn: q.txn, + input: q.input, + tracers: q.tracers, + instr: q.instr, + builtins: q.builtins, + builtinCache: builtins.Cache{}, + virtualCache: newVirtualCache(), + comprehensionCache: newComprehensionCache(), + genvarprefix: q.genvarprefix, + runtime: q.runtime, + indexing: q.indexing, } e.caller = e q.startTimer(metrics.RegoQueryEval) diff --git a/topdown/topdown_test.go b/topdown/topdown_test.go index 29e212b58b..c0f2a884b9 100644 --- a/topdown/topdown_test.go +++ b/topdown/topdown_test.go @@ -2385,6 +2385,25 @@ func TestTopDownWithKeyword(t *testing.T) { setl[x] { data.foo[x] }`}, rules: []string{`p = true { data.ex.setl[1] with data.foo as {1} }`}, }, + { + // NOTE(tsandall): This case assumes that partial sets are not memoized. + // If we change that, it'll be harder to test that the comprehension + // cache is invalidated. + note: "invalidate comprehension cache", + exp: `[[{"b": ["a", "c"]}], [{"b": ["a"]}]]`, + modules: []string{`package ex + s[x] { + x = {v: ks | + v = input[i] + ks = {k | v = input[k]} + } + } + `}, + rules: []string{`p = [x, y] { + x = data.ex.s with input as {"a": "b", "c": "b"} + y = data.ex.s with input as {"a": "b"} + }`}, + }, } for _, tc := range tests {