feat(caveats): cache compiled caveats per schema version

ComputeCheck constructs a fresh CaveatRunner per call, so its per-runner
deserialized-caveat cache never persists across checks. Deserializing a
caveat rebuilds its CEL environment, which is expensive and, with a rich
caveat type set, dominated check cost.

Cache compiled caveats on the schema-derived cache (CachedSchema), keyed by
caveat name, so the CEL environment is built once per schema version rather
than once per check. Falls back to per-runner caching when the reader is not
backed by a unified stored schema.

Author: josephschorr

internal/caveats/run.go

@@ -11,6 +11,7 @@ import (
 	"github.com/authzed/spicedb/internal/telemetry/otelconv"
 	"github.com/authzed/spicedb/pkg/caveats"
 	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
+	"github.com/authzed/spicedb/pkg/datalayer"
 	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
@@ -51,11 +52,24 @@ func RunSingleCaveatExpression(
 	return runner.RunCaveatExpression(ctx, expr, context, reader, debugOption)
 }
 
+// cachedSchemaProvider is satisfied (structurally) by SchemaReaders backed by a unified
+// stored schema. It exposes the shared, per-schema-version CachedSchema, which hosts
+// schema-derived caches such as the compiled-caveat cache.
+type cachedSchemaProvider interface {
+	CachedSchema() *datalayer.CachedSchema
+}
+
 // CaveatRunner is a helper for running caveats, providing a cache for deserialized caveats.
 type CaveatRunner struct {
 	caveatTypeSet       *caveattypes.TypeSet
 	caveatDefs          map[string]*core.CaveatDefinition
 	deserializedCaveats map[string]*caveats.CompiledCaveat
+
+	// schemaCache, when non-nil, is a compiled-caveat cache tied to the stored schema
+	// (and thus shared across checks and invalidated on schema change). It is discovered
+	// from the reader on first use. When nil, deserializedCaveats provides per-runner
+	// caching only (the legacy behavior).
+	schemaCache *CompiledCaveatCache
 }
 
 // NewCaveatRunner creates a new CaveatRunner.
@@ -91,6 +105,17 @@ func (cr *CaveatRunner) PopulateCaveatDefinitionsForExpr(ctx context.Context, ex
 	ctx, span := tracer.Start(ctx, "PopulateCaveatDefinitions")
 	defer span.End()
 
+	// If the reader is backed by a unified stored schema, use the compiled-caveat cache
+	// tied to that schema version so deserialization (which rebuilds the CEL environment)
+	// is paid once per schema rather than once per check.
+	if cr.schemaCache == nil {
+		if provider, ok := reader.(cachedSchemaProvider); ok {
+			if cached := provider.CachedSchema(); cached != nil {
+				cr.schemaCache = CompiledCaveatCacheFor(cached)
+			}
+		}
+	}
+
 	// Collect all referenced caveat definitions in the expression.
 	caveatNames := mapz.NewSet[string]()
 	collectCaveatNames(expr, caveatNames)
@@ -138,12 +163,23 @@ func (cr *CaveatRunner) get(caveatDefName string) (*core.CaveatDefinition, *cave
 		return caveat, deserialized, nil
 	}
 
-	parameterTypes, err := caveattypes.DecodeParameterTypes(cr.caveatTypeSet, caveat.ParameterTypes)
-	if err != nil {
-		return nil, nil, err
+	compile := func() (*caveats.CompiledCaveat, error) {
+		parameterTypes, err := caveattypes.DecodeParameterTypes(cr.caveatTypeSet, caveat.ParameterTypes)
+		if err != nil {
+			return nil, err
+		}
+		return caveats.DeserializeCaveatWithTypeSet(cr.caveatTypeSet, caveat.SerializedExpression, parameterTypes)
 	}
 
-	justDeserialized, err := caveats.DeserializeCaveatWithTypeSet(cr.caveatTypeSet, caveat.SerializedExpression, parameterTypes)
+	// Prefer the schema-tied cache (shared across checks) when available; fall back to
+	// per-runner compilation otherwise.
+	var justDeserialized *caveats.CompiledCaveat
+	var err error
+	if cr.schemaCache != nil {
+		justDeserialized, err = cr.schemaCache.GetOrCompile(caveatDefName, compile)
+	} else {
+		justDeserialized, err = compile()
+	}
 	if err != nil {
 		return caveat, nil, err
 	}

internal/caveats/schemacache.go

@@ -0,0 +1,49 @@
+package caveats
+
+import (
+	"sync"
+
+	"github.com/authzed/spicedb/pkg/caveats"
+	"github.com/authzed/spicedb/pkg/datalayer"
+	"github.com/authzed/spicedb/pkg/spiceerrors"
+)
+
+// compiledCaveatCacheKey identifies the schema-derived cache of compiled (deserialized)
+// caveats. The cache is tied to a single stored-schema version (via datalayer.CachedSchema)
+// and is discarded when the schema changes.
+var compiledCaveatCacheKey = datalayer.NewDerivedCacheKey("caveats.compiled")
+
+func init() {
+	if err := datalayer.RegisterDerivedCache(compiledCaveatCacheKey, func() any { return &CompiledCaveatCache{} }); err != nil {
+		spiceerrors.MustPanicf("failed to register compiled caveat cache: %v", err)
+	}
+}
+
+// CompiledCaveatCache caches deserialized caveats (which embed a built CEL environment) by
+// caveat name, for a single schema version. Deserializing a caveat rebuilds its CEL
+// environment, which is expensive; caching it on the (shared) stored schema avoids paying
+// that cost on every check.
+type CompiledCaveatCache struct {
+	m sync.Map // map[string]*caveats.CompiledCaveat
+}
+
+// GetOrCompile returns the cached compiled caveat for name, or invokes compile and caches
+// the result. compile is only called on a miss; concurrent misses may call compile more
+// than once but only one result is retained.
+func (c *CompiledCaveatCache) GetOrCompile(name string, compile func() (*caveats.CompiledCaveat, error)) (*caveats.CompiledCaveat, error) {
+	if v, ok := c.m.Load(name); ok {
+		return v.(*caveats.CompiledCaveat), nil
+	}
+	compiled, err := compile()
+	if err != nil {
+		return nil, err
+	}
+	actual, _ := c.m.LoadOrStore(name, compiled)
+	return actual.(*caveats.CompiledCaveat), nil
+}
+
+// CompiledCaveatCacheFor returns the compiled-caveat cache tied to the given cached schema,
+// building it lazily on first access.
+func CompiledCaveatCacheFor(s *datalayer.CachedSchema) *CompiledCaveatCache {
+	return datalayer.GetDerivedCache[*CompiledCaveatCache](s, compiledCaveatCacheKey)
+}

internal/caveats/schemacache_test.go

@@ -0,0 +1,46 @@
+package caveats
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/authzed/spicedb/pkg/caveats"
+	caveattypes "github.com/authzed/spicedb/pkg/caveats/types"
+)
+
+func TestCompiledCaveatCacheGetOrCompile(t *testing.T) {
+	env := caveats.NewEnvironmentWithTypeSet(caveattypes.Default.TypeSet)
+	compiled, err := caveats.CompileCaveatWithName(env, "1 == 1", "test")
+	require.NoError(t, err)
+
+	c := &CompiledCaveatCache{}
+	calls := 0
+	compile := func() (*caveats.CompiledCaveat, error) {
+		calls++
+		return compiled, nil
+	}
+
+	// First access compiles; subsequent accesses return the cached instance without recompiling.
+	got1, err := c.GetOrCompile("a", compile)
+	require.NoError(t, err)
+	require.Same(t, compiled, got1)
+
+	got2, err := c.GetOrCompile("a", compile)
+	require.NoError(t, err)
+	require.Same(t, compiled, got2)
+	require.Equal(t, 1, calls, "compile should be invoked once per name")
+
+	// Distinct names compile independently.
+	_, err = c.GetOrCompile("b", compile)
+	require.NoError(t, err)
+	require.Equal(t, 2, calls)
+
+	// Errors propagate and are not cached.
+	boom := errors.New("boom")
+	_, err = c.GetOrCompile("c", func() (*caveats.CompiledCaveat, error) {
+		return nil, boom
+	})
+	require.ErrorIs(t, err, boom)
+}