Skip to content

Commit 69fd4ac

Browse files
ggreerggreer
authored
Grant expansion: Remove expandable annotation after processing expandable grant. (#601)
When compacting syncs, we spend a lot of time expanding grants that have already been expanded. So now we remove the expandable annotation after putting the grant in the graph. Also improved the grant expand test to validate that we get the correct number of grants and that expandable annotations are removed.
1 parent d6c929e commit 69fd4ac

File tree

3 files changed

+95
-50
lines changed

3 files changed

+95
-50
lines changed

‎pkg/sync/expand/expander.go‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,7 @@ func (e *Expander) runAction(ctx context.Context, action *EntitlementGraphAction
218218
return "", fmt.Errorf("runAction: error creating new grant: %w", err)
219219
}
220220
newGrants = append(newGrants, descendantGrant)
221-
newGrants, err = e.putGrantsInChunks(ctx, newGrants, 10000)
221+
newGrants, err = PutGrantsInChunks(ctx, e.store, newGrants, 10000)
222222
if err != nil {
223223
l.Error("runAction: error updating descendant grants", zap.Error(err))
224224
return "", fmt.Errorf("runAction: error updating descendant grants: %w", err)
@@ -255,7 +255,7 @@ func (e *Expander) runAction(ctx context.Context, action *EntitlementGraphAction
255255
}
256256
newGrants = append(newGrants, grantsToUpdate...)
257257

258-
newGrants, err = e.putGrantsInChunks(ctx, newGrants, 10000)
258+
newGrants, err = PutGrantsInChunks(ctx, e.store, newGrants, 10000)
259259
if err != nil {
260260
l.Error("runAction: error updating descendant grants", zap.Error(err))
261261
return "", fmt.Errorf("runAction: error updating descendant grants: %w", err)
@@ -268,7 +268,7 @@ func (e *Expander) runAction(ctx context.Context, action *EntitlementGraphAction
268268
}
269269
}
270270

271-
_, err = e.putGrantsInChunks(ctx, newGrants, 0)
271+
_, err = PutGrantsInChunks(ctx, e.store, newGrants, 0)
272272
if err != nil {
273273
l.Error("runAction: error updating descendant grants", zap.Error(err))
274274
return "", fmt.Errorf("runAction: error updating descendant grants: %w", err)
@@ -277,16 +277,16 @@ func (e *Expander) runAction(ctx context.Context, action *EntitlementGraphAction
277277
return sourceGrants.GetNextPageToken(), nil
278278
}
279279

280-
// putGrantsInChunks accumulates grants until the buffer exceeds minChunkSize,
280+
// PutGrantsInChunks accumulates grants until the buffer exceeds minChunkSize,
281281
// then writes all grants to the store at once.
282-
func (e *Expander) putGrantsInChunks(ctx context.Context, grants []*v2.Grant, minChunkSize int) ([]*v2.Grant, error) {
282+
func PutGrantsInChunks(ctx context.Context, store ExpanderStore, grants []*v2.Grant, minChunkSize int) ([]*v2.Grant, error) {
283283
if len(grants) < minChunkSize {
284284
return grants, nil
285285
}
286286

287-
err := e.store.PutGrants(ctx, grants...)
287+
err := store.PutGrants(ctx, grants...)
288288
if err != nil {
289-
return nil, fmt.Errorf("putGrantsInChunks: error putting grants: %w", err)
289+
return nil, fmt.Errorf("PutGrantsInChunks: error putting grants: %w", err)
290290
}
291291

292292
return make([]*v2.Grant, 0), nil

‎pkg/sync/syncer.go‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1658,11 +1658,40 @@ func (s *syncer) loadEntitlementGraph(ctx context.Context, graph *expand.Entitle
16581658
}
16591659

16601660
// Process grants and add edges to the graph
1661+
updatedGrants := make([]*v2.Grant, 0)
16611662
for _, grant := range resp.GetList() {
16621663
err := s.processGrantForGraph(ctx, grant, graph)
16631664
if err != nil {
16641665
return err
16651666
}
1667+
1668+
// Remove expandable annotation from descendant grant now that we've added it to the graph.
1669+
// That way if this sync is part of a compaction, expanding grants at the end of compaction won't redo work.
1670+
newAnnos := make(annotations.Annotations, 0)
1671+
updated := false
1672+
for _, anno := range grant.GetAnnotations() {
1673+
if anno.MessageIs(&v2.GrantExpandable{}) {
1674+
updated = true
1675+
} else {
1676+
newAnnos = append(newAnnos, anno)
1677+
}
1678+
}
1679+
if !updated {
1680+
continue
1681+
}
1682+
1683+
grant.SetAnnotations(newAnnos)
1684+
l.Debug("removed expandable annotation from grant", zap.String("grant_id", grant.GetId()))
1685+
updatedGrants = append(updatedGrants, grant)
1686+
updatedGrants, err = expand.PutGrantsInChunks(ctx, s.store, updatedGrants, 10000)
1687+
if err != nil {
1688+
return err
1689+
}
1690+
}
1691+
1692+
_, err = expand.PutGrantsInChunks(ctx, s.store, updatedGrants, 0)
1693+
if err != nil {
1694+
return err
16661695
}
16671696

16681697
if graph.Loaded {

‎pkg/sync/syncer_test.go‎

Lines changed: 59 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,7 @@ var userResourceType = v2.ResourceType_builder{
3636
}.Build()
3737

3838
func TestExpandGrants(t *testing.T) {
39-
ctx, cancel := context.WithCancel(context.Background())
40-
defer cancel()
39+
ctx := t.Context()
4140

4241
// 2500 * 4 = 10K - used to cause an infinite loop on pagition
4342
usersPerLayer := 2500
@@ -46,48 +45,79 @@ func TestExpandGrants(t *testing.T) {
4645
mc := newMockConnector()
4746

4847
mc.rtDB = append(mc.rtDB, groupResourceType, userResourceType)
49-
type asdf struct {
48+
type grantData struct {
5049
r *v2.Resource
5150
e *v2.Entitlement
5251
}
53-
groups := make([]*asdf, 0)
54-
for i := 0; i < groupCount; i++ {
52+
expectedGrantCount := 0
53+
groups := make([]*grantData, 0)
54+
for i := range groupCount {
5555
groupId := "group_" + strconv.Itoa(i)
5656
group, groupEnt, err := mc.AddGroup(ctx, groupId)
5757
for _, g := range groups {
58+
// Add previous groups to the new group. Make the entitlement expandable so that users in previous groups are added to the new group.
5859
_ = mc.AddGroupMember(ctx, g.r, group, groupEnt)
60+
expectedGrantCount++
5961
}
60-
groups = append(groups, &asdf{
62+
groups = append(groups, &grantData{
6163
r: group,
6264
e: groupEnt,
6365
})
6466
require.NoError(t, err)
6567

66-
for j := 0; j < usersPerLayer; j++ {
68+
for j := range usersPerLayer {
6769
pid := fmt.Sprintf("user_%d_%d_%d", i, usersPerLayer, j)
6870
principal, err := mc.AddUser(ctx, pid)
6971
require.NoError(t, err)
7072

7173
_ = mc.AddGroupMember(ctx, group, principal)
7274
}
75+
expectedGrantCount += usersPerLayer * (i + 1)
7376
}
7477

7578
tempDir, err := os.MkdirTemp("", "baton-benchmark-expand-grants")
7679
require.NoError(t, err)
7780
defer os.RemoveAll(tempDir)
7881
c1zpath := filepath.Join(tempDir, "expand-grants.c1z")
7982
syncer, err := NewSyncer(ctx, mc, WithC1ZPath(c1zpath), WithTmpDir(tempDir))
83+
defer func() { _ = os.Remove(c1zpath) }()
8084
require.NoError(t, err)
8185
err = syncer.Sync(ctx)
8286
require.NoError(t, err)
87+
8388
err = syncer.Close(ctx)
8489
require.NoError(t, err)
85-
_ = os.Remove(c1zpath)
90+
91+
// Validate that grants got expanded
92+
c1zManager, err := manager.New(ctx, c1zpath)
93+
require.NoError(t, err)
94+
store, err := c1zManager.LoadC1Z(ctx)
95+
require.NoError(t, err)
96+
97+
// Yes it's wasteful to load all grants into memory, but this connector doesn't make a ton of grants.
98+
allGrants := make([]*v2.Grant, 0)
99+
pageToken := ""
100+
for {
101+
resp, err := store.ListGrants(ctx, v2.GrantsServiceListGrantsRequest_builder{PageToken: pageToken}.Build())
102+
require.NoError(t, err)
103+
allGrants = append(allGrants, resp.GetList()...)
104+
pageToken = resp.GetNextPageToken()
105+
if pageToken == "" {
106+
break
107+
}
108+
}
109+
require.Len(t, allGrants, expectedGrantCount, "should have %d grants but got %d", expectedGrantCount, len(allGrants))
110+
for _, grant := range allGrants {
111+
annos := annotations.Annotations(grant.GetAnnotations())
112+
expandable := &v2.GrantExpandable{}
113+
ok, err := annos.Pick(expandable)
114+
require.NoError(t, err)
115+
require.False(t, ok, "grants are expanded, but grant %s has expandable annotation with entitlement ids %v", grant.GetId(), expandable.GetEntitlementIds())
116+
}
86117
}
87118

88119
func TestInvalidResourceTypeFilter(t *testing.T) {
89-
ctx, cancel := context.WithCancel(context.Background())
90-
defer cancel()
120+
ctx := t.Context()
91121

92122
tempDir, err := os.MkdirTemp("", "baton-invalid-resource-type-filter-sync-test")
93123
require.NoError(t, err)
@@ -114,8 +144,7 @@ func TestInvalidResourceTypeFilter(t *testing.T) {
114144
}
115145

116146
func TestResourceTypeFilter(t *testing.T) {
117-
ctx, cancel := context.WithCancel(context.Background())
118-
defer cancel()
147+
ctx := t.Context()
119148

120149
tempDir, err := os.MkdirTemp("", "baton-resource-type-filter-sync-test")
121150
require.NoError(t, err)
@@ -142,8 +171,7 @@ func TestResourceTypeFilter(t *testing.T) {
142171
}
143172

144173
func TestExpandGrantBadEntitlement(t *testing.T) {
145-
ctx, cancel := context.WithCancel(context.Background())
146-
defer cancel()
174+
ctx := t.Context()
147175

148176
ctx, err := logging.Init(ctx)
149177
require.NoError(t, err)
@@ -217,8 +245,7 @@ func TestExpandGrantBadEntitlement(t *testing.T) {
217245
}
218246

219247
func TestExpandGrantImmutable(t *testing.T) {
220-
ctx, cancel := context.WithCancel(context.Background())
221-
defer cancel()
248+
ctx := t.Context()
222249

223250
mc := newMockConnector()
224251

@@ -317,8 +344,7 @@ func TestExpandGrantImmutable(t *testing.T) {
317344
}
318345

319346
func TestExpandGrantImmutableCycle(t *testing.T) {
320-
ctx, cancel := context.WithCancel(context.Background())
321-
defer cancel()
347+
ctx := t.Context()
322348

323349
mc := newMockConnector()
324350

@@ -430,8 +456,7 @@ func TestExpandGrantImmutableCycle(t *testing.T) {
430456
_ = os.Remove(c1zpath)
431457
}
432458
func BenchmarkExpandCircle(b *testing.B) {
433-
ctx, cancel := context.WithCancel(context.Background())
434-
defer cancel()
459+
ctx := b.Context()
435460

436461
// create a loop of N entitlements
437462
circleSize := 7
@@ -443,7 +468,7 @@ func BenchmarkExpandCircle(b *testing.B) {
443468

444469
mc.rtDB = append(mc.rtDB, groupResourceType, userResourceType)
445470

446-
for i := 0; i < groupCount; i++ {
471+
for i := range groupCount {
447472
groupId := "group_" + strconv.Itoa(i)
448473
group, _, err := mc.AddGroup(ctx, groupId)
449474
require.NoError(b, err)
@@ -454,7 +479,7 @@ func BenchmarkExpandCircle(b *testing.B) {
454479

455480
_ = mc.AddGroupMember(ctx, group, childGroup, childEnt)
456481

457-
for j := 0; j < usersPerLayer; j++ {
482+
for j := range usersPerLayer {
458483
pid := "user_" + strconv.Itoa(i*usersPerLayer+j)
459484
principal, err := mc.AddUser(ctx, pid)
460485
require.NoError(b, err)
@@ -466,7 +491,7 @@ func BenchmarkExpandCircle(b *testing.B) {
466491
}
467492

468493
// create the circle
469-
for i := 0; i < circleSize; i += 1 {
494+
for i := range circleSize {
470495
groupId := "group_" + strconv.Itoa(i)
471496
nextGroupId := "group_" + strconv.Itoa((i+1)%circleSize) // Wrap around to the start for the last element
472497
currentEnt := mc.entDB[groupId][0]
@@ -479,8 +504,8 @@ func BenchmarkExpandCircle(b *testing.B) {
479504
require.NoError(b, err)
480505
defer os.RemoveAll(tempDir)
481506
c1zpath := filepath.Join(tempDir, "expand-circle.c1z")
482-
b.ResetTimer()
483-
for i := 0; i < b.N; i++ {
507+
508+
for b.Loop() {
484509
syncer, err := NewSyncer(ctx, mc, WithC1ZPath(c1zpath), WithTmpDir(tempDir))
485510
require.NoError(b, err)
486511
err = syncer.Sync(ctx)
@@ -492,8 +517,7 @@ func BenchmarkExpandCircle(b *testing.B) {
492517
}
493518

494519
func TestExternalResourcePath(t *testing.T) {
495-
ctx, cancel := context.WithCancel(context.Background())
496-
defer cancel()
520+
ctx := t.Context()
497521

498522
tempDir, err := os.MkdirTemp("", "baton-id-test")
499523
require.NoError(t, err)
@@ -596,8 +620,7 @@ func TestExternalResourcePath(t *testing.T) {
596620
}
597621

598622
func TestPartialSync(t *testing.T) {
599-
ctx, cancel := context.WithCancel(context.Background())
600-
defer cancel()
623+
ctx := t.Context()
601624

602625
tempDir, err := os.MkdirTemp("", "baton-partial-sync-test")
603626
require.NoError(t, err)
@@ -687,8 +710,7 @@ func TestPartialSync(t *testing.T) {
687710
}
688711

689712
func TestPartialSyncSkipEntitlementsAndGrants(t *testing.T) {
690-
ctx, cancel := context.WithCancel(context.Background())
691-
defer cancel()
713+
ctx := t.Context()
692714

693715
tempDir, err := os.MkdirTemp("", "baton-partial-sync-test-skip-entitlements-and-grants")
694716
require.NoError(t, err)
@@ -747,8 +769,7 @@ func TestPartialSyncSkipEntitlementsAndGrants(t *testing.T) {
747769
}
748770

749771
func TestPartialSyncUnimplemented(t *testing.T) {
750-
ctx, cancel := context.WithCancel(context.Background())
751-
defer cancel()
772+
ctx := t.Context()
752773

753774
tempDir, err := os.MkdirTemp("", "baton-partial-sync-test-unimplemented")
754775
require.NoError(t, err)
@@ -818,8 +839,7 @@ func TestPartialSyncUnimplemented(t *testing.T) {
818839
}
819840

820841
func TestExternalResourceMatchAll(t *testing.T) {
821-
ctx, cancel := context.WithCancel(context.Background())
822-
defer cancel()
842+
ctx := t.Context()
823843

824844
tempDir, err := os.MkdirTemp("", "baton-external-match-all-test")
825845
require.NoError(t, err)
@@ -913,8 +933,7 @@ func TestExternalResourceMatchAll(t *testing.T) {
913933
}
914934

915935
func TestExternalResourceMatchID(t *testing.T) {
916-
ctx, cancel := context.WithCancel(context.Background())
917-
defer cancel()
936+
ctx := t.Context()
918937

919938
tempDir, err := os.MkdirTemp("", "baton-external-match-id-test")
920939
require.NoError(t, err)
@@ -976,8 +995,7 @@ func TestExternalResourceMatchID(t *testing.T) {
976995
}
977996

978997
func TestExternalResourceEmailMatch(t *testing.T) {
979-
ctx, cancel := context.WithCancel(context.Background())
980-
defer cancel()
998+
ctx := t.Context()
981999

9821000
tempDir, err := os.MkdirTemp("", "baton-external-email-match-test")
9831001
require.NoError(t, err)
@@ -1065,8 +1083,7 @@ func TestExternalResourceEmailMatch(t *testing.T) {
10651083
}
10661084

10671085
func TestExternalResourceGroupProfileMatch(t *testing.T) {
1068-
ctx, cancel := context.WithCancel(context.Background())
1069-
defer cancel()
1086+
ctx := t.Context()
10701087

10711088
tempDir, err := os.MkdirTemp("", "baton-external-group-profile-match-test")
10721089
require.NoError(t, err)
@@ -1179,8 +1196,7 @@ func TestExternalResourceGroupProfileMatch(t *testing.T) {
11791196
}
11801197

11811198
func TestExternalResourceWithGrantToEntitlement(t *testing.T) {
1182-
ctx, cancel := context.WithCancel(context.Background())
1183-
defer cancel()
1199+
ctx := t.Context()
11841200

11851201
tempDir, err := os.MkdirTemp("", "baton-external-grant-to-entitlement-test")
11861202
require.NoError(t, err)

0 commit comments

Comments
 (0)