Skip to content

Commit e95f9b9

Browse files
barakmichbarakmich
committed
chore(services): store observed counts in permissionsServer
1 parent a01eddc commit e95f9b9

File tree

3 files changed

+292
-9
lines changed

3 files changed

+292
-9
lines changed

‎internal/services/v1/permissions_queryplan.go‎

Lines changed: 63 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package v1
22

33
import (
44
"context"
5+
"sync"
56

67
v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
78

@@ -13,9 +14,57 @@ import (
1314
"github.com/authzed/spicedb/pkg/schema/v2"
1415
)
1516

17+
// QueryPlanMetadata aggregates CountStats for iterator canonical keys across multiple queries.
18+
// This allows tracking which parts of query plans are used most frequently.
19+
type QueryPlanMetadata struct {
20+
mu sync.Mutex
21+
stats map[query.CanonicalKey]query.CountStats // GUARDED_BY(mu)
22+
}
23+
24+
// NewQueryPlanMetadata creates a new QueryPlanMetadata tracker.
25+
func NewQueryPlanMetadata() *QueryPlanMetadata {
26+
return &QueryPlanMetadata{
27+
stats: make(map[query.CanonicalKey]query.CountStats),
28+
}
29+
}
30+
31+
// MergeCountStats merges CountStats from a query execution into the aggregated metadata.
32+
func (m *QueryPlanMetadata) MergeCountStats(counts map[query.CanonicalKey]query.CountStats) {
33+
m.mu.Lock()
34+
defer m.mu.Unlock()
35+
36+
for key, newStats := range counts {
37+
existing := m.stats[key]
38+
existing.CheckCalls += newStats.CheckCalls
39+
existing.IterSubjectsCalls += newStats.IterSubjectsCalls
40+
existing.IterResourcesCalls += newStats.IterResourcesCalls
41+
existing.CheckResults += newStats.CheckResults
42+
existing.IterSubjectsResults += newStats.IterSubjectsResults
43+
existing.IterResourcesResults += newStats.IterResourcesResults
44+
m.stats[key] = existing
45+
}
46+
}
47+
48+
// GetStats returns a copy of all aggregated stats.
49+
func (m *QueryPlanMetadata) GetStats() map[query.CanonicalKey]query.CountStats {
50+
m.mu.Lock()
51+
defer m.mu.Unlock()
52+
53+
result := make(map[query.CanonicalKey]query.CountStats, len(m.stats))
54+
for k, v := range m.stats {
55+
result[k] = v
56+
}
57+
return result
58+
}
59+
1660
// checkPermissionWithQueryPlan executes a permission check using the query plan API.
1761
// This builds an iterator tree from the schema and executes it against the datastore.
1862
func (ps *permissionServer) checkPermissionWithQueryPlan(ctx context.Context, req *v1.CheckPermissionRequest) (*v1.CheckPermissionResponse, error) {
63+
// Lazy initialization of QueryPlanMetadata if nil
64+
if ps.queryPlanMetadata == nil {
65+
ps.queryPlanMetadata = NewQueryPlanMetadata()
66+
}
67+
1968
atRevision, checkedAt, err := consistency.RevisionFromContext(ctx)
2069
if err != nil {
2170
return nil, ps.rewriteError(ctx, err)
@@ -69,14 +118,16 @@ func (ps *permissionServer) checkPermissionWithQueryPlan(ctx context.Context, re
69118
return nil, ps.rewriteError(ctx, err)
70119
}
71120

72-
// Create query context with optional tracing
73-
qctx := &query.Context{
74-
Context: ctx,
75-
Executor: query.LocalExecutor{},
76-
Reader: reader,
77-
CaveatContext: caveatContext,
78-
CaveatRunner: caveatsimpl.NewCaveatRunner(ps.config.CaveatTypeSet),
79-
}
121+
// Create count observer to track query execution statistics
122+
countObserver := query.NewCountObserver()
123+
124+
// Create query context with count observer
125+
qctx := query.NewLocalContext(ctx,
126+
query.WithReader(reader),
127+
query.WithCaveatContext(caveatContext),
128+
query.WithCaveatRunner(caveatsimpl.NewCaveatRunner(ps.config.CaveatTypeSet)),
129+
query.WithObserver(countObserver),
130+
)
80131

81132
// Execute the check
82133
resource := query.Object{
@@ -101,6 +152,10 @@ func (ps *permissionServer) checkPermissionWithQueryPlan(ctx context.Context, re
101152
return nil, ps.rewriteError(ctx, err)
102153
}
103154

155+
// Merge count statistics into metadata after query completes
156+
countStats := countObserver.GetStats()
157+
ps.queryPlanMetadata.MergeCountStats(countStats)
158+
104159
resp := &v1.CheckPermissionResponse{
105160
CheckedAt: checkedAt,
106161
Permissionship: permissionship,

‎internal/services/v1/permissions_queryplan_test.go‎

Lines changed: 226 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,226 @@
1+
package v1
2+
3+
import (
4+
"sync"
5+
"testing"
6+
7+
"github.com/stretchr/testify/require"
8+
9+
"github.com/authzed/spicedb/pkg/query"
10+
)
11+
12+
func TestNewQueryPlanMetadata(t *testing.T) {
13+
metadata := NewQueryPlanMetadata()
14+
require.NotNil(t, metadata)
15+
require.NotNil(t, metadata.stats)
16+
require.Empty(t, metadata.GetStats())
17+
}
18+
19+
func TestQueryPlanMetadataMergeCountStats(t *testing.T) {
20+
metadata := NewQueryPlanMetadata()
21+
key1 := query.CanonicalKey("test-key-1")
22+
key2 := query.CanonicalKey("test-key-2")
23+
24+
// First merge
25+
counts1 := map[query.CanonicalKey]query.CountStats{
26+
key1: {
27+
CheckCalls: 5,
28+
IterSubjectsCalls: 3,
29+
IterResourcesCalls: 2,
30+
CheckResults: 10,
31+
IterSubjectsResults: 6,
32+
IterResourcesResults: 4,
33+
},
34+
key2: {
35+
CheckCalls: 2,
36+
IterSubjectsCalls: 1,
37+
IterResourcesCalls: 1,
38+
CheckResults: 4,
39+
IterSubjectsResults: 2,
40+
IterResourcesResults: 2,
41+
},
42+
}
43+
44+
metadata.MergeCountStats(counts1)
45+
46+
stats := metadata.GetStats()
47+
require.Equal(t, counts1[key1], stats[key1])
48+
require.Equal(t, counts1[key2], stats[key2])
49+
50+
// Second merge - should accumulate
51+
counts2 := map[query.CanonicalKey]query.CountStats{
52+
key1: {
53+
CheckCalls: 3,
54+
IterSubjectsCalls: 2,
55+
IterResourcesCalls: 1,
56+
CheckResults: 6,
57+
IterSubjectsResults: 4,
58+
IterResourcesResults: 2,
59+
},
60+
}
61+
62+
metadata.MergeCountStats(counts2)
63+
64+
stats = metadata.GetStats()
65+
expected := query.CountStats{
66+
CheckCalls: 8, // 5 + 3
67+
IterSubjectsCalls: 5, // 3 + 2
68+
IterResourcesCalls: 3, // 2 + 1
69+
CheckResults: 16, // 10 + 6
70+
IterSubjectsResults: 10, // 6 + 4
71+
IterResourcesResults: 6, // 4 + 2
72+
}
73+
require.Equal(t, expected, stats[key1])
74+
require.Equal(t, counts1[key2], stats[key2]) // key2 unchanged
75+
}
76+
77+
func TestQueryPlanMetadataGetStatsReturnsACopy(t *testing.T) {
78+
metadata := NewQueryPlanMetadata()
79+
key := query.CanonicalKey("test-key")
80+
81+
counts := map[query.CanonicalKey]query.CountStats{
82+
key: {
83+
CheckCalls: 5,
84+
CheckResults: 10,
85+
},
86+
}
87+
metadata.MergeCountStats(counts)
88+
89+
stats1 := metadata.GetStats()
90+
stats2 := metadata.GetStats()
91+
92+
// Should be equal
93+
require.Equal(t, stats1, stats2)
94+
95+
// Modifying the returned map should not affect the metadata
96+
stats1[key] = query.CountStats{CheckCalls: 999}
97+
stats3 := metadata.GetStats()
98+
99+
require.NotEqual(t, 999, stats3[key].CheckCalls)
100+
require.Equal(t, 5, stats3[key].CheckCalls)
101+
}
102+
103+
func TestQueryPlanMetadataConcurrency(t *testing.T) {
104+
metadata := NewQueryPlanMetadata()
105+
key1 := query.CanonicalKey("test-key-1")
106+
key2 := query.CanonicalKey("test-key-2")
107+
108+
var wg sync.WaitGroup
109+
iterations := 100
110+
111+
// Goroutine 1 - merges stats for key1
112+
wg.Add(1)
113+
go func() {
114+
defer wg.Done()
115+
for i := 0; i < iterations; i++ {
116+
metadata.MergeCountStats(map[query.CanonicalKey]query.CountStats{
117+
key1: {CheckCalls: 1, CheckResults: 2},
118+
})
119+
}
120+
}()
121+
122+
// Goroutine 2 - merges stats for key2
123+
wg.Add(1)
124+
go func() {
125+
defer wg.Done()
126+
for i := 0; i < iterations; i++ {
127+
metadata.MergeCountStats(map[query.CanonicalKey]query.CountStats{
128+
key2: {IterSubjectsCalls: 1, IterSubjectsResults: 3},
129+
})
130+
}
131+
}()
132+
133+
// Goroutine 3 - reads stats
134+
wg.Add(1)
135+
go func() {
136+
defer wg.Done()
137+
for i := 0; i < iterations; i++ {
138+
_ = metadata.GetStats()
139+
}
140+
}()
141+
142+
wg.Wait()
143+
144+
// Verify final counts
145+
stats := metadata.GetStats()
146+
require.Equal(t, iterations, stats[key1].CheckCalls)
147+
require.Equal(t, iterations*2, stats[key1].CheckResults)
148+
require.Equal(t, iterations, stats[key2].IterSubjectsCalls)
149+
require.Equal(t, iterations*3, stats[key2].IterSubjectsResults)
150+
}
151+
152+
func TestQueryPlanMetadataEmptyMerge(t *testing.T) {
153+
metadata := NewQueryPlanMetadata()
154+
155+
// Merging empty map should not cause issues
156+
metadata.MergeCountStats(map[query.CanonicalKey]query.CountStats{})
157+
158+
stats := metadata.GetStats()
159+
require.Empty(t, stats)
160+
}
161+
162+
func TestQueryPlanMetadataMergeMultipleKeys(t *testing.T) {
163+
metadata := NewQueryPlanMetadata()
164+
165+
// Merge multiple keys at once
166+
counts := map[query.CanonicalKey]query.CountStats{
167+
query.CanonicalKey("key1"): {CheckCalls: 1, CheckResults: 2},
168+
query.CanonicalKey("key2"): {IterSubjectsCalls: 3, IterSubjectsResults: 4},
169+
query.CanonicalKey("key3"): {IterResourcesCalls: 5, IterResourcesResults: 6},
170+
}
171+
172+
metadata.MergeCountStats(counts)
173+
174+
stats := metadata.GetStats()
175+
require.Len(t, stats, 3)
176+
require.Equal(t, counts[query.CanonicalKey("key1")], stats[query.CanonicalKey("key1")])
177+
require.Equal(t, counts[query.CanonicalKey("key2")], stats[query.CanonicalKey("key2")])
178+
require.Equal(t, counts[query.CanonicalKey("key3")], stats[query.CanonicalKey("key3")])
179+
}
180+
181+
func TestPermissionServerQueryPlanMetadataInitialization(t *testing.T) {
182+
// Test that permissionServer properly initializes queryPlanMetadata
183+
ps := &permissionServer{
184+
queryPlanMetadata: nil, // Explicitly set to nil to test lazy init
185+
}
186+
187+
require.Nil(t, ps.queryPlanMetadata)
188+
189+
// The lazy initialization would happen in checkPermissionWithQueryPlan,
190+
// but we can't easily test that without a full integration test.
191+
// Instead, we verify that NewQueryPlanMetadata works correctly
192+
ps.queryPlanMetadata = NewQueryPlanMetadata()
193+
require.NotNil(t, ps.queryPlanMetadata)
194+
require.Empty(t, ps.queryPlanMetadata.GetStats())
195+
}
196+
197+
func TestQueryPlanMetadataAccumulatesAcrossQueries(t *testing.T) {
198+
metadata := NewQueryPlanMetadata()
199+
key := query.CanonicalKey("test-iterator")
200+
201+
// Simulate first query execution
202+
firstQueryStats := map[query.CanonicalKey]query.CountStats{
203+
key: {
204+
CheckCalls: 2,
205+
CheckResults: 3,
206+
},
207+
}
208+
metadata.MergeCountStats(firstQueryStats)
209+
210+
stats := metadata.GetStats()
211+
require.Equal(t, 2, stats[key].CheckCalls)
212+
require.Equal(t, 3, stats[key].CheckResults)
213+
214+
// Simulate second query execution - should accumulate
215+
secondQueryStats := map[query.CanonicalKey]query.CountStats{
216+
key: {
217+
CheckCalls: 1,
218+
CheckResults: 2,
219+
},
220+
}
221+
metadata.MergeCountStats(secondQueryStats)
222+
223+
stats = metadata.GetStats()
224+
require.Equal(t, 3, stats[key].CheckCalls, "Should accumulate: 2 + 1 = 3")
225+
require.Equal(t, 5, stats[key].CheckResults, "Should accumulate: 3 + 2 = 5")
226+
}

‎internal/services/v1/relationships.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,7 @@ func NewPermissionsServer(
172172
dispatchChunkSize: configWithDefaults.DispatchChunkSize,
173173
caveatTypeSet: configWithDefaults.CaveatTypeSet,
174174
},
175+
queryPlanMetadata: NewQueryPlanMetadata(),
175176
}
176177
}
177178

@@ -182,7 +183,8 @@ type permissionServer struct {
182183
dispatch dispatch.Dispatcher
183184
config PermissionsServerConfig
184185

185-
bulkChecker *bulkChecker
186+
bulkChecker *bulkChecker
187+
queryPlanMetadata *QueryPlanMetadata
186188
}
187189

188190
func (ps *permissionServer) ReadRelationships(req *v1.ReadRelationshipsRequest, resp v1.PermissionsService_ReadRelationshipsServer) error {

0 commit comments

Comments
 (0)