perf(opa): cache prepared eval queries to avoid per-Eval PrepareForEval

opa/eval_prepared_test.go

@@ -0,0 +1,141 @@
+package opa
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/boostsecurityio/poutine/models"
+	"github.com/boostsecurityio/poutine/results"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// These tests pin the invariants that Eval's cached query plan must preserve:
+//
+//  1. The same query on the same Opa, evaluated with DIFFERENT inputs, must
+//     return per-input-correct results — input must stay per-eval and must never
+//     be baked into a cached plan.
+//  2. Concurrent Eval calls on a shared Opa (as AnalyzeOrg does across repo
+//     goroutines) must be race-free and correct.
+//  3. Config changes (WithConfig) and rule-set changes (Compile) must be
+//     reflected by subsequent Evals — a stale cached plan bound to the old
+//     compiler would be a regression.
+
+func TestEvalSameQueryDifferentInputs(t *testing.T) {
+	o, err := NewOpa(context.TODO(), &models.Config{Include: []models.ConfigInclude{}})
+	noOpaErrors(t, err)
+	ctx := context.TODO()
+
+	// Reuse the exact same query string across calls, varying only the input.
+	const q = `utils.job_uses_self_hosted_runner(input)`
+	cases := map[string]bool{
+		"self-hosted":   true,
+		"ubuntu-latest": false,
+		"random-name":   true,
+		"windows-2022":  false,
+	}
+	// Run each case twice, interleaved, to catch any input carried across calls.
+	for pass := 0; pass < 2; pass++ {
+		for runner, want := range cases {
+			var got bool
+			err := o.Eval(ctx, q, map[string]interface{}{"runs_on": []string{runner}}, &got)
+			noOpaErrors(t, err)
+			assert.Equal(t, want, got, "runner=%s pass=%d", runner, pass)
+		}
+	}
+}
+
+func TestEvalConcurrentSharedOpa(t *testing.T) {
+	o, err := NewOpa(context.TODO(), &models.Config{Include: []models.ConfigInclude{}})
+	noOpaErrors(t, err)
+	ctx := context.TODO()
+
+	const q = `utils.job_uses_self_hosted_runner(input)`
+	type tc struct {
+		runner string
+		want   bool
+	}
+	cases := []tc{
+		{"self-hosted", true},
+		{"ubuntu-latest", false},
+		{"random-name", true},
+		{"windows-2019", false},
+	}
+
+	var wg sync.WaitGroup
+	for i := 0; i < 64; i++ {
+		c := cases[i%len(cases)]
+		wg.Add(1)
+		go func(c tc) {
+			defer wg.Done()
+			var got bool
+			if err := o.Eval(ctx, q, map[string]interface{}{"runs_on": []string{c.runner}}, &got); err != nil {
+				t.Errorf("eval error: %v", err)
+				return
+			}
+			assert.Equal(t, c.want, got, "runner=%s", c.runner)
+		}(c)
+	}
+	wg.Wait()
+}
+
+func TestEvalReflectsRecompile(t *testing.T) {
+	o, err := NewOpa(context.TODO(), &models.Config{Include: []models.ConfigInclude{}})
+	noOpaErrors(t, err)
+	ctx := context.TODO()
+
+	const q = `data.rules.pr_runs_on_self_hosted.rule`
+
+	var before *results.Rule
+	err = o.Eval(ctx, q, nil, &before)
+	noOpaErrors(t, err)
+	require.NotNil(t, before)
+	assert.Equal(t, []interface{}{}, before.Config["allowed_runners"].Value)
+
+	// Change config (which rewrites the store) and recompile the modules.
+	err = o.WithConfig(ctx, &models.Config{
+		RulesConfig: map[string]map[string]interface{}{
+			"pr_runs_on_self_hosted": {"allowed_runners": []string{"self-hosted"}},
+		},
+	})
+	require.NoError(t, err)
+	require.NoError(t, o.Compile(ctx, nil, nil))
+
+	// The same query on the same Opa must now reflect the new config, not a
+	// stale cached result/plan.
+	var after *results.Rule
+	err = o.Eval(ctx, q, nil, &after)
+	noOpaErrors(t, err)
+	require.NotNil(t, after)
+	assert.Equal(t, []interface{}{"self-hosted"}, after.Config["allowed_runners"].Value)
+}
+
+// TestEvalRecompileInvalidatesCache specifically guards the cache-invalidation in
+// Compile. Unlike TestEvalReflectsRecompile (whose config change is read from the
+// store on every eval and so would pass even without invalidation), this changes
+// the *rule set*: it primes the cache with a query against a rule, then recompiles
+// with that rule skipped. A cached plan bound to the old compiler would still
+// resolve the rule; a correctly invalidated cache re-prepares and the rule is
+// gone. This test FAILS if the invalidation in Compile is removed.
+func TestEvalRecompileInvalidatesCache(t *testing.T) {
+	o, err := NewOpa(context.TODO(), &models.Config{Include: []models.ConfigInclude{}})
+	noOpaErrors(t, err)
+	ctx := context.TODO()
+
+	const q = `data.rules.pr_runs_on_self_hosted.rule`
+
+	// Prime the cache: the rule resolves.
+	var before *results.Rule
+	require.NoError(t, o.Eval(ctx, q, nil, &before))
+	require.NotNil(t, before)
+
+	// Recompile with the rule skipped → new compiler, rule removed from the set.
+	require.NoError(t, o.Compile(ctx, []string{"pr_runs_on_self_hosted"}, nil))
+
+	// With a correctly invalidated cache, the query now resolves to nothing and
+	// Eval reports the empty result set. A stale plan would still succeed.
+	var after *results.Rule
+	err = o.Eval(ctx, q, nil, &after)
+	require.Error(t, err, "skipped rule must not resolve after recompile; stale cache would still find it")
+}

opa/opa.go

@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"slices"
 	"strings"
+	"sync"
 
 	"github.com/boostsecurityio/poutine/models"
 	"github.com/open-policy-agent/opa/v1/ast"
@@ -37,6 +38,15 @@ type Opa struct {
 	Store               storage.Store
 	LoadPaths           []string
 	customEmbeddedRules []embeddedSource
+
+	// preparedMu guards prepared. prepared caches one compiled query plan
+	// (rego.PreparedEvalQuery) per query string so Eval does not re-run
+	// PrepareForEval — the query compile + evaluation-plan build — on every call.
+	// The cache is invalidated in Compile, which is the only thing that changes
+	// the rule set the plans are bound to (config lives in the store and is read
+	// fresh on every eval, so it needs no invalidation).
+	preparedMu sync.RWMutex
+	prepared   map[string]rego.PreparedEvalQuery
 }
 
 func NewOpa(ctx context.Context, config *models.Config) (*Opa, error) {
@@ -232,20 +242,26 @@ func (o *Opa) Compile(ctx context.Context, skip []string, allowed []string) erro
 	}
 
 	o.Compiler = compiler
+
+	// Cached plans are bound to the previous compiler (and thus the previous rule
+	// set); drop them so subsequent Evals re-prepare against the new compiler.
+	o.preparedMu.Lock()
+	o.prepared = nil
+	o.preparedMu.Unlock()
+
 	return nil
 }
 
 func (o *Opa) Eval(ctx context.Context, query string, input map[string]interface{}, result interface{}) error {
-	regoInstance := rego.New(
-		rego.Query(query),
-		rego.Compiler(o.Compiler),
-		rego.PrintHook(o),
-		rego.Input(input),
-		rego.Imports([]string{"data.poutine.utils"}),
-		rego.Store(o.Store),
-	)
+	pq, err := o.preparedQuery(ctx, query)
+	if err != nil {
+		return err
+	}
 
-	rs, err := regoInstance.Eval(ctx)
+	// Input is supplied per call via EvalInput; the cached plan holds no input,
+	// and each Eval opens a fresh store transaction, so config changes written by
+	// WithConfig are observed here without invalidating the cache.
+	rs, err := pq.Eval(ctx, rego.EvalInput(input))
 	if err != nil {
 		return err
 	}
@@ -263,6 +279,43 @@ func (o *Opa) Eval(ctx context.Context, query string, input map[string]interface
 	return json.Unmarshal(data, result)
 }
 
+// preparedQuery returns the cached compiled plan for query, preparing (and
+// caching) it on first use. The compiled plan is bound to the current compiler
+// and is safe for concurrent evaluation; Compile clears the cache when it swaps
+// the compiler out.
+func (o *Opa) preparedQuery(ctx context.Context, query string) (rego.PreparedEvalQuery, error) {
+	o.preparedMu.RLock()
+	pq, ok := o.prepared[query]
+	o.preparedMu.RUnlock()
+	if ok {
+		return pq, nil
+	}
+
+	pq, err := rego.New(
+		rego.Query(query),
+		rego.Compiler(o.Compiler),
+		rego.PrintHook(o),
+		rego.Imports([]string{"data.poutine.utils"}),
+		rego.Store(o.Store),
+	).PrepareForEval(ctx)
+	if err != nil {
+		return rego.PreparedEvalQuery{}, fmt.Errorf("failed to prepare query: %w", err)
+	}
+
+	o.preparedMu.Lock()
+	defer o.preparedMu.Unlock()
+	// Another goroutine may have prepared the same query while we were preparing;
+	// prefer the already-cached one so all callers share a single plan.
+	if existing, ok := o.prepared[query]; ok {
+		return existing, nil
+	}
+	if o.prepared == nil {
+		o.prepared = make(map[string]rego.PreparedEvalQuery)
+	}
+	o.prepared[query] = pq
+	return pq, nil
+}
+
 func Capabilities() (*ast.Capabilities, error) {
 	capabilites := &ast.Capabilities{}
 	err := json.Unmarshal(capabilitiesJson, capabilites)