feat: add recursive direction strategies, and fix IS BFS (#2891)

Author: barakmich

pkg/query/benchmarks/check_deep_arrow_benchmark_test.go

@@ -0,0 +1,117 @@
+package benchmarks
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/authzed/spicedb/internal/datastore/common"
+	"github.com/authzed/spicedb/internal/datastore/memdb"
+	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/query"
+	"github.com/authzed/spicedb/pkg/schema/v2"
+	"github.com/authzed/spicedb/pkg/schemadsl/compiler"
+	"github.com/authzed/spicedb/pkg/schemadsl/input"
+	"github.com/authzed/spicedb/pkg/tuple"
+)
+
+// BenchmarkCheckDeepArrow benchmarks permission checking through a deep recursive chain.
+// This recreates the testharness scenario with:
+// - A 30+ level deep parent chain: document:target -> document:1 -> ... -> document:29
+// - document:29#view@user:slow
+// - Checking if user:slow has viewer permission on document:target
+//
+// The permission viewer = view + parent->viewer creates a recursive traversal through
+// all 30+ levels to find the view relationship at the end of the chain.
+func BenchmarkCheckDeepArrow(b *testing.B) {
+	// Create an in-memory datastore
+	rawDS, err := memdb.NewMemdbDatastore(0, 0, memdb.DisableGC)
+	require.NoError(b, err)
+
+	ctx := context.Background()
+
+	schemaText := `
+		definition user {}
+
+		definition document {
+			relation parent: document
+			relation view: user
+			permission viewer = view + parent->viewer
+		}
+	`
+
+	// Compile the schema
+	compiled, err := compiler.Compile(compiler.InputSchema{
+		Source:       input.Source("benchmark"),
+		SchemaString: schemaText,
+	}, compiler.AllowUnprefixedObjectType())
+	require.NoError(b, err)
+
+	// Write the schema
+	_, err = rawDS.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
+		return rwt.LegacyWriteNamespaces(ctx, compiled.ObjectDefinitions...)
+	})
+	require.NoError(b, err)
+
+	// Build relationships for the deep arrow scenario
+	// Create a chain: document:target -> document:1 -> document:2 -> ... -> document:30 -> document:31
+	// Plus: document:29#view@user:slow
+	relationships := make([]tuple.Relationship, 0, 33)
+
+	// document:target#parent@document:1
+	relationships = append(relationships, tuple.MustParse("document:target#parent@document:1"))
+
+	// Chain: document:1 through document:30
+	for i := 1; i <= 30; i++ {
+		rel := fmt.Sprintf("document:%d#parent@document:%d", i, i+1)
+		relationships = append(relationships, tuple.MustParse(rel))
+	}
+
+	// The view relationship at the end of the chain
+	relationships = append(relationships, tuple.MustParse("document:29#view@user:slow"))
+
+	// Write all relationships to the datastore
+	revision, err := common.WriteRelationships(ctx, rawDS, tuple.UpdateOperationCreate, relationships...)
+	require.NoError(b, err)
+
+	// Build schema for querying
+	dsSchema, err := schema.BuildSchemaFromDefinitions(compiled.ObjectDefinitions, nil)
+	require.NoError(b, err)
+
+	// Create the iterator tree for the viewer permission using BuildIteratorFromSchema
+	viewerIterator, err := query.BuildIteratorFromSchema(dsSchema, "document", "viewer")
+	require.NoError(b, err)
+
+	// Create query context
+	queryCtx := query.NewLocalContext(ctx,
+		query.WithReader(rawDS.SnapshotReader(revision)),
+		query.WithMaxRecursionDepth(50),
+	)
+
+	// The resource we're checking: document:target
+	resources := query.NewObjects("document", "target")
+
+	// The subject we're checking: user:slow
+	subject := query.NewObject("user", "slow").WithEllipses()
+
+	// Reset the timer - everything before this is setup
+	b.ResetTimer()
+
+	// Run the benchmark
+	for b.Loop() {
+		// Check if user:slow can view document:target
+		// This will traverse the entire 30+ level chain
+		seq, err := queryCtx.Check(viewerIterator, resources, subject)
+		require.NoError(b, err)
+
+		// Collect all results (should find user:slow at the end of the chain)
+		paths, err := query.CollectAll(seq)
+		require.NoError(b, err)
+
+		// Verify we found the expected result
+		require.Len(b, paths, 1)
+		require.Equal(b, "slow", paths[0].Subject.ObjectID)
+	}
+}

pkg/query/context.go

@@ -3,21 +3,25 @@ package query
 import (
 	"context"
 	"fmt"
+	"io"
 	"maps"
 	"strings"
 	"sync"
 	"time"
 
 	"github.com/authzed/spicedb/internal/caveats"
 	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/datastore/options"
 	"github.com/authzed/spicedb/pkg/spiceerrors"
+	"github.com/authzed/spicedb/pkg/tuple"
 )
 
 // TraceLogger is used for debugging iterator execution
 type TraceLogger struct {
 	traces []string
 	depth  int
 	stack  []Iterator // Stack of iterator pointers for proper indentation context
+	writer io.Writer  // Optional writer to output traces in real-time
 }
 
 // NewTraceLogger creates a new trace logger
@@ -26,6 +30,28 @@ func NewTraceLogger() *TraceLogger {
 		traces: make([]string, 0),
 		depth:  0,
 		stack:  make([]Iterator, 0),
+		writer: nil,
+	}
+}
+
+// NewTraceLoggerWithWriter creates a new trace logger with an optional writer
+// for real-time trace output
+func NewTraceLoggerWithWriter(w io.Writer) *TraceLogger {
+	return &TraceLogger{
+		traces: make([]string, 0),
+		depth:  0,
+		stack:  make([]Iterator, 0),
+		writer: w,
+	}
+}
+
+// appendTrace appends a trace line to the traces slice and optionally writes it
+// to the writer if one is configured
+func (t *TraceLogger) appendTrace(line string) {
+	t.traces = append(t.traces, line)
+	if t.writer != nil {
+		// Write the line with a newline
+		fmt.Fprintln(t.writer, line)
 	}
 }
 
@@ -53,14 +79,11 @@ func (t *TraceLogger) EnterIterator(it Iterator, traceString string) {
 	indent := strings.Repeat("  ", t.depth)
 	idPrefix := iteratorIDPrefix(it)
 
-	t.traces = append(
-		t.traces,
-		fmt.Sprintf("%s-> %s: %s",
-			indent,
-			idPrefix,
-			traceString,
-		),
-	)
+	t.appendTrace(fmt.Sprintf("%s-> %s: %s",
+		indent,
+		idPrefix,
+		traceString,
+	))
 	t.depth++
 	t.stack = append(t.stack, it) // Push iterator pointer onto stack
 }
@@ -129,7 +152,7 @@ func (t *TraceLogger) ExitIterator(it Iterator, paths []Path) {
 			p.Resource.ObjectType, p.Resource.ObjectID, p.Relation,
 			p.Subject.ObjectType, p.Subject.ObjectID, caveatInfo)
 	}
-	t.traces = append(t.traces, fmt.Sprintf("%s<- %s: returned %d paths: [%s]",
+	t.appendTrace(fmt.Sprintf("%s<- %s: returned %d paths: [%s]",
 		indent, idPrefix, len(paths), strings.Join(pathStrs, ", ")))
 }
 
@@ -155,7 +178,7 @@ func (t *TraceLogger) LogStep(it Iterator, step string, data ...any) {
 	indent := strings.Repeat("  ", indentLevel)
 	idPrefix := iteratorIDPrefix(it)
 	message := fmt.Sprintf(step, data...)
-	t.traces = append(t.traces, fmt.Sprintf("%s   %s: %s", indent, idPrefix, message))
+	t.appendTrace(fmt.Sprintf("%s   %s: %s", indent, idPrefix, message))
 }
 
 // DumpTrace returns all traces as a string
@@ -258,6 +281,17 @@ type Context struct {
 	TraceLogger       *TraceLogger      // For debugging iterator execution
 	Analyze           *AnalyzeCollector // Thread-safe collector for query analysis stats
 	MaxRecursionDepth int               // Maximum depth for recursive iterators (0 = use default of 10)
+
+	// Pagination options for IterSubjects and IterResources
+	PaginationCursors map[string]*tuple.Relationship // Cursors for pagination, keyed by iterator ID
+	PaginationLimit   *uint64                        // Limit for pagination (max number of results to return)
+	PaginationSort    options.SortOrder              // Sort order for pagination
+
+	// recursiveFrontierCollectors holds frontier collections for BFS IterSubjects.
+	// Key: RecursiveIterator.ID()
+	// Value: collected Objects for the next frontier
+	// A non-nil entry for an ID enables collection mode for that RecursiveIterator.
+	recursiveFrontierCollectors map[string][]Object
 }
 
 // NewLocalContext creates a new query execution context with a LocalExecutor.
@@ -306,6 +340,32 @@ func WithMaxRecursionDepth(depth int) ContextOption {
 	return func(ctx *Context) { ctx.MaxRecursionDepth = depth }
 }
 
+// WithPaginationLimit sets the pagination limit for the context.
+func WithPaginationLimit(limit uint64) ContextOption {
+	return func(ctx *Context) { ctx.PaginationLimit = &limit }
+}
+
+// WithPaginationSort sets the pagination sort order for the context.
+func WithPaginationSort(sort options.SortOrder) ContextOption {
+	return func(ctx *Context) { ctx.PaginationSort = sort }
+}
+
+// GetPaginationCursor retrieves the cursor for a specific iterator ID.
+func (ctx *Context) GetPaginationCursor(iteratorID string) *tuple.Relationship {
+	if ctx.PaginationCursors == nil {
+		return nil
+	}
+	return ctx.PaginationCursors[iteratorID]
+}
+
+// SetPaginationCursor sets the cursor for a specific iterator ID.
+func (ctx *Context) SetPaginationCursor(iteratorID string, cursor *tuple.Relationship) {
+	if ctx.PaginationCursors == nil {
+		ctx.PaginationCursors = make(map[string]*tuple.Relationship)
+	}
+	ctx.PaginationCursors[iteratorID] = cursor
+}
+
 func (ctx *Context) TraceStep(it Iterator, step string, data ...any) {
 	if ctx.TraceLogger != nil {
 		ctx.TraceLogger.LogStep(it, step, data...)
@@ -483,3 +543,42 @@ type Executor interface {
 	// specified ObjectType. If filterResourceType.Type is empty, no filtering is applied.
 	IterResources(ctx *Context, it Iterator, subject ObjectAndRelation, filterResourceType ObjectType) (PathSeq, error)
 }
+
+// EnableFrontierCollection enables frontier collection for a RecursiveIterator.
+// Creates a non-nil entry in the map, which signals collection mode.
+func (ctx *Context) EnableFrontierCollection(iteratorID string) {
+	if ctx.recursiveFrontierCollectors == nil {
+		ctx.recursiveFrontierCollectors = make(map[string][]Object)
+	}
+	ctx.recursiveFrontierCollectors[iteratorID] = []Object{}
+}
+
+// CollectFrontierObject appends an object to the frontier collection.
+// Only appends if collection mode is enabled (non-nil entry exists).
+func (ctx *Context) CollectFrontierObject(iteratorID string, obj Object) {
+	if ctx.recursiveFrontierCollectors == nil {
+		return
+	}
+	if collection, exists := ctx.recursiveFrontierCollectors[iteratorID]; exists {
+		ctx.recursiveFrontierCollectors[iteratorID] = append(collection, obj)
+	}
+}
+
+// ExtractFrontierCollection retrieves and removes the collected frontier.
+func (ctx *Context) ExtractFrontierCollection(iteratorID string) []Object {
+	if ctx.recursiveFrontierCollectors == nil {
+		return nil
+	}
+	collection := ctx.recursiveFrontierCollectors[iteratorID]
+	delete(ctx.recursiveFrontierCollectors, iteratorID)
+	return collection
+}
+
+// IsCollectingFrontier checks if collection mode is enabled (non-nil entry exists).
+func (ctx *Context) IsCollectingFrontier(iteratorID string) bool {
+	if ctx.recursiveFrontierCollectors == nil {
+		return false
+	}
+	_, exists := ctx.recursiveFrontierCollectors[iteratorID]
+	return exists
+}