feat: add query plan cache

internal/services/v1/permissions_queryplan.go

@@ -9,11 +9,9 @@ import (
 	caveatsimpl "github.com/authzed/spicedb/internal/caveats"
 	"github.com/authzed/spicedb/internal/dispatch"
 	"github.com/authzed/spicedb/pkg/datalayer"
-	"github.com/authzed/spicedb/pkg/datastore"
 	"github.com/authzed/spicedb/pkg/middleware/consistency"
 	"github.com/authzed/spicedb/pkg/query"
 	"github.com/authzed/spicedb/pkg/query/queryopt"
-	"github.com/authzed/spicedb/pkg/schema/v2"
 	"github.com/authzed/spicedb/pkg/tuple"
 )
 
@@ -28,35 +26,7 @@ func (ps *permissionServer) checkPermissionWithQueryPlan(ctx context.Context, re
 	dl := datalayer.MustFromContext(ctx)
 	reader := dl.SnapshotReader(atRevision, schemaHash)
 
-	// Load all namespace and caveat definitions to build the schema
-	// TODO: Better schema caching
-	sr, err := reader.ReadSchema(ctx)
-	if err != nil {
-		return nil, ps.rewriteError(ctx, err)
-	}
-
-	namespaces, err := sr.ListAllTypeDefinitions(ctx)
-	if err != nil {
-		return nil, ps.rewriteError(ctx, err)
-	}
-
-	caveats, err := sr.ListAllCaveatDefinitions(ctx)
-	if err != nil {
-		return nil, ps.rewriteError(ctx, err)
-	}
-
-	// Build schema from definitions
-	fullSchema, err := schema.BuildSchemaFromDefinitions(
-		datastore.DefinitionsOf(namespaces),
-		datastore.DefinitionsOf(caveats),
-	)
-	if err != nil {
-		return nil, ps.rewriteError(ctx, err)
-	}
-
-	// Build and optimize the outline, then compile to an iterator tree.
-	// TODO: Better outline caching
-	co, err := query.BuildOutlineFromSchema(fullSchema, req.Resource.ObjectType, req.Permission)
+	co, err := ps.queryPlanCache.getOrBuildOutline(ctx, reader, schemaHash, req.Resource.ObjectType, req.Permission)
 	if err != nil {
 		return nil, ps.rewriteError(ctx, err)
 	}
@@ -173,32 +143,7 @@ func (ps *permissionServer) lookupResourcesWithQueryPlan(req *v1.LookupResources
 	dl := datalayer.MustFromContext(ctx)
 	reader := dl.SnapshotReader(atRevision, schemaHash)
 
-	// Load schema
-	sr, err := reader.ReadSchema(ctx)
-	if err != nil {
-		return ps.rewriteError(ctx, err)
-	}
-
-	namespaces, err := sr.ListAllTypeDefinitions(ctx)
-	if err != nil {
-		return ps.rewriteError(ctx, err)
-	}
-
-	caveats, err := sr.ListAllCaveatDefinitions(ctx)
-	if err != nil {
-		return ps.rewriteError(ctx, err)
-	}
-
-	fullSchema, err := schema.BuildSchemaFromDefinitions(
-		datastore.DefinitionsOf(namespaces),
-		datastore.DefinitionsOf(caveats),
-	)
-	if err != nil {
-		return ps.rewriteError(ctx, err)
-	}
-
-	// Build and compile the iterator tree for the requested resource type and permission
-	co, err := query.BuildOutlineFromSchema(fullSchema, req.ResourceObjectType, req.Permission)
+	co, err := ps.queryPlanCache.getOrBuildOutline(ctx, reader, schemaHash, req.ResourceObjectType, req.Permission)
 	if err != nil {
 		return ps.rewriteError(ctx, err)
 	}
@@ -308,32 +253,7 @@ func (ps *permissionServer) lookupSubjectsWithQueryPlan(req *v1.LookupSubjectsRe
 	dl := datalayer.MustFromContext(ctx)
 	reader := dl.SnapshotReader(atRevision, schemaHash)
 
-	// Load schema
-	sr, err := reader.ReadSchema(ctx)
-	if err != nil {
-		return ps.rewriteError(ctx, err)
-	}
-
-	namespaces, err := sr.ListAllTypeDefinitions(ctx)
-	if err != nil {
-		return ps.rewriteError(ctx, err)
-	}
-
-	caveats, err := sr.ListAllCaveatDefinitions(ctx)
-	if err != nil {
-		return ps.rewriteError(ctx, err)
-	}
-
-	fullSchema, err := schema.BuildSchemaFromDefinitions(
-		datastore.DefinitionsOf(namespaces),
-		datastore.DefinitionsOf(caveats),
-	)
-	if err != nil {
-		return ps.rewriteError(ctx, err)
-	}
-
-	// Build and compile the iterator tree for the resource type and permission
-	co, err := query.BuildOutlineFromSchema(fullSchema, req.Resource.ObjectType, req.Permission)
+	co, err := ps.queryPlanCache.getOrBuildOutline(ctx, reader, schemaHash, req.Resource.ObjectType, req.Permission)
 	if err != nil {
 		return ps.rewriteError(ctx, err)
 	}

internal/services/v1/queryplan_cache.go

@@ -0,0 +1,151 @@
+package v1
+
+import (
+	"context"
+
+	"resenje.org/singleflight"
+
+	"github.com/authzed/spicedb/pkg/cache"
+	"github.com/authzed/spicedb/pkg/datalayer"
+	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/query"
+	schema "github.com/authzed/spicedb/pkg/schema/v2"
+)
+
+type queryPlanCache struct {
+	schemas      cache.Cache[cache.StringKey, *schema.Schema]
+	schemaGroup  singleflight.Group[cache.StringKey, *schema.Schema]
+	outlines     cache.Cache[cache.StringKey, query.CanonicalOutline]
+	outlineGroup singleflight.Group[cache.StringKey, query.CanonicalOutline]
+}
+
+func newQueryPlanCache() *queryPlanCache {
+	// MaxCost must account for the per-entry weight added by the cache layer,
+	// which includes the key string length on top of the caller-supplied cost
+	// of 1. Schema keys are ~64-char hashes; outline keys are ~86+ chars.
+	schemas, err := cache.NewStandardCache[cache.StringKey, *schema.Schema](&cache.Config{
+		MaxCost: 16 * 100,
+	})
+	if err != nil {
+		panic("failed to create query plan schema cache: " + err.Error())
+	}
+
+	outlines, err := cache.NewStandardCache[cache.StringKey, query.CanonicalOutline](&cache.Config{
+		MaxCost: 256 * 100,
+	})
+	if err != nil {
+		panic("failed to create query plan outline cache: " + err.Error())
+	}
+
+	return &queryPlanCache{
+		schemas:  schemas,
+		outlines: outlines,
+	}
+}
+
+func (c *queryPlanCache) getOrBuildOutline(
+	ctx context.Context,
+	reader datalayer.RevisionedReader,
+	schemaHash datalayer.SchemaHash,
+	resourceType string,
+	permission string,
+) (query.CanonicalOutline, error) {
+	if schemaHash.IsBypassSentinel() {
+		return c.buildOutline(ctx, reader, schemaHash, resourceType, permission)
+	}
+
+	outlineKey := cache.StringKey(string(schemaHash) + ":" + resourceType + "#" + permission)
+	if co, ok := c.outlines.Get(outlineKey); ok {
+		return co, nil
+	}
+
+	co, _, err := c.outlineGroup.Do(ctx, outlineKey, func(ctx context.Context) (query.CanonicalOutline, error) {
+		if co, ok := c.outlines.Get(outlineKey); ok {
+			return co, nil
+		}
+		co, err := c.buildOutline(ctx, reader, schemaHash, resourceType, permission)
+		if err != nil {
+			return query.CanonicalOutline{}, err
+		}
+		c.outlines.Set(outlineKey, co, 1)
+		return co, nil
+	})
+	if err != nil {
+		return query.CanonicalOutline{}, err
+	}
+
+	return co, nil
+}
+
+func (c *queryPlanCache) buildOutline(
+	ctx context.Context,
+	reader datalayer.RevisionedReader,
+	schemaHash datalayer.SchemaHash,
+	resourceType string,
+	permission string,
+) (query.CanonicalOutline, error) {
+	fullSchema, err := c.getOrBuildSchema(ctx, reader, schemaHash)
+	if err != nil {
+		return query.CanonicalOutline{}, err
+	}
+	return query.BuildOutlineFromSchema(fullSchema, resourceType, permission)
+}
+
+func (c *queryPlanCache) getOrBuildSchema(
+	ctx context.Context,
+	reader datalayer.RevisionedReader,
+	schemaHash datalayer.SchemaHash,
+) (*schema.Schema, error) {
+	if schemaHash.IsBypassSentinel() {
+		return c.loadSchema(ctx, reader)
+	}
+
+	schemaKey := cache.StringKey(schemaHash)
+	if s, ok := c.schemas.Get(schemaKey); ok {
+		return s, nil
+	}
+
+	s, _, err := c.schemaGroup.Do(ctx, schemaKey, func(ctx context.Context) (*schema.Schema, error) {
+		if s, ok := c.schemas.Get(schemaKey); ok {
+			return s, nil
+		}
+		s, err := c.loadSchema(ctx, reader)
+		if err != nil {
+			return nil, err
+		}
+		c.schemas.Set(schemaKey, s, 1)
+		return s, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return s, nil
+}
+
+func (c *queryPlanCache) loadSchema(ctx context.Context, reader datalayer.RevisionedReader) (*schema.Schema, error) {
+	sr, err := reader.ReadSchema(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	namespaces, err := sr.ListAllTypeDefinitions(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	caveats, err := sr.ListAllCaveatDefinitions(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return schema.BuildSchemaFromDefinitions(
+		datastore.DefinitionsOf(namespaces),
+		datastore.DefinitionsOf(caveats),
+	)
+}
+
+func (c *queryPlanCache) Close() {
+	c.schemas.Close()
+	c.outlines.Close()
+}