wasm/sdk: add benchmarks (open-policy-agent#3103)

This is porting some of the existing topdown tests --
but run with the Wasm SDK. I've factored out the data
generation bits into the util package; using it in both places.

I've come to believe that whatever memory usage data we're
collecting there is probably bogus: too many of my runs yield

    1597 B/op	      41 allocs/op

regardless of the input size. I'd think that the problem lies in
the boundary crossing to cgo-land, but I have yet to find a
source for that.

That aside, disabling the allocation reporting, and gathering
benchmark data for the run times is probably already useful,
so let's go with that.

Signed-off-by: Stephan Renatus <[email protected]>

internal/wasm/sdk/internal/wasm/pool.go

@@ -196,7 +196,6 @@ func (p *Pool) SetPolicyData(policy []byte, data []byte) error {
 
 	if bytes.Equal(policy, currentPolicy) && bytes.Equal(data, currentData) {
 		return nil
-
 	}
 
 	err := p.setPolicyData(policy, data)

internal/wasm/sdk/internal/wasm/util.go

@@ -4,12 +4,13 @@
 
 package wasm
 
-const wasmPageSize = 65535
+// PageSize represents the WASM page size in bytes.
+const PageSize = 65535
 
 // Pages converts a byte size to Pages, rounding up as necessary.
 func Pages(n uint32) uint32 {
-	pages := n / wasmPageSize
-	if pages*wasmPageSize == n {
+	pages := n / PageSize
+	if pages*PageSize == n {
 		return pages
 	}
 

internal/wasm/sdk/internal/wasm/vm.go

@@ -8,7 +8,6 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"strconv"
 	"time"
@@ -259,7 +258,7 @@ func (i *VM) Eval(ctx context.Context, entrypoint int32, input *interface{}, met
 			if e := recover(); e != nil {
 				switch e := e.(type) {
 				case abortError:
-					err = errors.New(e.message)
+					err = fmt.Errorf(e.message)
 				case builtinError:
 					err = e.err
 					if _, ok := err.(topdown.Halt); !ok {

internal/wasm/sdk/opa/config.go

@@ -64,7 +64,7 @@ func (o *OPA) WithDataJSON(data interface{}) *OPA {
 // WithMemoryLimits configures the memory limits (in bytes) for a single policy
 // evaluation.
 func (o *OPA) WithMemoryLimits(min, max uint32) *OPA {
-	if min < 2*65535 {
+	if min < 2*wasm.PageSize {
 		o.configErr = fmt.Errorf("too low minimum memory limit: %w", errors.ErrInvalidConfig)
 		return o
 	}

internal/wasm/sdk/opa/opa.go

@@ -182,7 +182,7 @@ func (o *OPA) Eval(ctx context.Context, opts EvalOpts) (*Result, error) {
 		return nil, fmt.Errorf("%v: %w", err, errors.ErrInternal)
 	}
 
-	return &Result{result}, nil
+	return &Result{Result: result}, nil
 }
 
 // Close waits until all the pending evaluations complete and then

internal/wasm/sdk/opa/opa_bench_test.go

@@ -0,0 +1,194 @@
+package opa_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/open-policy-agent/opa/internal/wasm/sdk/internal/wasm"
+	"github.com/open-policy-agent/opa/internal/wasm/sdk/opa"
+	"github.com/open-policy-agent/opa/rego"
+	"github.com/open-policy-agent/opa/util/test"
+)
+
+func BenchmarkWasmRego(b *testing.B) {
+	policy := compileRegoToWasm("a = true", "data.p.a = x", false)
+	instance, _ := opa.New().
+		WithPolicyBytes(policy).
+		WithMemoryLimits(131070, 2*131070). // TODO: For some reason unlimited memory slows down the eval_ctx_new().
+		WithPoolSize(1).
+		Init()
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	ctx := context.Background()
+	var input interface{} = make(map[string]interface{})
+
+	for i := 0; i < b.N; i++ {
+		if _, err := instance.Eval(ctx, opa.EvalOpts{Input: &input}); err != nil {
+			panic(err)
+		}
+	}
+}
+
+func BenchmarkGoRego(b *testing.B) {
+	pq := compileRego(`package p
+
+a = true`, "data.p.a = x")
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	ctx := context.Background()
+	input := make(map[string]interface{})
+
+	for i := 0; i < b.N; i++ {
+		if _, err := pq.Eval(ctx, rego.EvalInput(input)); err != nil {
+			panic(err)
+		}
+	}
+}
+
+func BenchmarkWASMArrayIteration(b *testing.B) {
+	sizes := []int{10, 100, 1000, 10000}
+	for _, n := range sizes {
+		b.Run(fmt.Sprint(n), func(b *testing.B) {
+			benchmarkIteration(b, test.ArrayIterationBenchmarkModule(n))
+		})
+	}
+}
+
+func BenchmarkWASMSetIteration(b *testing.B) {
+	sizes := []int{10, 100, 1000, 10000}
+	for _, n := range sizes {
+		b.Run(fmt.Sprint(n), func(b *testing.B) {
+			benchmarkIteration(b, test.SetIterationBenchmarkModule(n))
+		})
+	}
+}
+
+func BenchmarkWASMObjectIteration(b *testing.B) {
+	sizes := []int{10, 100, 1000, 10000}
+	for _, n := range sizes {
+		b.Run(fmt.Sprint(n), func(b *testing.B) {
+			benchmarkIteration(b, test.ObjectIterationBenchmarkModule(n))
+		})
+	}
+}
+
+var r *opa.Result
+
+func benchmarkIteration(b *testing.B, module string) {
+	query := "data.test.main = x"
+	policy := compileRegoToWasm(module, query, false)
+
+	instance, err := opa.New().
+		WithPolicyBytes(policy).
+		WithMemoryLimits(2*wasm.PageSize, 47*wasm.PageSize).
+		WithPoolSize(1).
+		Init()
+	if err != nil {
+		b.Fatalf("init sdk: %v", err)
+	}
+
+	b.ResetTimer()
+	ctx := context.Background()
+	var input interface{} = make(map[string]interface{})
+
+	for i := 0; i < b.N; i++ {
+		r, err = instance.Eval(ctx, opa.EvalOpts{Input: &input})
+		if err != nil {
+			b.Fatalf("Unexpected query error: %v", err)
+		}
+		if string(r.Result) != `{{"x":true}}` {
+			b.Errorf("unexpected result: %s", string(r.Result))
+		}
+	}
+}
+
+func BenchmarkWASMLargeJSON(b *testing.B) {
+	for _, kv := range []struct{ key, val int }{
+		{10, 10},
+		{10, 100},
+		{10, 1000},
+		{10, 10000},
+		{100, 100},
+		{100, 1000},
+	} {
+		b.Run(fmt.Sprintf("%dx%d", kv.key, kv.val), func(b *testing.B) {
+			ctx := context.Background()
+			data := test.GenerateJSONBenchmarkData(kv.key, kv.val)
+
+			// Read data.values N times inside query.
+			query := "data.keys[_] = x; data.values = y"
+			policy := compileRegoToWasm("", query, false)
+
+			instance, err := opa.New().
+				WithPolicyBytes(policy).
+				WithDataJSON(data).
+				WithMemoryLimits(200*wasm.PageSize, 600*wasm.PageSize). // This is rather much
+				WithPoolSize(1).
+				Init()
+			if err != nil {
+				b.Fatalf("init sdk: %v", err)
+			}
+
+			b.ResetTimer()
+			var input interface{} = make(map[string]interface{})
+
+			for i := 0; i < b.N; i++ {
+				r, err = instance.Eval(ctx, opa.EvalOpts{Input: &input})
+				if err != nil {
+					b.Fatalf("Unexpected query error: %v", err)
+				}
+			}
+		})
+	}
+}
+
+func BenchmarkWASMVirtualDocs(b *testing.B) {
+	for _, kv := range []struct{ total, hit int }{
+		{1, 1},
+		{10, 1},
+		{100, 1},
+		{1000, 1},
+		{10, 10},
+		{100, 10},
+		{1000, 10},
+		{100, 100},
+		{1000, 100},
+		{1000, 1000},
+	} {
+		b.Run(fmt.Sprintf("total=%d/hit=%d", kv.total, kv.hit), func(b *testing.B) {
+			runVirtualDocsBenchmark(b, kv.total, kv.hit)
+		})
+	}
+}
+
+func runVirtualDocsBenchmark(b *testing.B, numTotalRules, numHitRules int) {
+	ctx := context.Background()
+	module, input := test.GenerateVirtualDocsBenchmarkData(numTotalRules, numHitRules)
+	query := "data.a.b.c.allow = x"
+
+	policy := compileRegoToWasm(module, query, false)
+
+	instance, err := opa.New().
+		WithPolicyBytes(policy).
+		WithMemoryLimits(8*wasm.PageSize, 8*wasm.PageSize).
+		WithPoolSize(1).
+		Init()
+	if err != nil {
+		b.Fatalf("init sdk: %v", err)
+	}
+
+	b.ResetTimer()
+	var inp interface{} = input
+
+	for i := 0; i < b.N; i++ {
+		r, err = instance.Eval(ctx, opa.EvalOpts{Input: &inp})
+		if err != nil {
+			b.Fatalf("Unexpected query error: %v", err)
+		}
+	}
+}

internal/wasm/sdk/opa/opa_test.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -17,6 +18,9 @@ import (
 	"github.com/open-policy-agent/opa/util"
 )
 
+// control dumping in this file
+const dump = true
+
 func TestOPA(t *testing.T) {
 	type Eval struct {
 		NewPolicy string
@@ -220,7 +224,7 @@ a = "c" { input > 2 }`,
 
 	for _, test := range tests {
 		t.Run(test.Description, func(t *testing.T) {
-			policy := compileRegoToWasm(test.Policy, test.Query)
+			policy := compileRegoToWasm(test.Policy, test.Query, dump)
 			data := []byte(test.Data)
 			if len(data) == 0 {
 				data = nil
@@ -240,14 +244,14 @@ a = "c" { input > 2 }`,
 			for _, eval := range test.Evals {
 				switch {
 				case eval.NewPolicy != "" && eval.NewData != "":
-					policy := compileRegoToWasm(eval.NewPolicy, test.Query)
+					policy := compileRegoToWasm(eval.NewPolicy, test.Query, dump)
 					data := parseJSON(eval.NewData)
 					if err := instance.SetPolicyData(policy, data); err != nil {
 						t.Errorf(err.Error())
 					}
 
 				case eval.NewPolicy != "":
-					policy := compileRegoToWasm(eval.NewPolicy, test.Query)
+					policy := compileRegoToWasm(eval.NewPolicy, test.Query, dump)
 					if err := instance.SetPolicy(policy); err != nil {
 						t.Errorf(err.Error())
 					}
@@ -349,51 +353,22 @@ func TestNamedEntrypoint(t *testing.T) {
 	}
 }
 
-func BenchmarkWasmRego(b *testing.B) {
-	policy := compileRegoToWasm("a = true", "data.p.a = x")
-	instance, _ := opa.New().
-		WithPolicyBytes(policy).
-		WithMemoryLimits(131070, 2*131070). // TODO: For some reason unlimited memory slows down the eval_ctx_new().
-		WithPoolSize(1).
-		Init()
-
-	b.ReportAllocs()
-	b.ResetTimer()
-
-	ctx := context.Background()
-	var input interface{} = make(map[string]interface{})
-
-	for i := 0; i < b.N; i++ {
-		if _, err := instance.Eval(ctx, opa.EvalOpts{Input: &input}); err != nil {
-			panic(err)
-		}
-	}
-}
-
-func BenchmarkGoRego(b *testing.B) {
-	pq := compileRego(`package p
-
-a = true`, "data.p.a = x")
-
-	b.ReportAllocs()
-	b.ResetTimer()
-
-	input := make(map[string]interface{})
-
-	for i := 0; i < b.N; i++ {
-		if _, err := pq.Eval(context.Background(), rego.EvalInput(input)); err != nil {
-			panic(err)
-		}
+// compileRegoToWasm is shared with the benchmarking functions in opa_bench_test.go;
+// those function use helpers shared with topdown_bench_test.go, and they all use
+// `package test` -- whereas the callers in this file don't provide the package at
+// all and assume it'll be `p`.
+func compileRegoToWasm(module string, query string, dump bool) []byte {
+	if !strings.HasPrefix(module, "package") {
+		module = fmt.Sprintf("package p\n%s", module)
 	}
-}
-
-func compileRegoToWasm(policy string, query string) []byte {
-	module := fmt.Sprintf("package p\n%s", policy)
-	cr, err := rego.New(
+	opts := []func(*rego.Rego){
 		rego.Query(query),
 		rego.Module("module.rego", module),
-		rego.Dump(os.Stderr),
-	).Compile(context.Background(), rego.CompilePartial(false))
+	}
+	if dump {
+		opts = append(opts, rego.Dump(os.Stderr))
+	}
+	cr, err := rego.New(opts...).Compile(context.Background(), rego.CompilePartial(false))
 	if err != nil {
 		panic(err)
 	}

rego/rego.go

@@ -1257,7 +1257,7 @@ func (r *Rego) Compile(ctx context.Context, opts ...CompileOption) (*CompileResu
 		}
 	} else {
 		var err error
-		// If creating a new transacation it should be closed before calling the
+		// If creating a new transaction it should be closed before calling the
 		// planner to avoid holding open the transaction longer than needed.
 		//
 		// TODO(tsandall): in future, planner could make use of store, in which