Merge pull request #1397 from josephschorr/lr-context-and-perf

Further fixes to context handling in LookupResources and performance improvements
authzed · Jun 19, 2023 · 666d083 · 666d083
2 parents 8321363 + 2937947
commit 666d083
Show file tree

Hide file tree

Showing 14 changed files with 1,118 additions and 453 deletions.
diff --git a/internal/dispatch/caching/caching.go b/internal/dispatch/caching/caching.go
@@ -16,7 +16,6 @@ import (
 	"github.com/authzed/spicedb/internal/dispatch/keys"
 	"github.com/authzed/spicedb/pkg/cache"
 	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
-	"github.com/authzed/spicedb/pkg/spiceerrors"
 )
 
 const (
@@ -224,10 +223,6 @@ func (cd *Dispatcher) DispatchExpand(ctx context.Context, req *v1.DispatchExpand
 func (cd *Dispatcher) DispatchReachableResources(req *v1.DispatchReachableResourcesRequest, stream dispatch.ReachableResourcesStream) error {
 	cd.reachableResourcesTotalCounter.Inc()
 
-	if req.OptionalLimit == 0 {
-		return spiceerrors.MustBugf("a limit must be specified on reachable resources to use with the caching dispatcher")
-	}
-
 	requestKey, err := cd.keyHandler.ReachableResourcesCacheKey(stream.Context(), req)
 	if err != nil {
 		return err
@@ -296,10 +291,6 @@ func sliceSize(xs []byte) int64 {
 func (cd *Dispatcher) DispatchLookupResources(req *v1.DispatchLookupResourcesRequest, stream dispatch.LookupResourcesStream) error {
 	cd.lookupResourcesTotalCounter.Inc()
 
-	if req.OptionalLimit == 0 {
-		return spiceerrors.MustBugf("a limit must be specified on lookup resources to use with the caching dispatcher")
-	}
-
 	requestKey, err := cd.keyHandler.LookupResourcesCacheKey(stream.Context(), req)
 	if err != nil {
 		return err

diff --git a/internal/dispatch/graph/lookupresources_test.go b/internal/dispatch/graph/lookupresources_test.go
@@ -398,7 +398,7 @@ func genResourceIds(resourceName string, number int) []string {
 	return resourceIDs
 }
 
-func TestLookupResourcesOverSchema(t *testing.T) {
+func TestLookupResourcesOverSchemaWithCursors(t *testing.T) {
 	testCases := []struct {
 		name                string
 		schema              string
@@ -545,44 +545,80 @@ func TestLookupResourcesOverSchema(t *testing.T) {
 			ONR("user", "tom", "..."),
 			genResourceIds("document", 150),
 		},
+		{
+			"big",
+			`definition user {}
+		
+		 	 definition document {
+				relation editor: user
+				relation viewer: user
+				permission view = viewer + editor
+  			 }`,
+			joinTuples(
+				genTuples("document", "viewer", "user", "tom", 15100),
+				genTuples("document", "editor", "user", "tom", 15100),
+			),
+			RR("document", "view"),
+			ONR("user", "tom", "..."),
+			genResourceIds("document", 15100),
+		},
 	}
 
 	for _, tc := range testCases {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
-			require := require.New(t)
-
-			dispatcher := NewLocalOnlyDispatcher(10)
-
-			ds, err := memdb.NewMemdbDatastore(0, 0, memdb.DisableGC)
-			require.NoError(err)
-
-			ds, revision := testfixtures.DatastoreFromSchemaAndTestRelationships(ds, tc.schema, tc.relationships, require)
-
-			ctx := datastoremw.ContextWithHandle(context.Background())
-			require.NoError(datastoremw.SetInContext(ctx, ds))
-
-			stream := dispatch.NewCollectingDispatchStream[*v1.DispatchLookupResourcesResponse](ctx)
-			err = dispatcher.DispatchLookupResources(&v1.DispatchLookupResourcesRequest{
-				ObjectRelation: tc.permission,
-				Subject:        tc.subject,
-				Metadata: &v1.ResolverMeta{
-					AtRevision:     revision.String(),
-					DepthRemaining: 50,
-				},
-			}, stream)
-			require.NoError(err)
-
-			foundResourceIDs := util.NewSet[string]()
-			for _, result := range stream.Results() {
-				foundResourceIDs.Add(result.ResolvedResource.ResourceId)
+			for _, pageSize := range []int{0, 104, 1023} {
+				pageSize := pageSize
+				t.Run(fmt.Sprintf("ps-%d_", pageSize), func(t *testing.T) {
+					require := require.New(t)
+
+					dispatcher := NewLocalOnlyDispatcher(10)
+
+					ds, err := memdb.NewMemdbDatastore(0, 0, memdb.DisableGC)
+					require.NoError(err)
+
+					ds, revision := testfixtures.DatastoreFromSchemaAndTestRelationships(ds, tc.schema, tc.relationships, require)
+
+					ctx := datastoremw.ContextWithHandle(context.Background())
+					require.NoError(datastoremw.SetInContext(ctx, ds))
+
+					var currentCursor *v1.Cursor
+					foundResourceIDs := util.NewSet[string]()
+					for {
+						stream := dispatch.NewCollectingDispatchStream[*v1.DispatchLookupResourcesResponse](ctx)
+						err = dispatcher.DispatchLookupResources(&v1.DispatchLookupResourcesRequest{
+							ObjectRelation: tc.permission,
+							Subject:        tc.subject,
+							Metadata: &v1.ResolverMeta{
+								AtRevision:     revision.String(),
+								DepthRemaining: 50,
+							},
+							OptionalLimit:  uint32(pageSize),
+							OptionalCursor: currentCursor,
+						}, stream)
+						require.NoError(err)
+
+						if pageSize > 0 {
+							require.LessOrEqual(len(stream.Results()), pageSize)
+						}
+
+						for _, result := range stream.Results() {
+							foundResourceIDs.Add(result.ResolvedResource.ResourceId)
+							currentCursor = result.AfterResponseCursor
+						}
+
+						if pageSize == 0 || len(stream.Results()) < pageSize {
+							break
+						}
+					}
+
+					foundResourceIDsSlice := foundResourceIDs.AsSlice()
+					sort.Strings(foundResourceIDsSlice)
+					sort.Strings(tc.expectedResourceIDs)
+
+					require.Equal(tc.expectedResourceIDs, foundResourceIDsSlice)
+				})
 			}
-
-			foundResourceIDsSlice := foundResourceIDs.AsSlice()
-			sort.Strings(foundResourceIDsSlice)
-			sort.Strings(tc.expectedResourceIDs)
-
-			require.Equal(tc.expectedResourceIDs, foundResourceIDsSlice)
 		})
 	}
 }