Merge pull request #1568 from bradengroom/improve-grouping

Update groupItems to use less memory
authzed · Oct 11, 2023 · 7ac20aa · 7ac20aa
2 parents 786555c + 4f0692f
commit 7ac20aa
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 36 deletions.
diff --git a/internal/services/v1/experimental.go b/internal/services/v1/experimental.go
@@ -469,7 +469,7 @@ func (es *experimentalServer) BulkCheckPermission(ctx context.Context, req *v1.B
 		return nil
 	}
 
-	appendResultsForError := func(params computed.CheckParameters, resourceIDs []string, err error) error {
+	appendResultsForError := func(params *computed.CheckParameters, resourceIDs []string, err error) error {
 		rewritten := es.rewriteError(ctx, err)
 		statusResp, ok := status.FromError(rewritten)
 		if !ok {
@@ -499,7 +499,7 @@ func (es *experimentalServer) BulkCheckPermission(ctx context.Context, req *v1.B
 		return nil
 	}
 
-	appendResultsForCheck := func(params computed.CheckParameters, resourceIDs []string, metadata *dispatchv1.ResponseMeta, results map[string]*dispatchv1.ResourceCheckResult) error {
+	appendResultsForCheck := func(params *computed.CheckParameters, resourceIDs []string, metadata *dispatchv1.ResponseMeta, results map[string]*dispatchv1.ResourceCheckResult) error {
 		bulkResponseMutex.Lock()
 		defer bulkResponseMutex.Unlock()
 
@@ -548,7 +548,7 @@ func (es *experimentalServer) BulkCheckPermission(ctx context.Context, req *v1.B
 				}
 
 				// Call bulk check to compute the check result(s) for the resource ID(s).
-				rcr, metadata, err := computed.ComputeBulkCheck(ctx, es.dispatch, group.params, resourceIDs)
+				rcr, metadata, err := computed.ComputeBulkCheck(ctx, es.dispatch, *group.params, resourceIDs)
 				if err != nil {
 					return appendResultsForError(group.params, resourceIDs, err)
 				}
@@ -576,7 +576,7 @@ func pairItemFromCheckResult(checkResult *dispatchv1.ResourceCheckResult) *v1.Bu
 	}
 }
 
-func requestItemFromResourceAndParameters(params computed.CheckParameters, resourceID string) (*v1.BulkCheckPermissionRequestItem, error) {
+func requestItemFromResourceAndParameters(params *computed.CheckParameters, resourceID string) (*v1.BulkCheckPermissionRequestItem, error) {
 	item := &v1.BulkCheckPermissionRequestItem{
 		Resource: &v1.ObjectReference{
 			ObjectType: params.ResourceType.Namespace,

diff --git a/internal/services/v1/grouping.go b/internal/services/v1/grouping.go
@@ -7,14 +7,13 @@ import (
 
 	"github.com/authzed/spicedb/internal/graph/computed"
 	"github.com/authzed/spicedb/pkg/datastore"
-	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 )
 
 var MaxBulkCheckDispatchChunkSize = datastore.FilterMaximumIDCount
 
 type groupedCheckParameters struct {
-	params      computed.CheckParameters
+	params      *computed.CheckParameters
 	resourceIDs []string
 }
 
@@ -26,50 +25,51 @@ type groupingParameters struct {
 
 // groupItems takes a slice of BulkCheckPermissionRequestItem and groups them based
 // on using the same permission, subject type, subject id, and caveat.
-func groupItems(ctx context.Context, params groupingParameters, items []*v1.BulkCheckPermissionRequestItem) ([]groupedCheckParameters, error) {
-	groups := mapz.NewMultiMap[string, string]()
-	groupParams := make(map[string]computed.CheckParameters, len(items))
+func groupItems(ctx context.Context, params groupingParameters, items []*v1.BulkCheckPermissionRequestItem) (map[string]*groupedCheckParameters, error) {
+	res := make(map[string]*groupedCheckParameters)
 
 	for _, item := range items {
 		hash, err := computeBulkCheckPermissionItemHashWithoutResourceID(item)
 		if err != nil {
 			return nil, err
 		}
 
-		if _, ok := groupParams[hash]; !ok {
+		if _, ok := res[hash]; !ok {
 			caveatContext, err := GetCaveatContext(ctx, item.Context, params.maxCaveatContextSize)
 			if err != nil {
 				return nil, err
 			}
 
-			groupParams[hash] = computed.CheckParameters{
-				ResourceType: &core.RelationReference{
-					Namespace: item.Resource.ObjectType,
-					Relation:  item.Permission,
-				},
-				Subject: &core.ObjectAndRelation{
-					Namespace: item.Subject.Object.ObjectType,
-					ObjectId:  item.Subject.Object.ObjectId,
-					Relation:  normalizeSubjectRelation(item.Subject),
-				},
-				CaveatContext: caveatContext,
-				AtRevision:    params.atRevision,
-				MaximumDepth:  params.maximumAPIDepth,
-				DebugOption:   computed.NoDebugging,
+			res[hash] = &groupedCheckParameters{
+				params:      checkParametersFromBulkCheckPermissionRequestItem(item, params, caveatContext),
+				resourceIDs: []string{item.Resource.ObjectId},
 			}
+		} else {
+			res[hash].resourceIDs = append(res[hash].resourceIDs, item.Resource.ObjectId)
 		}
-
-		groups.Add(hash, item.Resource.ObjectId)
 	}
 
-	grouped := make([]groupedCheckParameters, 0, groups.Len())
-	for _, key := range groups.Keys() {
-		resourceIDs, _ := groups.Get(key)
-		grouped = append(grouped, groupedCheckParameters{
-			params:      groupParams[key],
-			resourceIDs: resourceIDs,
-		})
-	}
+	return res, nil
+}
 
-	return grouped, nil
+func checkParametersFromBulkCheckPermissionRequestItem(
+	bc *v1.BulkCheckPermissionRequestItem,
+	params groupingParameters,
+	caveatContext map[string]any,
+) *computed.CheckParameters {
+	return &computed.CheckParameters{
+		ResourceType: &core.RelationReference{
+			Namespace: bc.Resource.ObjectType,
+			Relation:  bc.Permission,
+		},
+		Subject: &core.ObjectAndRelation{
+			Namespace: bc.Subject.Object.ObjectType,
+			ObjectId:  bc.Subject.Object.ObjectId,
+			Relation:  normalizeSubjectRelation(bc.Subject),
+		},
+		CaveatContext: caveatContext,
+		AtRevision:    params.atRevision,
+		MaximumDepth:  params.maximumAPIDepth,
+		DebugOption:   computed.NoDebugging,
+	}
 }
diff --git a/internal/services/v1/grouping_test.go b/internal/services/v1/grouping_test.go
@@ -9,6 +9,7 @@ import (
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/maps"
 
 	"github.com/authzed/spicedb/internal/graph/computed"
 	"github.com/authzed/spicedb/pkg/datastore"
@@ -198,10 +199,11 @@ func TestGroupItems(t *testing.T) {
 				maximumAPIDepth:      1,
 			}
 
-			ccp, err := groupItems(context.Background(), cp, items)
+			ccpByHash, err := groupItems(context.Background(), cp, items)
 			if tt.err != "" {
 				require.ErrorContains(t, err, tt.err)
 			} else {
+				ccp := maps.Values(ccpByHash)
 				require.NoError(t, err)
 				require.Equal(t, len(tt.groupings), len(ccp))