feat: in-process embedded permissions and per-schema-version caching (#3166)

Author: josephschorr

CHANGELOG.md

@@ -8,6 +8,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Schema: reads inside write transactions now use a cheap hash-only lookup (`schema_revision`) to check the cache before loading the full schema blob, reducing DB round-trips on cache hits (https://github.com/authzed/spicedb/pull/3160)
 - Updated the Prometheus buckets for `grpc_server_handling_seconds` and `spicedb_datastore_query_latency` to be able to correlate them (https://github.com/authzed/spicedb/pull/3188)
 - Use `testcontainers` instead of `ory/dockertest` for running containers in integration tests (https://github.com/authzed/spicedb/pull/2782)
+- Embedded: add `pkg/embedded`, an in-process library for running permission checks against a datastore via the dispatch engine, without standing up a gRPC server (https://github.com/authzed/spicedb/pull/3166)
+- Caveats: compiled caveats (and their CEL environments) are now cached per schema version — hung off the stored schema (`ReadOnlyStoredSchema`) and rebuilt only when the schema changes — rather than rebuilt on every check, reducing check cost for schemas with many caveats (https://github.com/authzed/spicedb/pull/3166)
 
 ### Fixed
 - Fixed a nil pointer dereference panic in `CheckBulkPermissions` that could occur under concurrent load when a tracing-enabled check shared a singleflight dispatch with a non-tracing bulk check. Debug-enabled checks are no longer singleflighted together with non-debug checks. (https://github.com/authzed/spicedb/pull/3174)

internal/caveats/run.go

@@ -51,11 +51,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 stored schema, which hosts
+// schema-derived caches such as the compiled-caveat cache.
+type cachedSchemaProvider interface {
+	StoredSchema() *datastore.ReadOnlyStoredSchema
+}
+
 // 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
+
+	// compiledCaveatCache, 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).
+	compiledCaveatCache *CompiledCaveatCache
 }
 
 // NewCaveatRunner creates a new CaveatRunner.
@@ -91,6 +104,21 @@ 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.compiledCaveatCache == nil {
+		if provider, ok := reader.(cachedSchemaProvider); ok {
+			if stored := provider.StoredSchema(); stored != nil {
+				compiledCaveatCache, err := CompiledCaveatCacheFor(stored)
+				if err != nil {
+					return err
+				}
+				cr.compiledCaveatCache = compiledCaveatCache
+			}
+		}
+	}
+
 	// Collect all referenced caveat definitions in the expression.
 	caveatNames := mapz.NewSet[string]()
 	collectCaveatNames(expr, caveatNames)
@@ -138,12 +166,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.compiledCaveatCache != nil {
+		justDeserialized, err = cr.compiledCaveatCache.GetOrCompile(caveatDefName, compile)
+	} else {
+		justDeserialized, err = compile()
+	}
 	if err != nil {
 		return caveat, nil, err
 	}

internal/caveats/run_test.go

@@ -731,6 +731,70 @@ func TestCaveatRunnerPopulateCaveatDefinitionsForExpr(t *testing.T) {
 	req.True(result.IsPartial())
 }
 
+// storedSchemaCaveatReader wraps a CaveatDefinitionLookup so it also advertises a stored
+// schema. That is what makes CaveatRunner discover and use the schema-tied compiled-caveat
+// cache (the cachedSchemaProvider path) rather than per-runner compilation.
+type storedSchemaCaveatReader struct {
+	CaveatDefinitionLookup
+	stored *datastore.ReadOnlyStoredSchema
+}
+
+func (r storedSchemaCaveatReader) StoredSchema() *datastore.ReadOnlyStoredSchema {
+	return r.stored
+}
+
+func TestCaveatRunnerUsesSchemaTiedCompiledCaveatCache(t *testing.T) {
+	req := require.New(t)
+
+	rawDS, err := dsfortesting.NewMemDBDatastoreForTesting(t, 0, 0, memdb.DisableGC)
+	req.NoError(err)
+
+	ds, _ := testfixtures.DatastoreFromSchemaAndTestRelationships(t, rawDS, `
+		caveat first_caveat(firstparam int) {
+			firstparam == 42
+		}
+		`, nil)
+
+	headRevisionResult, err := ds.HeadRevision(t.Context())
+	req.NoError(err)
+
+	dl := datalayer.NewDataLayer(ds)
+	sr, err := dl.SnapshotReader(headRevisionResult.Revision, datalayer.NoSchemaHashForTesting).ReadSchema(t.Context())
+	req.NoError(err)
+
+	// A stored schema is what a unified-schema reader exposes; wrapping the real reader with
+	// one drives the schema-tied cache path.
+	stored := datastore.NewReadOnlyStoredSchema(&core.StoredSchema{
+		VersionOneof: &core.StoredSchema_V1{
+			V1: &core.StoredSchema_V1StoredSchema{SchemaText: "caveat first_caveat(firstparam int) { firstparam == 42 }"},
+		},
+	})
+	reader := storedSchemaCaveatReader{CaveatDefinitionLookup: sr, stored: stored}
+
+	runner := NewCaveatRunner(types.Default.TypeSet)
+	expr := caveatexpr("first_caveat")
+
+	result, err := runner.RunCaveatExpression(t.Context(), expr,
+		map[string]any{"firstparam": int64(42)}, reader, RunCaveatExpressionNoDebugging)
+	req.NoError(err)
+	req.True(result.Value())
+
+	// The runner discovered the schema-tied cache...
+	req.NotNil(runner.compiledCaveatCache)
+	shared, err := CompiledCaveatCacheFor(stored)
+	req.NoError(err)
+	req.Same(runner.compiledCaveatCache, shared) // ...and it is THE cache hung off that schema.
+
+	// ...and the run populated it: the compiled caveat is now cached, so GetOrCompile
+	// returns it without recompiling (proving the cache is actually used across checks).
+	got, err := shared.GetOrCompile("first_caveat", func() (*pkgcaveats.CompiledCaveat, error) {
+		t.Fatal("compile should not be called; first_caveat should already be cached")
+		return nil, nil
+	})
+	req.NoError(err)
+	req.NotNil(got)
+}
+
 func TestCaveatRunnerEmptyExpression(t *testing.T) {
 	req := require.New(t)
 

internal/caveats/schemacache.go

@@ -0,0 +1,44 @@
+package caveats
+
+import (
+	"sync"
+
+	"github.com/authzed/spicedb/pkg/caveats"
+	"github.com/authzed/spicedb/pkg/datastore"
+)
+
+// compiledCaveatCacheKey identifies the schema-derived cache of compiled (deserialized) caveats
+// hung off a datastore.ReadOnlyStoredSchema. The cache is tied to a single stored-schema version
+// and is discarded when the schema changes.
+var compiledCaveatCacheKey = datastore.NewDerivedCacheKey[*CompiledCaveatCache]()
+
+// 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 stored schema,
+// building it lazily on first access.
+func CompiledCaveatCacheFor(s *datastore.ReadOnlyStoredSchema) (*CompiledCaveatCache, error) {
+	return datastore.LoadOrStoreDerived(s, compiledCaveatCacheKey, func() *CompiledCaveatCache {
+		return &CompiledCaveatCache{}
+	})
+}

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)
+}

internal/datastore/common/sqlschema.go

@@ -59,7 +59,8 @@ func (s *SQLSingleStoreSchemaReaderWriter[T]) ReadStoredSchema(ctx context.Conte
 		return nil, fmt.Errorf("failed to unmarshal schema: %w", err)
 	}
 
-	return datastore.NewReadOnlyStoredSchema(storedSchema), nil
+	// len(data) is the exact serialized size, used as the rough schema-size base for cache cost.
+	return datastore.NewReadOnlyStoredSchemaWithSize(storedSchema, len(data)), nil
 }
 
 // WriteStoredSchema writes the stored schema to the unified schema table.

internal/datastore/memdb/storedschema.go

@@ -39,7 +39,8 @@ func (r *memdbReader) ReadStoredSchema(_ context.Context) (*datastore.ReadOnlySt
 		return nil, fmt.Errorf("failed to unmarshal schema: %w", err)
 	}
 
-	return datastore.NewReadOnlyStoredSchema(storedSchema), nil
+	// len(sd.data) is the exact serialized size, used as the rough schema-size base for cache cost.
+	return datastore.NewReadOnlyStoredSchemaWithSize(storedSchema, len(sd.data)), nil
 }
 
 // assertSchemaHash verifies the stored schema hash matches expectedHash.

pkg/cmd/datastore/datastore.go

@@ -138,6 +138,9 @@ type Config struct {
 	BootstrapOverwrite    bool                 `debugmap:"visible"`
 	BootstrapTimeout      time.Duration        `debugmap:"visible"        default:"10s"`
 	CaveatTypeSet         *caveattypes.TypeSet `debugmap:"hidden"`
+	// BootstrapSchemaMode controls the schema storage mode used when writing bootstrap
+	// data. The zero value (SchemaModeReadLegacyWriteLegacy) preserves prior behavior.
+	BootstrapSchemaMode datalayer.SchemaMode `debugmap:"visible"`
 
 	// Hedging
 	RequestHedgingEnabled          bool          `debugmap:"visible"`
@@ -529,7 +532,7 @@ func NewDatastore(ctx context.Context, options ...ConfigOption) (datastore.Datas
 		}
 
 		if len(bootstrapContents) > 0 {
-			bootstrapDL := datalayer.NewDataLayer(ds)
+			bootstrapDL := datalayer.NewDataLayer(ds, datalayer.WithSchemaMode(opts.BootstrapSchemaMode))
 			_, _, err = validationfile.PopulateFromFilesContents(ctx, bootstrapDL, opts.CaveatTypeSet, bootstrapContents)
 			if err != nil {
 				return nil, fmt.Errorf("failed to load bootstrap data: %w", err)

pkg/cmd/datastore/zz_generated.options.go

@@ -113,7 +113,14 @@ func (c *schemaHashCache) Set(schemaHash SchemaHash, schema *datastore.ReadOnlyS
 		schema: schema,
 	})
 
-	c.cache.Set(SchemaCacheKey(schemaHash), schema, int64(schema.EstimatedSize()))
+	// Cost the entry by the schema's estimated byte size (schema blob plus the schema-derived
+	// caches it will accrete), so the cache's max-cost budget is in bytes. Floor at 1 so an
+	// (effectively empty) schema is still admitted with a non-zero weight.
+	cost := schema.EstimatedSize()
+	if cost < 1 {
+		cost = 1
+	}
+	c.cache.Set(SchemaCacheKey(schemaHash), schema, cost)
 	return nil
 }