profiler: Disable local var plugging on trace

This adds support for the `topdown.CustomTracer` to the `Profiler` so
that we can signal via config to not plug local vars on trace events.

This improves the performance substantially in cases where there are
many/large local variables.

Fixes: #2245
Signed-off-by: Patrick East <east.patrick@gmail.com>

profiler/profiler.go

@@ -38,6 +38,13 @@ func (p *Profiler) Enabled() bool {
 	return true
 }
 
+// Config returns the standard Tracer configuration for the profiler
+func (p *Profiler) Config() topdown.TraceConfig {
+	return topdown.TraceConfig{
+		PlugLocalVars: false, // Event variable metadata is not required for the Profiler
+	}
+}
+
 // ReportByFile returns a profiler report for expressions grouped by the
 // file name. For each file the results are sorted by increasing row number.
 func (p *Profiler) ReportByFile() (report Report) {

profiler/profiler_bench_test.go

@@ -0,0 +1,78 @@
+// 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 profiler
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"testing"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/rego"
+)
+
+func BenchmarkProfilerBigLocalVar(b *testing.B) {
+	iterations := []int{1, 100, 1000}
+	vars := []int{1, 10}
+
+	for _, iterationCount := range iterations {
+		for _, varCount := range vars {
+			name := fmt.Sprintf("%dVars%dIterations", varCount, iterationCount)
+			b.Run(name, func(b *testing.B) {
+				profiler := New()
+				module := generateModule(varCount, iterationCount)
+
+				_, err := ast.ParseModule("test.rego", module)
+				if err != nil {
+					b.Fatal(err)
+				}
+
+				ctx := context.Background()
+
+				pq, err := rego.New(
+					rego.Module("test.rego", module),
+					rego.Query("data.test.p"),
+				).PrepareForEval(ctx)
+
+				if err != nil {
+					b.Fatal(err)
+				}
+
+				b.ResetTimer()
+
+				for i := 0; i < b.N; i++ {
+					b.StartTimer()
+					_, err = pq.Eval(ctx, rego.EvalTracer(profiler))
+					b.StopTimer()
+
+					if err != nil {
+						b.Fatal(err)
+					}
+				}
+			})
+		}
+	}
+}
+
+func generateModule(numVars int, dataSize int) string {
+	sb := strings.Builder{}
+	sb.WriteString(`package test
+
+p {
+	x := a
+	v := x[i]
+`)
+	for i := 0; i < numVars; i++ {
+		sb.WriteString(fmt.Sprintf("\tv%d := x[i+%d]\n", i, i))
+	}
+	sb.WriteString("\tfalse\n}\n")
+	sb.WriteString("\na := [\n")
+	for i := 0; i < dataSize; i++ {
+		sb.WriteString(fmt.Sprintf("\t%d,\n", i))
+	}
+	sb.WriteString("]\n")
+	return sb.String()
+}

profiler/profiler_test.go

@@ -7,6 +7,7 @@ package profiler
 import (
 	"context"
 	_ "encoding/json"
+	"reflect"
 	"testing"
 	"time"
 
@@ -479,3 +480,16 @@ allowed_operations = [
 
 	}
 }
+
+func TestProfilerTraceConfig(t *testing.T) {
+	ct := topdown.CustomTracer(New())
+	conf := ct.Config()
+
+	expected := topdown.TraceConfig{
+		PlugLocalVars: false,
+	}
+
+	if !reflect.DeepEqual(expected, conf) {
+		t.Fatalf("Expected config: %+v, got %+v", expected, conf)
+	}
+}