diff --git a/cmd/eval.go b/cmd/eval.go index a257687eaf..f2ae75e416 100644 --- a/cmd/eval.go +++ b/cmd/eval.go @@ -34,6 +34,7 @@ type evalCommandParams struct { partial bool unknowns []string disableInlining []string + shallowInlining bool disableIndexing bool dataPaths repeatedStringFlag inputPath string @@ -219,6 +220,7 @@ Set the output format with the --format flag. evalCommand.Flags().BoolVarP(¶ms.partial, "partial", "p", false, "perform partial evaluation") evalCommand.Flags().StringArrayVarP(¶ms.unknowns, "unknowns", "u", []string{"input"}, "set paths to treat as unknown during partial evaluation") evalCommand.Flags().StringArrayVarP(¶ms.disableInlining, "disable-inlining", "", []string{}, "set paths of documents to exclude from inlining") + evalCommand.Flags().BoolVarP(¶ms.shallowInlining, "shallow-inlining", "", false, "disable inlining of rules that depend on unknowns") evalCommand.Flags().BoolVar(¶ms.disableIndexing, "disable-indexing", false, "disable indexing optimizations") evalCommand.Flags().BoolVarP(¶ms.instrument, "instrument", "", false, "enable query instrumentation metrics (implies --metrics)") evalCommand.Flags().BoolVarP(¶ms.profile, "profile", "", false, "perform expression profiling") @@ -436,7 +438,7 @@ func setupEval(args []string, params evalCommandParams) (*evalContext, error) { regoArgs = append(regoArgs, rego.Unknowns(params.unknowns)) } - regoArgs = append(regoArgs, rego.DisableInlining(params.disableInlining)) + regoArgs = append(regoArgs, rego.DisableInlining(params.disableInlining), rego.ShallowInlining(params.shallowInlining)) var c *cover.Cover diff --git a/compile/compile.go b/compile/compile.go index e3ea903e7d..380e72de06 100644 --- a/compile/compile.go +++ b/compile/compile.go @@ -115,8 +115,7 @@ func (c *Compiler) WithEntrypoints(e ...string) *Compiler { // WithOptimizationLevel sets the optimization level on the compiler. By default // optimizations are disabled. Higher levels apply more aggressive optimizations -// but can take longer. Currently only two levels are supported: 0 (disabled) and -// 1 (enabled). +// but can take longer. func (c *Compiler) WithOptimizationLevel(n int) *Compiler { c.optimizationLevel = n return c @@ -277,7 +276,8 @@ func (c *Compiler) optimize(ctx context.Context) error { o := newOptimizer(c.bundle). WithEntrypoints(c.entrypointrefs). - WithDebug(c.debug) + WithDebug(c.debug). + WithShallowInlining(c.optimizationLevel <= 1) err := o.Do(ctx) if err != nil { @@ -334,6 +334,7 @@ type optimizer struct { nsprefix string resultsymprefix string outputprefix string + shallow bool debug *debugEvents } @@ -356,13 +357,12 @@ func (o *optimizer) WithEntrypoints(es []*ast.Term) *optimizer { return o } -func (o *optimizer) Do(ctx context.Context) error { +func (o *optimizer) WithShallowInlining(yes bool) *optimizer { + o.shallow = yes + return o +} - // TODO(tsandall): implement optimization levels. These will just be params on partial evaluation for now. - // - // Level 1: PE w/ constant folding. Only inline rules that are completely known. - // Level 2: L1 except inlining of rules with unknowns. - // Level 3: L2 except aggressive inlining using negation and copy propagation optimizations. +func (o *optimizer) Do(ctx context.Context) error { // NOTE(tsandall): if there are multiple entrypoints, copy the bundle because // if any of the optimization steps fail, we do not want to leave the caller's @@ -399,6 +399,7 @@ func (o *optimizer) Do(ctx context.Context) error { rego.ParsedQuery(ast.NewBody(ast.Equality.Expr(resultsym, e))), rego.PartialNamespace(o.nsprefix), rego.DisableInlining(o.findRequiredDocuments(e)), + rego.ShallowInlining(o.shallow), rego.SkipPartialNamespace(true), rego.Compiler(o.compiler), rego.Store(store), diff --git a/compile/compile_test.go b/compile/compile_test.go index d363e9f5bd..399b7b5d9b 100644 --- a/compile/compile_test.go +++ b/compile/compile_test.go @@ -268,12 +268,14 @@ func TestCompilerError(t *testing.T) { }) } -func TestCompilerOptimization(t *testing.T) { +func TestCompilerOptimizationL1(t *testing.T) { + files := map[string]string{ "test.rego": ` package test default p = false - p { input.x = data.foo }`, + p { q } + q { input.x = data.foo }`, "data.json": ` {"foo": 1}`, } @@ -290,15 +292,82 @@ func TestCompilerOptimization(t *testing.T) { t.Fatal(err) } - exp := ast.MustParseModule(` + optimizedExp := ast.MustParseModule(` package test - p { input.x = 1} + q { input.x = 1 } + p { data.test.q = X; X } default p = false `) - if len(compiler.bundle.Modules) != 1 || !compiler.bundle.Modules[0].Parsed.Equal(exp) { - t.Fatalf("expected one module but got: %v", compiler.bundle.Modules) + // NOTE(tsandall): PE generates vars with wildcard prefix. Instead of + // constructing the AST manually, just rewrite to the expected value + // here. If this becomes a common pattern, we could refactor (e.g., + // allow caller to control var prefix, split into a reusable function, + // etc.) + ast.TransformVars(optimizedExp, func(x ast.Var) (ast.Value, error) { + if x == ast.Var("X") { + return ast.Var("$_term_1_01"), nil + } + return x, nil + }) + + if len(compiler.bundle.Modules) != 1 { + t.Fatalf("expected 1 module but got: %v", compiler.bundle.Modules) + } + + if !compiler.bundle.Modules[0].Parsed.Equal(optimizedExp) { + t.Fatalf("expected optimized module to be:\n\n%v\n\ngot:\n\n%v", optimizedExp, compiler.bundle.Modules[0]) + } + }) +} + +func TestCompilerOptimizationL2(t *testing.T) { + + files := map[string]string{ + "test.rego": ` + package test + default p = false + p { q } + q { input.x = data.foo }`, + "data.json": ` + {"foo": 1}`, + } + + test.WithTempFS(files, func(root string) { + + compiler := New(). + WithPaths(root). + WithOptimizationLevel(2). + WithEntrypoints("test/p") + + err := compiler.Build(context.Background()) + if err != nil { + t.Fatal(err) + } + + prunedExp := ast.MustParseModule(` + package test + + q { input.x = data.foo }`) + + optimizedExp := ast.MustParseModule(` + package test + + p { input.x = 1 } + default p = false + `) + + if len(compiler.bundle.Modules) != 2 { + t.Fatalf("expected two modules but got: %v", compiler.bundle.Modules) + } + + if !compiler.bundle.Modules[0].Parsed.Equal(prunedExp) { + t.Fatalf("expected pruned module to be:\n\n%v\n\ngot:\n\n%v", prunedExp, compiler.bundle.Modules[0]) + } + + if !compiler.bundle.Modules[1].Parsed.Equal(optimizedExp) { + t.Fatalf("expected optimized module to be:\n\n%v\n\ngot:\n\n%v", optimizedExp, compiler.bundle.Modules[1]) } }) } diff --git a/docs/content/policy-performance.md b/docs/content/policy-performance.md index 47408ff270..d67c5f38de 100644 --- a/docs/content/policy-performance.md +++ b/docs/content/policy-performance.md @@ -679,6 +679,39 @@ loading 10,000 rules that implement an ACL-style authorization policy consumes a 130MB of RAM while 100,000 rules implementing the same policy (but with 10x more tuples to check) consumes approximately 1.1GB of RAM. +## Optimization Levels + +The `--optimize` (or `-O`) flag on the `opa build` command controls how bundles are optimized. + +> Optimization applies partial evaluation to precompute _known_ values in the policy. The goal of +partial evaluation is to convert non-linear-time policies into linear-time policies. + +By specifying the `--optimize` flag, users can control how much time and resources are spent +attempting to optimize the bundle. Generally, higher optimization levels require more time +and resources. Currently, OPA supports three optimization levels. The exact optimizations applied +in each level may change over time. + +### -O=0 (default) + +By default optimizations are disabled. + +### -O=1 (recommended) + +Policies are partially evaluated. Rules that do not depend on unknowns are evaluated and the +virtual documents they produce are inlined into call sites. If a virtual virtual is required at +evaluation time (e.g., because it is targetted by a `with` statement), then it will not be inlined. + +Rules that DO NOT depend on unknowns are also partially evaluated however the virtual documents +they produce ARE NOT inlined into call sites. The output policy should be structurally similar +to the input policy. + +### -O=2 (aggressive) + +Same as `-O=1` except virtual documents produced by rules that depend on unknowns may be inlined +into call sites. In addition, more aggressive inlining is applied within rules. This includes +[copy propagation](https://en.wikipedia.org/wiki/Copy_propagation) and inlining of certain negated +statements that would otherwise generate support rules. + ## Key Takeaways For high-performance use cases: diff --git a/rego/rego.go b/rego/rego.go index c49bf65dfe..3d32ff91ed 100644 --- a/rego/rego.go +++ b/rego/rego.go @@ -447,6 +447,7 @@ type Rego struct { unknowns []string parsedUnknowns []*ast.Term disableInlining []string + shallowInlining bool skipPartialNamespace bool partialNamespace string modules []rawModule @@ -750,6 +751,14 @@ func DisableInlining(paths []string) func(r *Rego) { } } +// ShallowInlining prevents rules that depend on unknown values from being inlined. +// Rules that only depend on known values are inlined. +func ShallowInlining(yes bool) func(r *Rego) { + return func(r *Rego) { + r.shallowInlining = yes + } +} + // SkipPartialNamespace disables namespacing of partial evalution results for support // rules generated from policy. Synthetic support rules are still namespaced. func SkipPartialNamespace(yes bool) func(r *Rego) { @@ -1800,7 +1809,8 @@ func (r *Rego) partial(ctx context.Context, ectx *EvalContext) (*PartialQueries, WithRuntime(r.runtime). WithIndexing(ectx.indexing). WithPartialNamespace(ectx.partialNamespace). - WithSkipPartialNamespace(r.skipPartialNamespace) + WithSkipPartialNamespace(r.skipPartialNamespace). + WithShallowInlining(r.shallowInlining) for i := range ectx.tracers { q = q.WithTracer(ectx.tracers[i]) diff --git a/rego/rego_test.go b/rego/rego_test.go index 2bb6511e7b..2091c4e245 100644 --- a/rego/rego_test.go +++ b/rego/rego_test.go @@ -1590,6 +1590,42 @@ func TestSkipPartialNamespaceOption(t *testing.T) { } } +func TestShallowInliningOption(t *testing.T) { + r := New(Query("data.test.p = true"), Module("example.rego", ` + package test + + p { + q = true + } + + q { + input.x = r + } + + r = 7 + `), ShallowInlining(true)) + + pq, err := r.Partial(context.Background()) + if err != nil { + t.Fatal(err) + } + + if len(pq.Queries) != 1 || !pq.Queries[0].Equal(ast.MustParseBody("data.partial.test.p = true")) { + t.Fatal("expected exactly one query and ref to be rewritten but got:", pq.Queries) + } + + exp := ast.MustParseModule(` + package partial.test + + q { 7 = input.x } + p { data.partial.test.q = true } + `) + + if len(pq.Support) != 1 || !pq.Support[0].Equal(exp) { + t.Fatal("expected module:", exp, "\n\ngot module:", pq.Support[0]) + } +} + func TestPrepareWithEmptyModule(t *testing.T) { _, err := New( Query("d"),