feat(datastore): cost the stored-schema cache by estimated byte size

The stored-schema cache previously admitted every entry at a fixed cost of 1, so a
MaxCost configured in bytes (e.g. the server's 32MiB default, or the embedded
SchemaCacheMaxCostBytes) actually bounded the number of cached schema versions rather
than memory.

Cost each entry by ReadOnlyStoredSchema.EstimatedSize() instead: a rough schema byte
size plus, for each registered schema-derived cache kind, an estimate of the bytes it
adds once populated. Datastore readers pass the exact serialized size they just read;
other callers fall back to a cheap schema-text proxy, avoiding a full proto walk
(proto.Size is ~1.5ms on a 1k-definition schema). RegisterDerivedCache now takes a size
estimator, and the compiled-caveat cache estimates per-caveat CEL overhead.

Author: josephschorr

internal/caveats/schemacache.go

@@ -13,12 +13,33 @@ import (
 // and is discarded when the schema changes.
 var compiledCaveatCacheKey = datastore.NewDerivedCacheKey("caveats.compiled")
 
+// compiledCaveatOverheadBytes is a rough per-caveat estimate of the memory a compiled caveat
+// adds on top of its serialized expression. A compiled caveat embeds a built CEL environment
+// and program, which dwarf the serialized expression; this is a deliberately conservative,
+// order-of-magnitude figure used only for cache-cost budgeting.
+const compiledCaveatOverheadBytes = 8 * 1024
+
 func init() {
-	if err := datastore.RegisterDerivedCache(compiledCaveatCacheKey, func() any { return &CompiledCaveatCache{} }); err != nil {
+	if err := datastore.RegisterDerivedCache(compiledCaveatCacheKey, func() any { return &CompiledCaveatCache{} }, estimateCompiledCaveatCacheSize); err != nil {
 		spiceerrors.MustPanicf("failed to register compiled caveat cache: %v", err)
 	}
 }
 
+// estimateCompiledCaveatCacheSize roughly estimates the bytes the compiled-caveat cache adds for
+// the given schema when fully populated: per caveat, the serialized expression plus a fixed
+// overhead for its compiled CEL environment.
+func estimateCompiledCaveatCacheSize(s *datastore.ReadOnlyStoredSchema) int64 {
+	v1 := s.Get().GetV1()
+	if v1 == nil {
+		return 0
+	}
+	total := 0
+	for _, caveat := range v1.GetCaveatDefinitions() {
+		total += len(caveat.GetSerializedExpression()) + compiledCaveatOverheadBytes
+	}
+	return int64(total)
+}
+
 // 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

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/datalayer/hashcache.go

@@ -113,7 +113,14 @@ func (c *schemaHashCache) Set(schemaHash SchemaHash, schema *datastore.ReadOnlyS
 		schema: schema,
 	})
 
-	c.cache.Set(SchemaCacheKey(schemaHash), schema, 1)
+	// 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
 }
 

pkg/datastore/datastore.go

@@ -526,17 +526,31 @@ type RevisionedNamespace = RevisionedDefinition[*core.NamespaceDefinition]
 // when the schema changes. The underlying schema itself remains immutable; only the
 // concurrency-safe derived caches are populated on demand.
 type ReadOnlyStoredSchema struct {
-	schema  *core.StoredSchema
-	derived sync.Map // map[DerivedCacheKey]any
+	schema     *core.StoredSchema
+	schemaSize int64    // rough byte size of the schema, used as a cache-cost base (see EstimatedSize)
+	derived    sync.Map // map[DerivedCacheKey]any
 }
 
-// NewReadOnlyStoredSchema wraps a StoredSchema as read-only.
-// Returns nil if the provided schema is nil.
+// NewReadOnlyStoredSchema wraps a StoredSchema as read-only. Returns nil if the provided
+// schema is nil. The schema's byte size is estimated cheaply from its schema text; callers
+// that have the exact serialized size on hand (e.g. datastore readers that just unmarshaled
+// it) should prefer NewReadOnlyStoredSchemaWithSize.
 func NewReadOnlyStoredSchema(schema *core.StoredSchema) *ReadOnlyStoredSchema {
 	if schema == nil {
 		return nil
 	}
-	return &ReadOnlyStoredSchema{schema: schema}
+	return NewReadOnlyStoredSchemaWithSize(schema, len(schema.GetV1().GetSchemaText()))
+}
+
+// NewReadOnlyStoredSchemaWithSize wraps a StoredSchema as read-only, recording sizeBytes as a
+// rough byte size of the schema for cache-cost accounting (see EstimatedSize). Returns nil if
+// the provided schema is nil. sizeBytes need not be exact; the serialized length the schema
+// was read from is a good value.
+func NewReadOnlyStoredSchemaWithSize(schema *core.StoredSchema, sizeBytes int) *ReadOnlyStoredSchema {
+	if schema == nil {
+		return nil
+	}
+	return &ReadOnlyStoredSchema{schema: schema, schemaSize: int64(sizeBytes)}
 }
 
 // Get returns the underlying StoredSchema. Callers must not modify the returned value.

pkg/datastore/derivedcache.go

@@ -18,21 +18,51 @@ func NewDerivedCacheKey(name string) DerivedCacheKey { return DerivedCacheKey{na
 // Name returns the human-readable name of the key.
 func (k DerivedCacheKey) Name() string { return k.name }
 
-// derivedCacheFactories maps a DerivedCacheKey to a factory that builds an empty cache
-// instance. Registered once at init time via RegisterDerivedCache.
-var derivedCacheFactories sync.Map // map[DerivedCacheKey]func() any
+// derivedCacheRegistration bundles the lazy factory for a derived cache kind with an estimator
+// of how many bytes that cache adds, when fully populated, on top of the schema it hangs off.
+type derivedCacheRegistration struct {
+	factory   func() any
+	estimator func(*ReadOnlyStoredSchema) int64
+}
+
+// derivedCacheFactories maps a DerivedCacheKey to its registration. Registered once at init
+// time via RegisterDerivedCache.
+var derivedCacheFactories sync.Map // map[DerivedCacheKey]derivedCacheRegistration
 
-// RegisterDerivedCache registers a factory used to lazily build a schema-derived cache of the
-// given kind. The factory returns a fresh, empty cache and is invoked at most once per
-// ReadOnlyStoredSchema instance (i.e. once per schema version). Intended to be called from an
-// init() function. It returns an error if a factory is already registered for the key.
-func RegisterDerivedCache(key DerivedCacheKey, factory func() any) error {
-	if _, loaded := derivedCacheFactories.LoadOrStore(key, factory); loaded {
+// RegisterDerivedCache registers a derived cache kind. factory returns a fresh, empty cache and
+// is invoked at most once per ReadOnlyStoredSchema instance (i.e. once per schema version).
+// estimator returns a rough byte size that this cache adds, when populated, on top of the
+// schema; it is summed into the schema's cache cost (see ReadOnlyStoredSchema.EstimatedSize) and
+// may be nil to contribute nothing. Intended to be called from an init() function. It returns an
+// error if a kind is already registered for the key.
+func RegisterDerivedCache(key DerivedCacheKey, factory func() any, estimator func(*ReadOnlyStoredSchema) int64) error {
+	reg := derivedCacheRegistration{factory: factory, estimator: estimator}
+	if _, loaded := derivedCacheFactories.LoadOrStore(key, reg); loaded {
 		return fmt.Errorf("derived schema cache already registered for key %q", key.name)
 	}
 	return nil
 }
 
+// EstimatedSize returns a rough byte size for this stored schema: the schema's own size plus,
+// for every registered derived cache kind, that kind's estimate of the additional bytes it adds
+// when populated. It is intended as the cost when caching the schema, so the cache's max-cost
+// budget accounts for the derived caches the schema will accrete (compiled caveats, etc.), not
+// just the schema blob. The estimate is deliberately rough and conservative (it assumes every
+// kind will be populated).
+func (r *ReadOnlyStoredSchema) EstimatedSize() int64 {
+	if r == nil {
+		return 0
+	}
+	size := r.schemaSize
+	derivedCacheFactories.Range(func(_, v any) bool {
+		if reg := v.(derivedCacheRegistration); reg.estimator != nil {
+			size += reg.estimator(r)
+		}
+		return true
+	})
+	return size
+}
+
 // derivedCache returns the derived cache registered under key for this schema, building it
 // once (lazily) on first access. It returns an error if no factory is registered for key,
 // which indicates a programming error (a cache kind accessed without being registered at
@@ -41,11 +71,11 @@ func (r *ReadOnlyStoredSchema) derivedCache(key DerivedCacheKey) (any, error) {
 	if v, ok := r.derived.Load(key); ok {
 		return v, nil
 	}
-	factory, ok := derivedCacheFactories.Load(key)
+	reg, ok := derivedCacheFactories.Load(key)
 	if !ok {
 		return nil, spiceerrors.MustBugf("no derived schema cache registered for key %q", key.name)
 	}
-	built := factory.(func() any)()
+	built := reg.(derivedCacheRegistration).factory()
 	actual, _ := r.derived.LoadOrStore(key, built)
 	return actual, nil
 }

pkg/datastore/derivedcache_test.go

@@ -31,7 +31,7 @@ func TestDerivedCacheLazyAndShared(t *testing.T) {
 	require.NoError(t, datastore.RegisterDerivedCache(key, func() any {
 		built++
 		return &testCache{id: built}
-	}))
+	}, nil))
 
 	s := newStoredSchema()
 
@@ -62,6 +62,20 @@ func TestDerivedCacheUnregisteredKeyErrors(t *testing.T) {
 
 func TestDerivedCacheDuplicateRegistrationErrors(t *testing.T) {
 	key := uniqueDerivedCacheKey("dup")
-	require.NoError(t, datastore.RegisterDerivedCache(key, func() any { return &testCache{} }))
-	require.Error(t, datastore.RegisterDerivedCache(key, func() any { return &testCache{} }))
+	require.NoError(t, datastore.RegisterDerivedCache(key, func() any { return &testCache{} }, nil))
+	require.Error(t, datastore.RegisterDerivedCache(key, func() any { return &testCache{} }, nil))
+}
+
+func TestEstimatedSizeIncludesSchemaAndDerivedEstimators(t *testing.T) {
+	// Base size comes from the explicit byte size; registered estimators add on top. The
+	// registry is process-global, so assert against the delta rather than an absolute total.
+	s := datastore.NewReadOnlyStoredSchemaWithSize(&core.StoredSchema{}, 1000)
+	before := s.EstimatedSize()
+	require.GreaterOrEqual(t, before, int64(1000), "estimated size includes the schema byte size")
+
+	key := uniqueDerivedCacheKey("estimator")
+	require.NoError(t, datastore.RegisterDerivedCache(key, func() any { return &testCache{} },
+		func(*datastore.ReadOnlyStoredSchema) int64 { return 250 }))
+
+	require.Equal(t, before+250, s.EstimatedSize(), "a registered estimator adds to estimated size")
 }