chore(services): store observed counts in permissionsServer

internal/services/v1/permissions_queryplan.go

@@ -2,6 +2,7 @@
 
 import (
 	"context"
+	"sync"
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 
@@ -13,9 +14,57 @@
 	"github.com/authzed/spicedb/pkg/schema/v2"
 )
 
+// QueryPlanMetadata aggregates CountStats for iterator canonical keys across multiple queries.
+// This allows tracking which parts of query plans are used most frequently.
+type QueryPlanMetadata struct {
+	mu    sync.Mutex
+	stats map[query.CanonicalKey]query.CountStats // GUARDED_BY(mu)
+}
+
+// NewQueryPlanMetadata creates a new QueryPlanMetadata tracker.
+func NewQueryPlanMetadata() *QueryPlanMetadata {
+	return &QueryPlanMetadata{
+		stats: make(map[query.CanonicalKey]query.CountStats),
+	}
+}
+
+// MergeCountStats merges CountStats from a query execution into the aggregated metadata.
+func (m *QueryPlanMetadata) MergeCountStats(counts map[query.CanonicalKey]query.CountStats) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	for key, newStats := range counts {
+		existing := m.stats[key]
+		existing.CheckCalls += newStats.CheckCalls
+		existing.IterSubjectsCalls += newStats.IterSubjectsCalls
+		existing.IterResourcesCalls += newStats.IterResourcesCalls
+		existing.CheckResults += newStats.CheckResults
+		existing.IterSubjectsResults += newStats.IterSubjectsResults
+		existing.IterResourcesResults += newStats.IterResourcesResults
+		m.stats[key] = existing
+	}
+}
+
+// GetStats returns a copy of all aggregated stats.
+func (m *QueryPlanMetadata) GetStats() map[query.CanonicalKey]query.CountStats {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	result := make(map[query.CanonicalKey]query.CountStats, len(m.stats))
+	for k, v := range m.stats {
+		result[k] = v
+	}
+	return result
+}
+
 // checkPermissionWithQueryPlan executes a permission check using the query plan API.
 // This builds an iterator tree from the schema and executes it against the datastore.
 func (ps *permissionServer) checkPermissionWithQueryPlan(ctx context.Context, req *v1.CheckPermissionRequest) (*v1.CheckPermissionResponse, error) {
+	// Lazy initialization of QueryPlanMetadata if nil
+	if ps.queryPlanMetadata == nil {
+		ps.queryPlanMetadata = NewQueryPlanMetadata()
+	}
+
 	atRevision, checkedAt, err := consistency.RevisionFromContext(ctx)
 	if err != nil {
 		return nil, ps.rewriteError(ctx, err)
@@ -69,14 +118,16 @@
 		return nil, ps.rewriteError(ctx, err)
 	}
 
-	// Create query context with optional tracing
-	qctx := &query.Context{
-		Context:       ctx,
-		Executor:      query.LocalExecutor{},
-		Reader:        reader,
-		CaveatContext: caveatContext,
-		CaveatRunner:  caveatsimpl.NewCaveatRunner(ps.config.CaveatTypeSet),
-	}
+	// Create count observer to track query execution statistics
+	countObserver := query.NewCountObserver()
+
+	// Create query context with count observer
+	qctx := query.NewLocalContext(ctx,
+		query.WithReader(reader),
+		query.WithCaveatContext(caveatContext),
+		query.WithCaveatRunner(caveatsimpl.NewCaveatRunner(ps.config.CaveatTypeSet)),
+		query.WithObserver(countObserver),
+	)
 
 	// Execute the check
 	resource := query.Object{
@@ -101,6 +152,10 @@
 		return nil, ps.rewriteError(ctx, err)
 	}
 
+	// Merge count statistics into metadata after query completes
+	countStats := countObserver.GetStats()
+	ps.queryPlanMetadata.MergeCountStats(countStats)
+
 	resp := &v1.CheckPermissionResponse{
 		CheckedAt:         checkedAt,
 		Permissionship:    permissionship,

internal/services/v1/permissions_queryplan_test.go

@@ -0,0 +1,226 @@
+package v1
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/authzed/spicedb/pkg/query"
+)
+
+func TestNewQueryPlanMetadata(t *testing.T) {
+	metadata := NewQueryPlanMetadata()
+	require.NotNil(t, metadata)
+	require.NotNil(t, metadata.stats)
+	require.Empty(t, metadata.GetStats())
+}
+
+func TestQueryPlanMetadataMergeCountStats(t *testing.T) {
+	metadata := NewQueryPlanMetadata()
+	key1 := query.CanonicalKey("test-key-1")
+	key2 := query.CanonicalKey("test-key-2")
+
+	// First merge
+	counts1 := map[query.CanonicalKey]query.CountStats{
+		key1: {
+			CheckCalls:           5,
+			IterSubjectsCalls:    3,
+			IterResourcesCalls:   2,
+			CheckResults:         10,
+			IterSubjectsResults:  6,
+			IterResourcesResults: 4,
+		},
+		key2: {
+			CheckCalls:           2,
+			IterSubjectsCalls:    1,
+			IterResourcesCalls:   1,
+			CheckResults:         4,
+			IterSubjectsResults:  2,
+			IterResourcesResults: 2,
+		},
+	}
+
+	metadata.MergeCountStats(counts1)
+
+	stats := metadata.GetStats()
+	require.Equal(t, counts1[key1], stats[key1])
+	require.Equal(t, counts1[key2], stats[key2])
+
+	// Second merge - should accumulate
+	counts2 := map[query.CanonicalKey]query.CountStats{
+		key1: {
+			CheckCalls:           3,
+			IterSubjectsCalls:    2,
+			IterResourcesCalls:   1,
+			CheckResults:         6,
+			IterSubjectsResults:  4,
+			IterResourcesResults: 2,
+		},
+	}
+
+	metadata.MergeCountStats(counts2)
+
+	stats = metadata.GetStats()
+	expected := query.CountStats{
+		CheckCalls:           8,  // 5 + 3
+		IterSubjectsCalls:    5,  // 3 + 2
+		IterResourcesCalls:   3,  // 2 + 1
+		CheckResults:         16, // 10 + 6
+		IterSubjectsResults:  10, // 6 + 4
+		IterResourcesResults: 6,  // 4 + 2
+	}
+	require.Equal(t, expected, stats[key1])
+	require.Equal(t, counts1[key2], stats[key2]) // key2 unchanged
+}
+
+func TestQueryPlanMetadataGetStatsReturnsACopy(t *testing.T) {
+	metadata := NewQueryPlanMetadata()
+	key := query.CanonicalKey("test-key")
+
+	counts := map[query.CanonicalKey]query.CountStats{
+		key: {
+			CheckCalls:   5,
+			CheckResults: 10,
+		},
+	}
+	metadata.MergeCountStats(counts)
+
+	stats1 := metadata.GetStats()
+	stats2 := metadata.GetStats()
+
+	// Should be equal
+	require.Equal(t, stats1, stats2)
+
+	// Modifying the returned map should not affect the metadata
+	stats1[key] = query.CountStats{CheckCalls: 999}
+	stats3 := metadata.GetStats()
+
+	require.NotEqual(t, 999, stats3[key].CheckCalls)
+	require.Equal(t, 5, stats3[key].CheckCalls)
+}
+
+func TestQueryPlanMetadataConcurrency(t *testing.T) {
+	metadata := NewQueryPlanMetadata()
+	key1 := query.CanonicalKey("test-key-1")
+	key2 := query.CanonicalKey("test-key-2")
+
+	var wg sync.WaitGroup
+	iterations := 100
+
+	// Goroutine 1 - merges stats for key1
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < iterations; i++ {
+			metadata.MergeCountStats(map[query.CanonicalKey]query.CountStats{
+				key1: {CheckCalls: 1, CheckResults: 2},
+			})
+		}
+	}()
+
+	// Goroutine 2 - merges stats for key2
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < iterations; i++ {
+			metadata.MergeCountStats(map[query.CanonicalKey]query.CountStats{
+				key2: {IterSubjectsCalls: 1, IterSubjectsResults: 3},
+			})
+		}
+	}()
+
+	// Goroutine 3 - reads stats
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for i := 0; i < iterations; i++ {
+			_ = metadata.GetStats()
+		}
+	}()
+
+	wg.Wait()
+
+	// Verify final counts
+	stats := metadata.GetStats()
+	require.Equal(t, iterations, stats[key1].CheckCalls)
+	require.Equal(t, iterations*2, stats[key1].CheckResults)
+	require.Equal(t, iterations, stats[key2].IterSubjectsCalls)
+	require.Equal(t, iterations*3, stats[key2].IterSubjectsResults)
+}
+
+func TestQueryPlanMetadataEmptyMerge(t *testing.T) {
+	metadata := NewQueryPlanMetadata()
+
+	// Merging empty map should not cause issues
+	metadata.MergeCountStats(map[query.CanonicalKey]query.CountStats{})
+
+	stats := metadata.GetStats()
+	require.Empty(t, stats)
+}
+
+func TestQueryPlanMetadataMergeMultipleKeys(t *testing.T) {
+	metadata := NewQueryPlanMetadata()
+
+	// Merge multiple keys at once
+	counts := map[query.CanonicalKey]query.CountStats{
+		query.CanonicalKey("key1"): {CheckCalls: 1, CheckResults: 2},
+		query.CanonicalKey("key2"): {IterSubjectsCalls: 3, IterSubjectsResults: 4},
+		query.CanonicalKey("key3"): {IterResourcesCalls: 5, IterResourcesResults: 6},
+	}
+
+	metadata.MergeCountStats(counts)
+
+	stats := metadata.GetStats()
+	require.Len(t, stats, 3)
+	require.Equal(t, counts[query.CanonicalKey("key1")], stats[query.CanonicalKey("key1")])
+	require.Equal(t, counts[query.CanonicalKey("key2")], stats[query.CanonicalKey("key2")])
+	require.Equal(t, counts[query.CanonicalKey("key3")], stats[query.CanonicalKey("key3")])
+}
+
+func TestPermissionServerQueryPlanMetadataInitialization(t *testing.T) {
+	// Test that permissionServer properly initializes queryPlanMetadata
+	ps := &permissionServer{
+		queryPlanMetadata: nil, // Explicitly set to nil to test lazy init
+	}
+
+	require.Nil(t, ps.queryPlanMetadata)
+
+	// The lazy initialization would happen in checkPermissionWithQueryPlan,
+	// but we can't easily test that without a full integration test.
+	// Instead, we verify that NewQueryPlanMetadata works correctly
+	ps.queryPlanMetadata = NewQueryPlanMetadata()
+	require.NotNil(t, ps.queryPlanMetadata)
+	require.Empty(t, ps.queryPlanMetadata.GetStats())
+}
+
+func TestQueryPlanMetadataAccumulatesAcrossQueries(t *testing.T) {
+	metadata := NewQueryPlanMetadata()
+	key := query.CanonicalKey("test-iterator")
+
+	// Simulate first query execution
+	firstQueryStats := map[query.CanonicalKey]query.CountStats{
+		key: {
+			CheckCalls:   2,
+			CheckResults: 3,
+		},
+	}
+	metadata.MergeCountStats(firstQueryStats)
+
+	stats := metadata.GetStats()
+	require.Equal(t, 2, stats[key].CheckCalls)
+	require.Equal(t, 3, stats[key].CheckResults)
+
+	// Simulate second query execution - should accumulate
+	secondQueryStats := map[query.CanonicalKey]query.CountStats{
+		key: {
+			CheckCalls:   1,
+			CheckResults: 2,
+		},
+	}
+	metadata.MergeCountStats(secondQueryStats)
+
+	stats = metadata.GetStats()
+	require.Equal(t, 3, stats[key].CheckCalls, "Should accumulate: 2 + 1 = 3")
+	require.Equal(t, 5, stats[key].CheckResults, "Should accumulate: 3 + 2 = 5")
+}

internal/services/v1/relationships.go

@@ -172,6 +172,7 @@ func NewPermissionsServer(
 			dispatchChunkSize:    configWithDefaults.DispatchChunkSize,
 			caveatTypeSet:        configWithDefaults.CaveatTypeSet,
 		},
+		queryPlanMetadata: NewQueryPlanMetadata(),
 	}
 }
 
@@ -182,7 +183,8 @@ type permissionServer struct {
 	dispatch dispatch.Dispatcher
 	config   PermissionsServerConfig
 
-	bulkChecker *bulkChecker
+	bulkChecker       *bulkChecker
+	queryPlanMetadata *QueryPlanMetadata
 }
 
 func (ps *permissionServer) ReadRelationships(req *v1.ReadRelationshipsRequest, resp v1.PermissionsService_ReadRelationshipsServer) error {