From 14a29c91e331139ca69bf036376bf6f162735c3e Mon Sep 17 00:00:00 2001 From: Jimmy Zelinskie Date: Fri, 3 Jul 2026 11:12:10 -0700 Subject: [PATCH] feat: add query plan cache Introduces a queryPlanCache that caches parsed schemas and compiled outlines behind singleflight groups, replacing the per-request schema load and outline build in checkPermission, lookupResources, and lookupSubjects query plan paths. Key correctness considerations: - Bypass-sentinel schema hashes (transactions, dev mode, tests) skip both cache and singleflight to avoid coalescing requests that may read from different datastore revisions. - MaxCost accounts for the entryWeight overhead (key string length) so entries are not immediately evicted. - ApplyAdvisor clones CanonicalKeys before mutation so cached outlines are not corrupted by downstream advisor passes. --- internal/services/v1/permissions_queryplan.go | 86 +--- internal/services/v1/queryplan_cache.go | 151 +++++++ internal/services/v1/queryplan_cache_test.go | 426 ++++++++++++++++++ internal/services/v1/relationships.go | 2 + pkg/query/advisor.go | 26 +- 5 files changed, 598 insertions(+), 93 deletions(-) create mode 100644 internal/services/v1/queryplan_cache.go create mode 100644 internal/services/v1/queryplan_cache_test.go diff --git a/internal/services/v1/permissions_queryplan.go b/internal/services/v1/permissions_queryplan.go index c14753dea2..9915fd9295 100644 --- a/internal/services/v1/permissions_queryplan.go +++ b/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) } diff --git a/internal/services/v1/queryplan_cache.go b/internal/services/v1/queryplan_cache.go new file mode 100644 index 0000000000..e2387c1dc6 --- /dev/null +++ b/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() +} diff --git a/internal/services/v1/queryplan_cache_test.go b/internal/services/v1/queryplan_cache_test.go new file mode 100644 index 0000000000..ed7a869a8e --- /dev/null +++ b/internal/services/v1/queryplan_cache_test.go @@ -0,0 +1,426 @@ +package v1 + +import ( + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/authzed/spicedb/pkg/datalayer" + mock_datalayer "github.com/authzed/spicedb/pkg/datalayer/mocks" + "github.com/authzed/spicedb/pkg/datastore" + "github.com/authzed/spicedb/pkg/namespace" +) + +var ( + testUserDef = namespace.Namespace("user") + testDocDef = namespace.Namespace("document", + namespace.MustRelation("viewer", nil, namespace.AllowedRelation("user", "...")), + namespace.MustRelation("view", + namespace.Union(namespace.ComputedUserset("viewer")), + ), + ) + testTypeDefs = []datastore.RevisionedTypeDefinition{ + {Definition: testUserDef}, + {Definition: testDocDef}, + } +) + +func mockReaderExpectSchema(ctrl *gomock.Controller) *mock_datalayer.MockRevisionedReader { + sr := mock_datalayer.NewMockSchemaReader(ctrl) + sr.EXPECT().ListAllTypeDefinitions(gomock.Any()).Return(testTypeDefs, nil).AnyTimes() + sr.EXPECT().ListAllCaveatDefinitions(gomock.Any()).Return([]datastore.RevisionedCaveat{}, nil).AnyTimes() + + reader := mock_datalayer.NewMockRevisionedReader(ctrl) + reader.EXPECT().ReadSchema(gomock.Any()).Return(sr, nil).AnyTimes() + return reader +} + +func countingReader(ctrl *gomock.Controller, count *atomic.Int32) *mock_datalayer.MockRevisionedReader { + sr := mock_datalayer.NewMockSchemaReader(ctrl) + sr.EXPECT().ListAllTypeDefinitions(gomock.Any()).Return(testTypeDefs, nil).AnyTimes() + sr.EXPECT().ListAllCaveatDefinitions(gomock.Any()).Return([]datastore.RevisionedCaveat{}, nil).AnyTimes() + + reader := mock_datalayer.NewMockRevisionedReader(ctrl) + reader.EXPECT().ReadSchema(gomock.Any()).DoAndReturn(func(_ any) (datalayer.SchemaReader, error) { + count.Add(1) + return sr, nil + }).AnyTimes() + return reader +} + +func TestGetOrBuildSchema_CachesOnRealHash(t *testing.T) { + ctrl := gomock.NewController(t) + var readCount atomic.Int32 + reader := countingReader(ctrl, &readCount) + + c := newQueryPlanCache() + defer c.Close() + + hash := datalayer.SchemaHash("abc123") + ctx := t.Context() + + s1, err := c.getOrBuildSchema(ctx, reader, hash) + require.NoError(t, err) + require.NotNil(t, s1) + require.Equal(t, int32(1), readCount.Load()) + + s2, err := c.getOrBuildSchema(ctx, reader, hash) + require.NoError(t, err) + require.Equal(t, s1, s2, "second call should return cached schema") + require.Equal(t, int32(1), readCount.Load(), "ReadSchema should not be called again") +} + +func TestGetOrBuildSchema_BypassSentinelSkipsCache(t *testing.T) { + ctrl := gomock.NewController(t) + var readCount atomic.Int32 + reader := countingReader(ctrl, &readCount) + + c := newQueryPlanCache() + defer c.Close() + + ctx := t.Context() + + _, err := c.getOrBuildSchema(ctx, reader, datalayer.NoSchemaHashForTesting) + require.NoError(t, err) + require.Equal(t, int32(1), readCount.Load()) + + _, err = c.getOrBuildSchema(ctx, reader, datalayer.NoSchemaHashForTesting) + require.NoError(t, err) + require.Equal(t, int32(2), readCount.Load(), "bypass sentinel must not cache") +} + +func TestGetOrBuildSchema_DifferentHashesCacheSeparately(t *testing.T) { + ctrl := gomock.NewController(t) + var readCount atomic.Int32 + reader := countingReader(ctrl, &readCount) + + c := newQueryPlanCache() + defer c.Close() + + ctx := t.Context() + + s1, err := c.getOrBuildSchema(ctx, reader, "hash-a") + require.NoError(t, err) + + s2, err := c.getOrBuildSchema(ctx, reader, "hash-b") + require.NoError(t, err) + require.Equal(t, int32(2), readCount.Load()) + require.NotSame(t, s1, s2, "different hashes should produce independent entries") + + // Re-fetch both from cache. + _, err = c.getOrBuildSchema(ctx, reader, "hash-a") + require.NoError(t, err) + _, err = c.getOrBuildSchema(ctx, reader, "hash-b") + require.NoError(t, err) + require.Equal(t, int32(2), readCount.Load(), "both should be served from cache") +} + +func TestGetOrBuildOutline_CachesOnRealHash(t *testing.T) { + ctrl := gomock.NewController(t) + var readCount atomic.Int32 + reader := countingReader(ctrl, &readCount) + + c := newQueryPlanCache() + defer c.Close() + + hash := datalayer.SchemaHash("abc123") + ctx := t.Context() + + co1, err := c.getOrBuildOutline(ctx, reader, hash, "document", "view") + require.NoError(t, err) + require.Equal(t, int32(1), readCount.Load()) + + co2, err := c.getOrBuildOutline(ctx, reader, hash, "document", "view") + require.NoError(t, err) + require.Equal(t, co1, co2, "outline should be returned from cache") + require.Equal(t, int32(1), readCount.Load()) +} + +func TestGetOrBuildOutline_BypassSentinelSkipsCache(t *testing.T) { + ctrl := gomock.NewController(t) + var readCount atomic.Int32 + reader := countingReader(ctrl, &readCount) + + c := newQueryPlanCache() + defer c.Close() + + ctx := t.Context() + + _, err := c.getOrBuildOutline(ctx, reader, datalayer.NoSchemaHashInTransaction, "document", "view") + require.NoError(t, err) + + _, err = c.getOrBuildOutline(ctx, reader, datalayer.NoSchemaHashInTransaction, "document", "view") + require.NoError(t, err) + require.Equal(t, int32(2), readCount.Load(), "bypass sentinel must reload every time") +} + +func TestGetOrBuildOutline_DifferentPermissionsCacheSeparately(t *testing.T) { + ctrl := gomock.NewController(t) + reader := mockReaderExpectSchema(ctrl) + + c := newQueryPlanCache() + defer c.Close() + + hash := datalayer.SchemaHash("abc123") + ctx := t.Context() + + co1, err := c.getOrBuildOutline(ctx, reader, hash, "document", "view") + require.NoError(t, err) + + co2, err := c.getOrBuildOutline(ctx, reader, hash, "document", "viewer") + require.NoError(t, err) + require.NotEqual(t, co1, co2, "different permissions should have different outlines") +} + +func TestGetOrBuildSchema_ConcurrentSameHash(t *testing.T) { + ctrl := gomock.NewController(t) + var readCount atomic.Int32 + reader := countingReader(ctrl, &readCount) + + c := newQueryPlanCache() + defer c.Close() + + hash := datalayer.SchemaHash("concurrent-hash") + ctx := t.Context() + const goroutines = 20 + + var wg sync.WaitGroup + for range goroutines { + wg.Go(func() { + _, err := c.getOrBuildSchema(ctx, reader, hash) + // Cannot use require on a separate goroutines. Assert fails the + // test, but lets all the goroutines run before exit. + assert.NoError(t, err) + }) + } + wg.Wait() + + require.LessOrEqual(t, readCount.Load(), int32(2), + "singleflight should collapse most concurrent loads into one or two calls") +} + +func TestGetOrBuildSchema_ConcurrentBypassSentinel(t *testing.T) { + ctrl := gomock.NewController(t) + var readCount atomic.Int32 + reader := countingReader(ctrl, &readCount) + + c := newQueryPlanCache() + defer c.Close() + + ctx := t.Context() + const goroutines = 10 + + var wg sync.WaitGroup + for range goroutines { + wg.Go(func() { + _, err := c.getOrBuildSchema(ctx, reader, datalayer.NoSchemaHashForTesting) + // Cannot use require on a separate goroutines. Assert fails the + // test, but lets all the goroutines run before exit. + assert.NoError(t, err) + }) + } + wg.Wait() + + require.Equal(t, int32(goroutines), readCount.Load(), + "bypass sentinels must not be deduplicated — each call should load independently") +} + +func TestGetOrBuildOutline_SchemaIsCachedAcrossOutlineCalls(t *testing.T) { + ctrl := gomock.NewController(t) + var readCount atomic.Int32 + reader := countingReader(ctrl, &readCount) + + c := newQueryPlanCache() + defer c.Close() + + hash := datalayer.SchemaHash("shared-schema") + ctx := t.Context() + + _, err := c.getOrBuildOutline(ctx, reader, hash, "document", "view") + require.NoError(t, err) + + _, err = c.getOrBuildOutline(ctx, reader, hash, "document", "viewer") + require.NoError(t, err) + + require.Equal(t, int32(1), readCount.Load(), + "schema should be cached and reused across different outline builds for the same hash") +} + +func TestAllBypassSentinelsSkipCache(t *testing.T) { + sentinels := []datalayer.SchemaHash{ + datalayer.NoSchemaHashInTransaction, + datalayer.NoSchemaHashInDevelopment, + datalayer.NoSchemaHashForTesting, + datalayer.NoSchemaHashForWatch, + datalayer.NoSchemaHashForLegacyCursor, + datalayer.NoSchemaHashInLegacyMode, + datalayer.NoSchemaHashInLegacyZedToken, + } + + for _, sentinel := range sentinels { + t.Run(string(sentinel), func(t *testing.T) { + ctrl := gomock.NewController(t) + var readCount atomic.Int32 + reader := countingReader(ctrl, &readCount) + + c := newQueryPlanCache() + defer c.Close() + + ctx := t.Context() + _, err := c.getOrBuildSchema(ctx, reader, sentinel) + require.NoError(t, err) + _, err = c.getOrBuildSchema(ctx, reader, sentinel) + require.NoError(t, err) + + require.Equal(t, int32(2), readCount.Load()) + }) + } +} + +func TestGetOrBuildOutline_InvalidPermission(t *testing.T) { + ctrl := gomock.NewController(t) + reader := mockReaderExpectSchema(ctrl) + + c := newQueryPlanCache() + defer c.Close() + + ctx := t.Context() + _, err := c.getOrBuildOutline(ctx, reader, "hash", "document", "nonexistent") + require.Error(t, err) +} + +func TestGetOrBuildOutline_InvalidResourceType(t *testing.T) { + ctrl := gomock.NewController(t) + reader := mockReaderExpectSchema(ctrl) + + c := newQueryPlanCache() + defer c.Close() + + ctx := t.Context() + _, err := c.getOrBuildOutline(ctx, reader, "hash", "nonexistent", "view") + require.Error(t, err) +} + +func TestLoadSchema_ReturnsValidSchema(t *testing.T) { + ctrl := gomock.NewController(t) + reader := mockReaderExpectSchema(ctrl) + + c := newQueryPlanCache() + defer c.Close() + + s, err := c.loadSchema(t.Context(), reader) + require.NoError(t, err) + require.NotNil(t, s) + + docDef, ok := s.GetTypeDefinition("document") + require.True(t, ok) + require.NotNil(t, docDef) + + userDef, ok := s.GetTypeDefinition("user") + require.True(t, ok) + require.NotNil(t, userDef) +} + +func TestGetOrBuildSchema_ErrorOnReadSchema(t *testing.T) { + ctrl := gomock.NewController(t) + + reader := mock_datalayer.NewMockRevisionedReader(ctrl) + reader.EXPECT().ReadSchema(gomock.Any()).Return(nil, errForTesting).AnyTimes() + + c := newQueryPlanCache() + defer c.Close() + + _, err := c.getOrBuildSchema(t.Context(), reader, "hash-err") + require.ErrorIs(t, err, errForTesting) +} + +func TestGetOrBuildSchema_ErrorOnListTypeDefinitions(t *testing.T) { + ctrl := gomock.NewController(t) + + sr := mock_datalayer.NewMockSchemaReader(ctrl) + sr.EXPECT().ListAllTypeDefinitions(gomock.Any()).Return(nil, errForTesting) + + reader := mock_datalayer.NewMockRevisionedReader(ctrl) + reader.EXPECT().ReadSchema(gomock.Any()).Return(sr, nil) + + c := newQueryPlanCache() + defer c.Close() + + _, err := c.getOrBuildSchema(t.Context(), reader, "hash-err") + require.ErrorIs(t, err, errForTesting) +} + +func TestGetOrBuildOutline_ErrorNotCached(t *testing.T) { + ctrl := gomock.NewController(t) + + // First call: reader returns an error. + sr := mock_datalayer.NewMockSchemaReader(ctrl) + sr.EXPECT().ListAllTypeDefinitions(gomock.Any()).Return(nil, errForTesting) + + errReader := mock_datalayer.NewMockRevisionedReader(ctrl) + errReader.EXPECT().ReadSchema(gomock.Any()).Return(sr, nil) + + c := newQueryPlanCache() + defer c.Close() + + _, err := c.getOrBuildOutline(t.Context(), errReader, "hash-err", "document", "view") + require.Error(t, err) +} + +func TestGetOrBuildOutline_UsesSchemaCache(t *testing.T) { + // Verifies that two outline lookups for different permissions on the same + // schema hash only call ReadSchema once (the schema layer caches it). + ctrl := gomock.NewController(t) + var readCount atomic.Int32 + + docDef := namespace.Namespace("document", + namespace.MustRelation("viewer", nil, namespace.AllowedRelation("user", "...")), + namespace.MustRelation("editor", nil, namespace.AllowedRelation("user", "...")), + namespace.MustRelation("view", + namespace.Union(namespace.ComputedUserset("viewer")), + ), + namespace.MustRelation("edit", + namespace.Union(namespace.ComputedUserset("editor")), + ), + ) + + tdefs := []datastore.RevisionedTypeDefinition{ + {Definition: namespace.Namespace("user")}, + {Definition: docDef}, + } + + sr := mock_datalayer.NewMockSchemaReader(ctrl) + sr.EXPECT().ListAllTypeDefinitions(gomock.Any()).Return(tdefs, nil).AnyTimes() + sr.EXPECT().ListAllCaveatDefinitions(gomock.Any()).Return([]datastore.RevisionedCaveat{}, nil).AnyTimes() + + // Override the reader to use our custom type defs while still counting. + customReader := mock_datalayer.NewMockRevisionedReader(ctrl) + customReader.EXPECT().ReadSchema(gomock.Any()).DoAndReturn(func(_ any) (datalayer.SchemaReader, error) { + readCount.Add(1) + return sr, nil + }).AnyTimes() + + c := newQueryPlanCache() + defer c.Close() + + hash := datalayer.SchemaHash("multi-perm") + ctx := t.Context() + + _, err := c.getOrBuildOutline(ctx, customReader, hash, "document", "view") + require.NoError(t, err) + + _, err = c.getOrBuildOutline(ctx, customReader, hash, "document", "edit") + require.NoError(t, err) + + require.Equal(t, int32(1), readCount.Load()) +} + +var errForTesting = &testError{} + +type testError struct{} + +func (e *testError) Error() string { return "test error" } diff --git a/internal/services/v1/relationships.go b/internal/services/v1/relationships.go index 32a621bf39..bc29fac67f 100644 --- a/internal/services/v1/relationships.go +++ b/internal/services/v1/relationships.go @@ -198,6 +198,7 @@ func NewPermissionsServer( caveatTypeSet: configWithDefaults.CaveatTypeSet, }, queryPlanMetadata: configWithDefaults.QueryPlanMetadata, + queryPlanCache: newQueryPlanCache(), } } @@ -210,6 +211,7 @@ type permissionServer struct { bulkChecker *bulkChecker queryPlanMetadata *query.QueryPlanMetadata + queryPlanCache *queryPlanCache } func (ps *permissionServer) ReadRelationships(req *v1.ReadRelationshipsRequest, resp v1.PermissionsService_ReadRelationshipsServer) error { diff --git a/pkg/query/advisor.go b/pkg/query/advisor.go index 5489d51703..33c27f79c7 100644 --- a/pkg/query/advisor.go +++ b/pkg/query/advisor.go @@ -1,6 +1,8 @@ package query import ( + "maps" + "github.com/authzed/spicedb/pkg/spiceerrors" ) @@ -40,22 +42,24 @@ type PlanAdvisor interface { // and a fully populated Hints map. func ApplyAdvisor(co CanonicalOutline, advisor PlanAdvisor) (CanonicalOutline, error) { // Phase 1: apply mutations bottom-up. - mutatedRoot, err := applyAdvisorMutations(co.Root, co, advisor) + mutatedRoot, extendedKeys, err := applyAdvisorMutations(co.Root, co, advisor) if err != nil { return CanonicalOutline{}, err } + mutatedCO := CanonicalOutline{ + Root: mutatedRoot, + CanonicalKeys: extendedKeys, + } + // Phase 2: collect hints using the post-mutation outline as the key source. hints := make(map[OutlineNodeID][]Hint) - if err := collectAdvisorHints(mutatedRoot, co, advisor, hints); err != nil { + if err := collectAdvisorHints(mutatedRoot, mutatedCO, advisor, hints); err != nil { return CanonicalOutline{}, err } - return CanonicalOutline{ - Root: mutatedRoot, - CanonicalKeys: co.CanonicalKeys, - Hints: hints, - }, nil + mutatedCO.Hints = hints + return mutatedCO, nil } // applyAdvisorMutations walks the outline tree bottom-up via WalkOutlineBottomUp. @@ -65,7 +69,7 @@ func ApplyAdvisor(co CanonicalOutline, advisor PlanAdvisor) (CanonicalOutline, e // After all mutations are applied, any newly synthesized nodes (ID==0) receive // fresh IDs via FillMissingNodeIDs, and their CanonicalKeys are recorded in // the provided keys map. -func applyAdvisorMutations(outline Outline, co CanonicalOutline, advisor PlanAdvisor) (Outline, error) { +func applyAdvisorMutations(outline Outline, co CanonicalOutline, advisor PlanAdvisor) (Outline, map[OutlineNodeID]CanonicalKey, error) { mutated, err := WalkOutlineBottomUp(outline, func(node Outline) (Outline, error) { mutations, err := advisor.GetMutations(node, co) if err != nil { @@ -85,9 +89,11 @@ func applyAdvisorMutations(outline Outline, co CanonicalOutline, advisor PlanAdv return result, nil }) if err != nil { - return Outline{}, err + return Outline{}, nil, err } - return FillMissingNodeIDs(mutated, co.CanonicalKeys), nil + + extendedKeys := maps.Clone(co.CanonicalKeys) + return FillMissingNodeIDs(mutated, extendedKeys), extendedKeys, nil } // collectAdvisorHints walks the outline tree pre-order via WalkOutlinePreOrder,