feat: implement handling of dispatch and caching for DispatchQueryPlan (#3079)

# Description

Implements caching and cache keys for DispatchQueryPlan. With this PR, the alias iterator results, for Check alone, are cached in the usual dispatch cache for SpiceDB.

This also introduces an alias-chain-collapse optimization, which reduces multiple consecutive Alias boundaries into a single boundary (to simplify dispatch)

Author: barakmich

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ## [Unreleased]
 ### Added
 - Add DispatchExecutor, a query plan executor that is Dispatch-aware and sends subproblems on Alias boundaries (https://github.com/authzed/spicedb/pull/3074)
+- Implement Dispatch caching for query plan execution (https://github.com/authzed/spicedb/pull/3079)
 
 ### Changed
 - Build: strip quarantine attribute for MacOS (https://github.com/authzed/spicedb/pull/3082)

internal/dispatch/caching/caching.go

@@ -405,8 +405,55 @@ func (cd *Dispatcher) DispatchLookupSubjects(req *v1.DispatchLookupSubjectsReque
 }
 
 func (cd *Dispatcher) DispatchQueryPlan(req *v1.DispatchQueryPlanRequest, stream dispatch.PlanStream) error {
-	// TODO: add caching logic
-	return cd.d.DispatchQueryPlan(req, stream)
+	switch req.Operation {
+	case v1.PlanOperation_PLAN_OPERATION_CHECK:
+		return cd.dispatchQueryPlanCheckCached(req, stream)
+	default:
+		// TODO: add caching for LookupResources and LookupSubjects
+		return cd.d.DispatchQueryPlan(req, stream)
+	}
+}
+
+func (cd *Dispatcher) dispatchQueryPlanCheckCached(req *v1.DispatchQueryPlanRequest, stream dispatch.PlanStream) error {
+	cd.checkTotalCounter.Inc()
+
+	requestKey, err := cd.keyHandler.PlanCheckCacheKey(stream.Context(), req)
+	if err != nil {
+		return err
+	}
+
+	if cachedPathRaw, found := cd.c.Get(requestKey); found {
+		var cachedPath v1.ResultPath
+		if err := cachedPath.UnmarshalVT(cachedPathRaw.([]byte)); err != nil {
+			return err
+		}
+		cd.checkFromCacheCounter.Inc()
+		return stream.Publish(&v1.DispatchQueryPlanResponse{
+			Paths: []*v1.ResultPath{&cachedPath},
+		})
+	}
+
+	// Cache miss — collect the streamed result to cache the path.
+	// Check produces at most one response containing a single ResultPath.
+	collecting := dispatch.NewCollectingDispatchStream[*v1.DispatchQueryPlanResponse](stream.Context())
+	if err := cd.d.DispatchQueryPlan(req, collecting); err != nil {
+		return err
+	}
+
+	for _, resp := range collecting.Results() {
+		if len(resp.Paths) > 0 {
+			pathBytes, err := resp.Paths[0].MarshalVT()
+			if err == nil {
+				cd.c.Set(requestKey, pathBytes, sliceSize(pathBytes))
+			}
+		}
+
+		if err := stream.Publish(resp); err != nil {
+			return err
+		}
+	}
+
+	return nil
 }
 
 func (cd *Dispatcher) Close() error {

internal/dispatch/graph/graph.go

@@ -13,6 +13,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
+	"github.com/authzed/spicedb/internal/caveats"
 	"github.com/authzed/spicedb/internal/dispatch"
 	"github.com/authzed/spicedb/internal/graph"
 	log "github.com/authzed/spicedb/internal/logging"
@@ -24,6 +25,8 @@ import (
 	"github.com/authzed/spicedb/pkg/middleware/nodeid"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
+	"github.com/authzed/spicedb/pkg/query"
+	"github.com/authzed/spicedb/pkg/schema/v2"
 	"github.com/authzed/spicedb/pkg/tuple"
 )
 
@@ -471,12 +474,166 @@ func (ld *localDispatcher) DispatchLookupSubjects(
 	)
 }
 
-// DispatchQueryPlan implements dispatch.Plan interface
+// DispatchQueryPlan implements dispatch.Plan interface.
+// It loads the schema, compiles the plan, finds the subtree by canonical key,
+// and executes it locally. The Impl method is called directly on the found
+// iterator to avoid re-triggering the executor's dispatch decision on the
+// same alias boundary that was already dispatched by the caller.
 func (ld *localDispatcher) DispatchQueryPlan(
 	req *v1.DispatchQueryPlanRequest,
 	stream dispatch.PlanStream,
 ) error {
-	return errors.New("DispatchQueryPlan not yet implemented")
+	ctx := stream.Context()
+
+	planCtx := req.PlanContext
+	if planCtx == nil {
+		return errors.New("DispatchQueryPlan: missing plan_context")
+	}
+
+	revision, err := ld.parseRevision(ctx, planCtx.Revision)
+	if err != nil {
+		return err
+	}
+
+	// Load schema at the requested revision.
+	// TODO: use cached compiled plans (Phase 7) instead of recompiling each time.
+	it, err := ld.findIteratorByCanonicalKey(ctx, revision, query.CanonicalKey(req.CanonicalKey))
+	if err != nil {
+		return err
+	}
+
+	// Build execution context with DispatchExecutor so nested alias boundaries
+	// in the subtree can re-dispatch through the full dispatch chain.
+	dl := datalayer.MustFromContext(ctx)
+	executor := dispatch.NewDispatchExecutor(ld, planCtx)
+	qctx := &query.Context{
+		Context:       ctx,
+		Executor:      executor,
+		Reader:        query.NewQueryDatastoreReader(dl.SnapshotReader(revision)),
+		CaveatRunner:  caveats.NewCaveatRunner(caveattypes.Default.TypeSet),
+		CaveatContext: dispatch.CaveatContextFromPlanContext(planCtx),
+	}
+
+	resource := query.Object{ObjectType: req.Resource.Namespace, ObjectID: req.Resource.ObjectId}
+	subject := query.ObjectAndRelation{
+		ObjectType: req.Subject.Namespace,
+		ObjectID:   req.Subject.ObjectId,
+		Relation:   req.Subject.Relation,
+	}
+
+	// Call Impl directly — the dispatch boundary has already been crossed.
+	switch req.Operation {
+	case v1.PlanOperation_PLAN_OPERATION_CHECK:
+		path, err := it.CheckImpl(qctx, resource, subject)
+		if err != nil {
+			return err
+		}
+		if path != nil {
+			return stream.Publish(&v1.DispatchQueryPlanResponse{
+				Paths: []*v1.ResultPath{dispatch.QueryPathToResultPath(path)},
+			})
+		}
+		return nil
+
+	case v1.PlanOperation_PLAN_OPERATION_LOOKUP_RESOURCES:
+		// TODO: implement caching for LookupResources results
+		pathSeq, err := it.IterResourcesImpl(qctx, subject, query.NoObjectFilter())
+		if err != nil {
+			return err
+		}
+		for path, err := range pathSeq {
+			if err != nil {
+				return err
+			}
+			if err := stream.Publish(&v1.DispatchQueryPlanResponse{
+				Paths: []*v1.ResultPath{dispatch.QueryPathToResultPath(path)},
+			}); err != nil {
+				return err
+			}
+		}
+		return nil
+
+	case v1.PlanOperation_PLAN_OPERATION_LOOKUP_SUBJECTS:
+		// TODO: implement caching for LookupSubjects results
+		pathSeq, err := it.IterSubjectsImpl(qctx, resource, query.NoObjectFilter())
+		if err != nil {
+			return err
+		}
+		for path, err := range pathSeq {
+			if err != nil {
+				return err
+			}
+			if err := stream.Publish(&v1.DispatchQueryPlanResponse{
+				Paths: []*v1.ResultPath{dispatch.QueryPathToResultPath(path)},
+			}); err != nil {
+				return err
+			}
+		}
+		return nil
+
+	default:
+		return fmt.Errorf("DispatchQueryPlan: unknown operation %v", req.Operation)
+	}
+}
+
+// findIteratorByCanonicalKey loads the schema at the given revision, compiles
+// all permissions, and returns the iterator subtree matching the canonical key.
+func (ld *localDispatcher) findIteratorByCanonicalKey(ctx context.Context, revision datastore.Revision, targetKey query.CanonicalKey) (query.Iterator, error) {
+	dl := datalayer.MustFromContext(ctx)
+	reader := dl.SnapshotReader(revision)
+
+	sr, err := reader.ReadSchema(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("DispatchQueryPlan: failed to read schema: %w", err)
+	}
+
+	nsDefs, err := sr.ListAllTypeDefinitions(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("DispatchQueryPlan: failed to list type definitions: %w", err)
+	}
+
+	caveatDefs, err := sr.ListAllCaveatDefinitions(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("DispatchQueryPlan: failed to list caveat definitions: %w", err)
+	}
+
+	fullSchema, err := schema.BuildSchemaFromDefinitions(
+		datastore.DefinitionsOf(nsDefs),
+		datastore.DefinitionsOf(caveatDefs),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("DispatchQueryPlan: failed to build schema: %w", err)
+	}
+
+	for nsName, def := range fullSchema.Definitions() {
+		for permName := range def.Permissions() {
+			co, err := query.BuildOutlineFromSchema(fullSchema, nsName, permName)
+			if err != nil {
+				continue
+			}
+			it, err := co.Compile()
+			if err != nil {
+				continue
+			}
+			if found := findByCanonicalKey(it, targetKey); found != nil {
+				return found, nil
+			}
+		}
+	}
+	return nil, fmt.Errorf("DispatchQueryPlan: no iterator found for canonical key %q", targetKey)
+}
+
+// findByCanonicalKey recursively searches an iterator tree for a node matching the key.
+func findByCanonicalKey(it query.Iterator, key query.CanonicalKey) query.Iterator {
+	if it.CanonicalKey() == key {
+		return it
+	}
+	for _, sub := range it.Subiterators() {
+		if found := findByCanonicalKey(sub, key); found != nil {
+			return found
+		}
+	}
+	return nil
 }
 
 func (ld *localDispatcher) Close() error {

internal/dispatch/keys/computed.go

@@ -20,11 +20,14 @@ type cachePrefix string
 // Define the various prefixes for the cache entries. These must *all* be unique and must *all*
 // also be placed into the cachePrefixes slice below.
 const (
-	checkViaRelationPrefix  cachePrefix = "cr"
-	checkViaCanonicalPrefix cachePrefix = "cc"
-	lookupPrefix            cachePrefix = "l"
-	expandPrefix            cachePrefix = "e"
-	lookupSubjectsPrefix    cachePrefix = "ls"
+	checkViaRelationPrefix    cachePrefix = "cr"
+	checkViaCanonicalPrefix   cachePrefix = "cc"
+	lookupPrefix              cachePrefix = "l"
+	expandPrefix              cachePrefix = "e"
+	lookupSubjectsPrefix      cachePrefix = "ls"
+	planCheckPrefix           cachePrefix = "pc"
+	planLookupResourcesPrefix cachePrefix = "pr"
+	planLookupSubjectsPrefix  cachePrefix = "ps"
 )
 
 var cachePrefixes = []cachePrefix{
@@ -33,6 +36,9 @@ var cachePrefixes = []cachePrefix{
 	lookupPrefix,
 	expandPrefix,
 	lookupSubjectsPrefix,
+	planCheckPrefix,
+	planLookupResourcesPrefix,
+	planLookupSubjectsPrefix,
 }
 
 // checkRequestToKey converts a check request into a cache key based on the relation
@@ -113,3 +119,31 @@ func lookupSubjectsRequestToKey(req *v1.DispatchLookupSubjectsRequest, option di
 		hashableIds(req.ResourceIds),
 	)
 }
+
+// planCheckRequestToKey converts a plan check request into a cache key
+func planCheckRequestToKey(req *v1.DispatchQueryPlanRequest, option dispatchCacheKeyHashComputeOption) DispatchCacheKey {
+	return dispatchCacheKeyHash(planCheckPrefix, req.PlanContext.Revision, option,
+		hashableString(req.CanonicalKey),
+		hashableOnr{req.Resource},
+		hashableOnr{req.Subject},
+		hashableContext{Struct: req.PlanContext.CaveatContext},
+	)
+}
+
+// planLookupResourcesRequestToKey converts a plan lookup resources request into a cache key
+func planLookupResourcesRequestToKey(req *v1.DispatchQueryPlanRequest, option dispatchCacheKeyHashComputeOption) DispatchCacheKey {
+	return dispatchCacheKeyHash(planLookupResourcesPrefix, req.PlanContext.Revision, option,
+		hashableString(req.CanonicalKey),
+		hashableOnr{req.Subject},
+		hashableContext{Struct: req.PlanContext.CaveatContext},
+	)
+}
+
+// planLookupSubjectsRequestToKey converts a plan lookup subjects request into a cache key
+func planLookupSubjectsRequestToKey(req *v1.DispatchQueryPlanRequest, option dispatchCacheKeyHashComputeOption) DispatchCacheKey {
+	return dispatchCacheKeyHash(planLookupSubjectsPrefix, req.PlanContext.Revision, option,
+		hashableString(req.CanonicalKey),
+		hashableOnr{req.Resource},
+		hashableContext{Struct: req.PlanContext.CaveatContext},
+	)
+}

internal/dispatch/keys/computed_test.go

@@ -522,6 +522,68 @@ var generatorFuncs = map[string]generatorFunc{
 			}
 	},
 
+	// Plan Check.
+	string(planCheckPrefix): func(
+		resourceIds []string,
+		subjectIds []string,
+		resourceRelation *core.RelationReference,
+		subjectRelation *core.RelationReference,
+		metadata *v1.ResolverMeta,
+	) (DispatchCacheKey, []string) {
+		return planCheckRequestToKey(&v1.DispatchQueryPlanRequest{
+				CanonicalKey: resourceRelation.Relation,
+				Resource:     ONR(resourceRelation.Namespace, resourceIds[0], resourceRelation.Relation),
+				Subject:      ONR(subjectRelation.Namespace, subjectIds[0], subjectRelation.Relation),
+				PlanContext:  &v1.PlanContext{Revision: metadata.AtRevision},
+			}, computeBothHashes), []string{
+				resourceRelation.Relation,
+				resourceRelation.Namespace,
+				resourceIds[0],
+				subjectRelation.Namespace,
+				subjectIds[0],
+				subjectRelation.Relation,
+			}
+	},
+
+	// Plan Lookup Resources.
+	string(planLookupResourcesPrefix): func(
+		resourceIds []string,
+		subjectIds []string,
+		resourceRelation *core.RelationReference,
+		subjectRelation *core.RelationReference,
+		metadata *v1.ResolverMeta,
+	) (DispatchCacheKey, []string) {
+		return planLookupResourcesRequestToKey(&v1.DispatchQueryPlanRequest{
+				CanonicalKey: resourceRelation.Relation,
+				Subject:      ONR(subjectRelation.Namespace, subjectIds[0], subjectRelation.Relation),
+				PlanContext:  &v1.PlanContext{Revision: metadata.AtRevision},
+			}, computeBothHashes), []string{
+				resourceRelation.Relation,
+				subjectRelation.Namespace,
+				subjectIds[0],
+				subjectRelation.Relation,
+			}
+	},
+
+	// Plan Lookup Subjects.
+	string(planLookupSubjectsPrefix): func(
+		resourceIds []string,
+		subjectIds []string,
+		resourceRelation *core.RelationReference,
+		subjectRelation *core.RelationReference,
+		metadata *v1.ResolverMeta,
+	) (DispatchCacheKey, []string) {
+		return planLookupSubjectsRequestToKey(&v1.DispatchQueryPlanRequest{
+				CanonicalKey: resourceRelation.Relation,
+				Resource:     ONR(resourceRelation.Namespace, resourceIds[0], resourceRelation.Relation),
+				PlanContext:  &v1.PlanContext{Revision: metadata.AtRevision},
+			}, computeBothHashes), []string{
+				resourceRelation.Relation,
+				resourceRelation.Namespace,
+				resourceIds[0],
+			}
+	},
+
 	// Lookup Subjects.
 	string(lookupSubjectsPrefix): func(
 		resourceIds []string,