Benchmark fixes

A number of benchmarks did not have a `for range b.N` (or equivalent) loop in
them, leading to nothing being measured. This PR fixes that, along with some
cleanups in benchmarks found along the way.

Also remove `b.StopTimer` where not absolutely necessary, as that is
[notoriously buggy](golang/go#27217), and had some
benchmarks hang for a very long time.

Signed-off-by: Anders Eknert <anders@styra.com>

Author: anderseknert

v1/ast/strings_bench_test.go

@@ -5,7 +5,6 @@ import "testing"
 // BenchmarkTypeName-10    32207775	   38.93 ns/op    8 B/op    1 allocs/op
 func BenchmarkTypeName(b *testing.B) {
 	term := StringTerm("foo")
-	b.ResetTimer()
 
 	for range b.N {
 		name := TypeName(term.Value)
@@ -18,7 +17,6 @@ func BenchmarkTypeName(b *testing.B) {
 // BenchmarkValueName-10    508312227    2.374 ns/op    0 B/op    0 allocs/op
 func BenchmarkValueName(b *testing.B) {
 	term := StringTerm("foo")
-	b.ResetTimer()
 
 	for range b.N {
 		name := ValueName(term.Value)

v1/bundle/file_bench_test.go

@@ -52,31 +52,31 @@ func BenchmarkTarballLoader(b *testing.B) {
 			defer f.Close()
 
 			b.Run(strconv.Itoa(n), func(b *testing.B) {
-				// Reset the file reader.
-				if _, err := f.Seek(0, 0); err != nil {
-					b.Fatalf("Unexpected error: %s", err)
+				for range b.N {
+					// Reset the file reader.
+					if _, err := f.Seek(0, 0); err != nil {
+						b.Fatalf("Unexpected error: %s", err)
+					}
+					loader := NewTarballLoaderWithBaseURL(f, tarballFile)
+					benchTestLoader(b, loader)
 				}
-				loader := NewTarballLoaderWithBaseURL(f, tarballFile)
-				benchTestLoader(b, loader)
 			})
 		})
 	}
 }
 
 func BenchmarkDirectoryLoader(b *testing.B) {
-	sizes := []int{10000, 100000, 250000, 500000}
-
-	for _, n := range sizes {
+	for _, n := range []int{10000, 100000, 250000, 500000} {
 		expectedFiles := make(map[string]string, len(benchTestArchiveFiles)+1)
 		maps.Copy(expectedFiles, benchTestArchiveFiles)
 		expectedFiles["/x/data.json"] = benchTestGetFlatDataJSON(n)
 
 		test.WithTempFS(expectedFiles, func(rootDir string) {
 			b.ResetTimer()
-
 			b.Run(strconv.Itoa(n), func(b *testing.B) {
-				loader := NewDirectoryLoader(rootDir)
-				benchTestLoader(b, loader)
+				for range b.N {
+					benchTestLoader(b, NewDirectoryLoader(rootDir))
+				}
 			})
 		})
 	}

v1/compile/compile_bench_test.go

@@ -19,21 +19,19 @@ import (
 func BenchmarkCompileDynamicPolicy(b *testing.B) {
 	// This benchmarks the compiler against increasingly large numbers of dynamically-selected policies.
 	// See: https://github.com/open-policy-agent/opa/issues/5216
-
 	numPolicies := []int{1000, 2500, 5000, 7500, 10000}
 
 	for _, n := range numPolicies {
 		testcase := generateDynamicPolicyBenchmarkData(n)
 		test.WithTestFS(testcase, true, func(root string, fileSys fs.FS) {
 			b.ResetTimer()
 			b.Run(strconv.Itoa(n), func(b *testing.B) {
-				compiler := New().
-					WithFS(fileSys).
-					WithPaths(root)
+				for range b.N {
+					compiler := New().WithFS(fileSys).WithPaths(root)
 
-				err := compiler.Build(context.Background())
-				if err != nil {
-					b.Fatal("unexpected error", err)
+					if err := compiler.Build(context.Background()); err != nil {
+						b.Fatal("unexpected error", err)
+					}
 				}
 			})
 		})
@@ -83,12 +81,12 @@ func BenchmarkLargePartialRulePolicy(b *testing.B) {
 			test.WithTempFS(testcase, func(root string) {
 				b.ResetTimer()
 
-				compiler := New().
-					WithPaths(root)
+				for range b.N {
+					compiler := New().WithPaths(root)
 
-				err := compiler.Build(context.Background())
-				if err != nil {
-					b.Fatal("unexpected error", err)
+					if err := compiler.Build(context.Background()); err != nil {
+						b.Fatal("unexpected error", err)
+					}
 				}
 			})
 		})

v1/cover/cover_bench_test.go

@@ -22,11 +22,9 @@ func BenchmarkCoverBigLocalVar(b *testing.B) {
 		for _, varCount := range vars {
 			name := fmt.Sprintf("%dVars%dIterations", varCount, iterationCount)
 			b.Run(name, func(b *testing.B) {
-				cover := New()
 				module := generateModule(varCount, iterationCount)
 
-				_, err := ast.ParseModule("test.rego", module)
-				if err != nil {
+				if _, err := ast.ParseModule("test.rego", module); err != nil {
 					b.Fatal(err)
 				}
 
@@ -41,14 +39,12 @@ func BenchmarkCoverBigLocalVar(b *testing.B) {
 					b.Fatal(err)
 				}
 
+				cover := New()
+
 				b.ResetTimer()
 
 				for range b.N {
-					b.StartTimer()
-					_, err = pq.Eval(ctx, rego.EvalQueryTracer(cover))
-					b.StopTimer()
-
-					if err != nil {
+					if _, err = pq.Eval(ctx, rego.EvalQueryTracer(cover)); err != nil {
 						b.Fatal(err)
 					}
 				}

v1/dependencies/deps_bench_test.go

@@ -24,9 +24,10 @@ func BenchmarkBase(b *testing.B) {
 
 			b.ResetTimer()
 
-			_, err := Base(compiler, ref)
-			if err != nil {
-				b.Fatalf("Failed to compute base doc deps: %v", err)
+			for range b.N {
+				if _, err := Base(compiler, ref); err != nil {
+					b.Fatalf("Failed to compute base doc deps: %v", err)
+				}
 			}
 		})
 	}
@@ -36,8 +37,7 @@ func BenchmarkVirtual(b *testing.B) {
 	ruleCounts := []int{10, 20, 50}
 	for _, ruleCount := range ruleCounts {
 		b.Run(strconv.Itoa(ruleCount), func(b *testing.B) {
-			policy := makePolicy(ruleCount)
-			module := ast.MustParseModule(policy)
+			module := ast.MustParseModule(makePolicy(ruleCount))
 			compiler := ast.NewCompiler()
 			if compiler.Compile(map[string]*ast.Module{"test": module}); compiler.Failed() {
 				b.Fatalf("Failed to compile policy: %v", compiler.Errors)
@@ -47,9 +47,10 @@ func BenchmarkVirtual(b *testing.B) {
 
 			b.ResetTimer()
 
-			_, err := Virtual(compiler, ref)
-			if err != nil {
-				b.Fatalf("Failed to compute virtual doc deps: %v", err)
+			for range b.N {
+				if _, err := Virtual(compiler, ref); err != nil {
+					b.Fatalf("Failed to compute virtual doc deps: %v", err)
+				}
 			}
 		})
 	}

v1/metrics/metrics_bench_test.go

@@ -46,22 +46,16 @@ func BenchmarkMetricsMarshaling(b *testing.B) {
 			b.Fatalf("No output")
 		}
 	}
-
-	b.StopTimer()
 }
 
 func BenchmarkMetricsTimerStartStopRestart(b *testing.B) {
 	m := metrics.New()
 
-	b.ResetTimer()
-
 	for range b.N {
 		m.Timer("foo").Start()
 		_ = m.Timer("foo").Stop()
 		_ = m.Timer("foo").Stop() // Second stop to exercise the sync guard.
 		m.Timer("foo").Start()
 		_ = m.Timer("foo").Stop()
 	}
-
-	b.StopTimer()
 }

v1/profiler/profiler_bench_test.go

@@ -25,8 +25,7 @@ func BenchmarkProfilerBigLocalVar(b *testing.B) {
 				profiler := New()
 				module := generateModule(varCount, iterationCount)
 
-				_, err := ast.ParseModule("test.rego", module)
-				if err != nil {
+				if _, err := ast.ParseModule("test.rego", module); err != nil {
 					b.Fatal(err)
 				}
 
@@ -44,11 +43,7 @@ func BenchmarkProfilerBigLocalVar(b *testing.B) {
 				b.ResetTimer()
 
 				for range b.N {
-					b.StartTimer()
-					_, err = pq.Eval(ctx, rego.EvalQueryTracer(profiler))
-					b.StopTimer()
-
-					if err != nil {
+					if _, err = pq.Eval(ctx, rego.EvalQueryTracer(profiler)); err != nil {
 						b.Fatal(err)
 					}
 				}

v1/rego/rego_bench_test.go

@@ -344,7 +344,6 @@ r contains true if {
 			}
 
 			b.ResetTimer()
-			b.ReportAllocs()
 
 			for range b.N {
 				res, err := pq.Eval(ctx)
@@ -416,8 +415,7 @@ func BenchmarkTrivialPolicy(b *testing.B) {
 	}
 
 	for range b.N {
-		_, err := pq.Eval(ctx)
-		if err != nil {
+		if _, err := pq.Eval(ctx); err != nil {
 			b.Fatal(err)
 		}
 	}
@@ -439,8 +437,7 @@ func BenchmarkTrivialQuery(b *testing.B) {
 	}
 
 	for range b.N {
-		_, err := pq.Eval(ctx, EvalMetrics(m))
-		if err != nil {
+		if _, err := pq.Eval(ctx, EvalMetrics(m)); err != nil {
 			b.Fatal(err)
 		}
 	}
@@ -502,8 +499,7 @@ local_var if {
 	for i, pq := range []PreparedEvalQuery{pq1, pq2} {
 		b.Run(names[i], func(b *testing.B) {
 			for range b.N {
-				_, err := pq.Eval(ctx)
-				if err != nil {
+				if _, err := pq.Eval(ctx); err != nil {
 					b.Fatal(err)
 				}
 			}

v1/test/authz/authz_bench_test.go

@@ -46,76 +46,60 @@ func BenchmarkAuthzAllow1000Paths(b *testing.B) {
 }
 
 func runAuthzBenchmark(b *testing.B, mode InputMode, numPaths int, extras ...bool) {
-
 	profile := DataSetProfile{
 		NumTokens: 1000,
 		NumPaths:  numPaths,
 	}
 
 	ctx := context.Background()
 	data := GenerateDataset(profile)
-
-	useDisk := false
-	if len(extras) > 0 {
-		useDisk = extras[0]
-	}
+	useDisk := len(extras) > 0 && extras[0]
 
 	var store storage.Store
-	var err error
 	if useDisk {
-		rootDir := b.TempDir()
-		store, err = disk.New(ctx, logging.NewNoOpLogger(), nil, disk.Options{
-			Dir:        rootDir,
-			Partitions: nil,
-		})
-		if err != nil {
+		var err error
+		if store, err = disk.New(ctx, logging.NewNoOpLogger(), nil, disk.Options{Dir: b.TempDir()}); err != nil {
 			b.Fatal(err)
 		}
 
-		err = storage.WriteOne(ctx, store, storage.AddOp, storage.Path{}, data)
-		if err != nil {
+		if err = storage.WriteOne(ctx, store, storage.AddOp, storage.Path{}, data); err != nil {
 			b.Fatal(err)
 		}
 	} else {
-		store = inmem.NewFromObject(data)
+		store = inmem.NewFromObjectWithOpts(data)
 	}
 
-	txn := storage.NewTransactionOrDie(ctx, store)
 	compiler := ast.NewCompiler()
-	module := ast.MustParseModule(Policy)
-
-	compiler.Compile(map[string]*ast.Module{"": module})
-	if compiler.Failed() {
+	if compiler.Compile(map[string]*ast.Module{"": ast.MustParseModule(Policy)}); compiler.Failed() {
 		b.Fatalf("Unexpected error(s): %v", compiler.Errors)
 	}
 
 	r := rego.New(
 		rego.Compiler(compiler),
 		rego.Store(store),
-		rego.Transaction(txn),
+		rego.Transaction(storage.NewTransactionOrDie(ctx, store)),
 		rego.Query(AllowQuery),
 	)
 
-	// Pre-process as much as we can
 	pq, err := r.PrepareForEval(ctx)
 	if err != nil {
 		b.Fatalf("Unexpected error(s): %v", err)
 	}
 
 	input, expected := GenerateInput(profile, mode)
 
+	inputAST, err := ast.InterfaceToValue(input)
+	if err != nil {
+		b.Fatal(err)
+	}
+
 	b.ResetTimer()
 
 	for range b.N {
-		b.StartTimer()
-		rs, err := pq.Eval(
-			ctx,
-			rego.EvalInput(input),
-		)
+		rs, err := pq.Eval(ctx, rego.EvalParsedInput(inputAST))
 		if err != nil {
 			b.Fatalf("Unexpected error(s): %v", err)
 		}
-		b.StopTimer()
 
 		if rs.Allowed() != expected {
 			b.Fatalf("Unexpected result: %v", rs)

v1/test/authz/testing.go

@@ -17,16 +17,15 @@ import (
 // Policy is a test rego policy for a token based authz system
 const Policy = `package policy.restauthz
 
-import rego.v1
 import data.restauthz.tokens
 
-default allow = false
+default allow := false
 
 allow if {
-	tokens[input.token_id] = token
-	token.authz_profiles[_] = authz
+	token := tokens[input.token_id]
+	some authz in token.authz_profiles
 	regex.match(authz.path, input.path)
-	authz.methods[_] = input.method
+	input.method in authz.methods
 }`
 
 // AllowQuery is the test query that goes with the Policy
@@ -52,7 +51,6 @@ const (
 
 // GenerateInput will use a dataset profile and desired InputMode to generate inputs for testing
 func GenerateInput(profile DataSetProfile, mode InputMode) (any, any) {
-
 	var input string
 	var allow bool
 
@@ -132,11 +130,10 @@ func generateTokens(profile DataSetProfile) map[string]token {
 }
 
 func generateToken(profile DataSetProfile, i int) token {
-	token := token{
+	return token{
 		ID:            generateTokenID(i),
 		AuthzProfiles: generateAuthzProfiles(profile),
 	}
-	return token
 }
 
 func generateAuthzProfiles(profile DataSetProfile) []authzProfile {