From ec130ec2a2ad42220d2e32112c54885b670f7f36 Mon Sep 17 00:00:00 2001 From: Anders Eknert Date: Wed, 29 Jan 2025 10:14:50 +0100 Subject: [PATCH] eval: optimize iteration (#7327) Since we do a lot of this in Rego, I was curious to learn what the cost of that was, and if it was possible to do better. I identified the two main costs as these: - `handleErr` wrapper allocating in each iteration, even though the normal condition is to not handle an error - `biunify` callback function literal leaking to the heap in each iteration, with an _high_ cost associated for such a simple one-liner function. While it wasn't possible to change that (as it referenced the outer scope) it *was* possible to inline the whole biunify flow for the array iteration case, as that turned out to be quite simple when stepped through with a debugger. Ideally we'd be able to inline iteation over sets and objects too, but that's much more complex (see comment in code), so I left that for the future. The numbers benchmarked are in the PR, so take a look there. The array iteration case is *quite* the improvement! For reference, this amounts to about 2 million allocations less in the `regal lint bundle` benchmark, and a noticeable improvement in evaluation time. Signed-off-by: Anders Eknert --- v1/rego/rego_bench_test.go | 137 ++++++++++++++++++++++++++++++++++++- v1/topdown/eval.go | 49 ++++++++++--- 2 files changed, 173 insertions(+), 13 deletions(-) diff --git a/v1/rego/rego_bench_test.go b/v1/rego/rego_bench_test.go index ce44a1a277..61207fe1f7 100644 --- a/v1/rego/rego_bench_test.go +++ b/v1/rego/rego_bench_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "os" + "strconv" "testing" "github.com/open-policy-agent/opa/internal/runtime" @@ -175,16 +176,148 @@ func BenchmarkAciTestOnlyEval(b *testing.B) { b.ReportAllocs() for i := 0; i < b.N; i++ { - res, err := pq.Eval(ctx, EvalParsedInput(input.Value)) if err != nil { b.Fatal(err) } - _ = res } } +// BenchmarkArrayIteration-10 +// 15574 77121 ns/op 67249 B/op 1115 allocs/op // handleErr wrapping, not inlined +// 33862 35864 ns/op 5768 B/op 93 allocs/op // handleErr only on error, inlined +func BenchmarkArrayIteration(b *testing.B) { + ctx := context.Background() + + at := make([]*ast.Term, 512) + for i := 0; i < 511; i++ { + at[i] = ast.StringTerm("a") + } + at[511] = ast.StringTerm("v") + + input := ast.NewObject(ast.Item(ast.StringTerm("foo"), ast.ArrayTerm(at...))) + module := ast.MustParseModule(`package test + + default r := false + + r if input.foo[_] == "v"`) + + r := New(Query("data.test.r = x"), ParsedModule(module)) + + pq, err := r.PrepareForEval(ctx) + if err != nil { + b.Fatal(err) + } + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + res, err := pq.Eval(ctx, EvalParsedInput(input)) + if err != nil { + b.Fatal(err) + } + + if res == nil { + b.Fatal("expected result") + } + + if res[0].Bindings["x"].(bool) != true { + b.Fatalf("expected true, got %v", res[0].Bindings["x"]) + } + } +} + +// BenchmarkSetIteration-10 +// 4800 272403 ns/op 80875 B/op 1193 allocs/op // handleErr wrapping, not inlined +// 4933 223234 ns/op 76772 B/op 681 allocs/op // handleErr only on error, not inlined +func BenchmarkSetIteration(b *testing.B) { + ctx := context.Background() + + at := make([]*ast.Term, 512) + for i := 0; i < 512; i++ { + at[i] = ast.StringTerm(strconv.Itoa(i)) + } + + input := ast.NewObject(ast.Item(ast.StringTerm("foo"), ast.ArrayTerm(at...))) + module := ast.MustParseModule(`package test + + s := {x | x := input.foo[_]} + + default r := false + + r if s[_] == "not found"`) + + r := New(Query("data.test.r = x"), ParsedModule(module)) + + pq, err := r.PrepareForEval(ctx) + if err != nil { + b.Fatal(err) + } + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + res, err := pq.Eval(ctx, EvalParsedInput(input)) + if err != nil { + b.Fatal(err) + } + if res == nil { + b.Fatal("expected result") + } + if res[0].Bindings["x"].(bool) != false { + b.Fatalf("expected false, got %v", res[0].Bindings["x"]) + } + } +} + +// BenchmarkObjectIteration-10 +// 12067 99582 ns/op 72830 B/op 1126 allocs/op // handleErr wrapping, not inlined +// 15358 85080 ns/op 27752 B/op 615 allocs/op // handleErr only on error, not inlined +func BenchmarkObjectIteration(b *testing.B) { + ctx := context.Background() + + at := make([][2]*ast.Term, 512) + for i := 0; i < 512; i++ { + at[i] = ast.Item(ast.StringTerm(strconv.Itoa(i)), ast.StringTerm(strconv.Itoa(i))) + } + + input := ast.NewObject(ast.Item(ast.StringTerm("foo"), ast.ObjectTerm(at...))) + module := ast.MustParseModule(`package test + + default r := false + + r if { + input.foo[_] == "512" + } + `) + + r := New(Query("data.test.r = x"), ParsedModule(module)) + + pq, err := r.PrepareForEval(ctx) + if err != nil { + b.Fatal(err) + } + + b.ResetTimer() + b.ReportAllocs() + + for i := 0; i < b.N; i++ { + res, err := pq.Eval(ctx, EvalParsedInput(input)) + if err != nil { + b.Fatal(err) + } + if res == nil { + b.Fatal("expected result") + } + if res[0].Bindings["x"].(bool) != false { + b.Fatalf("expected false, got %v", res[0].Bindings["x"]) + } + } +} + func mustReadFileAsString(b *testing.B, path string) string { b.Helper() diff --git a/v1/topdown/eval.go b/v1/topdown/eval.go index 4758759e71..c762bce9d9 100644 --- a/v1/topdown/eval.go +++ b/v1/topdown/eval.go @@ -3650,28 +3650,55 @@ func (e evalTerm) enumerate(iter unifyIterator) error { switch v := e.term.Value.(type) { case *ast.Array: + // Note(anders): + // For this case (e.g. input.foo[_]), we can avoid the (quite expensive) overhead of a callback + // function literal escaping to the heap in each iteration by inlining the biunification logic, + // meaning a 10x reduction in both the number of allocations made as well as the memory consumed. + // It is possible that such inlining could be done for the set/object cases as well, and that's + // worth looking into later, as I imagine set iteration in particular would be an even greater + // win across most policies. Those cases are however much more complex, as we need to deal with + // any type on either side, not just int/var as is the case here. for i := 0; i < v.Len(); i++ { - k := ast.InternedIntNumberTerm(i) - if err := handleErr(e.e.biunify(k, e.ref[e.pos], e.bindings, e.bindings, func() error { - return e.next(iter, k) - })); err != nil { - return err + a := ast.InternedIntNumberTerm(i) + b := e.ref[e.pos] + + if _, ok := b.Value.(ast.Var); ok { + if e.e.traceEnabled { + e.e.traceUnify(a, b) + } + var undo undo + b, e.bindings = e.bindings.apply(b) + e.bindings.bind(b, a, e.bindings, &undo) + + err := e.next(iter, a) + undo.Undo() + if err != nil { + if err := handleErr(err); err != nil { + return err + } + } } } case ast.Object: for _, k := range v.Keys() { - if err := handleErr(e.e.biunify(k, e.ref[e.pos], e.termbindings, e.bindings, func() error { + err := e.e.biunify(k, e.ref[e.pos], e.termbindings, e.bindings, func() error { return e.next(iter, e.termbindings.Plug(k)) - })); err != nil { - return err + }) + if err != nil { + if err := handleErr(err); err != nil { + return err + } } } case ast.Set: for _, elem := range v.Slice() { - if err := handleErr(e.e.biunify(elem, e.ref[e.pos], e.termbindings, e.bindings, func() error { + err := e.e.biunify(elem, e.ref[e.pos], e.termbindings, e.bindings, func() error { return e.next(iter, e.termbindings.Plug(elem)) - })); err != nil { - return err + }) + if err != nil { + if err := handleErr(err); err != nil { + return err + } } } }