perf(datasets): batch wildcard exclusions in subject set subtraction

CHANGELOG.md

@@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - The `DeleteRelationships` API now supports resumable, cursored batch deletion. When `optional_limit` and `optional_allow_partial_deletions` are set, the response returns an `after_result_cursor`; passing it back as `optional_cursor` on the next call resumes the deletion after the last relationship removed, which on CockroachDB avoids rescanning already-deleted relationships. Datastores that cannot resume a cursored deletion return an error when a cursor is supplied; without a cursor they transparently fall back to the existing non-cursored partial deletion.
 
 ### Fixed
+- LookupSubjects: avoid quadratic exclusion-list copying when subtracting many concrete subjects from a wildcard.
 - Remote dispatch: canceled LookupResources and LookupSubjects streams now return an error instead of reporting success with incomplete results, preventing partial results from being cached and replayed. (https://github.com/authzed/spicedb/pull/3380)
 - Postgres: `DeleteRelationships` with a limit no longer takes time proportional to how many relationships share a subject. The rows to soft-delete were selected by their seven key columns and then re-found by the `UPDATE`, which on a large table resolves through the by-subject index and walks every entry for that subject before filtering; a batch of 1000 could take seconds for a widely-shared subject, or exceed a deadline entirely for a wildcard one. The `UPDATE` now keys on the physical row address (`tableoid`, `ctid`), as the garbage collector already did, so each row is fetched directly. (https://github.com/authzed/spicedb/pull/3309)
 - MemDB/Datastore: a revision reported with no validity is no longer kept by the optimized revision cache, and we fixed an issue where MemDB's reported validity caused a stale value to be cached. (https://github.com/authzed/spicedb/pull/3374)

internal/datasets/basesubjectset.go

@@ -2,7 +2,6 @@ package datasets
 
 import (
 	"maps"
-	"slices"
 
 	"github.com/authzed/spicedb/internal/caveats"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
@@ -144,9 +143,54 @@ func (bss BaseSubjectSet[T]) Subtract(toRemove T) {
 // SubtractAll subtracts the other set of subjects from this set of subtracts, modifying this
 // set *in place*.
 func (bss BaseSubjectSet[T]) SubtractAll(other BaseSubjectSet[T]) {
-	for _, otherSubject := range other.AsSlice() {
-		bss.Subtract(otherSubject)
+	// Batch concrete removals from a wildcard so we only copy its exclusion list
+	// once. The wildcard's caveat stays the same; removals affect its exclusions.
+	//
+	// Fall back to Subtract if there is no wildcard to update, or if other has a
+	// wildcard. Subtracting a wildcard can remove this set's wildcard and turn
+	// the other wildcard's exclusions into concrete results. A single concrete
+	// removal doesn't need batching.
+	wildcard, hasWildcard := bss.wildcard.get()
+	if !hasWildcard || other.wildcard.getOrNil() != nil || len(other.concrete) < 2 {
+		for _, subject := range other.AsSlice() {
+			bss.Subtract(subject)
+		}
+		return
+	}
+
+	// The exclusion slice may be shared with a cloned set, so allocate a new one.
+	// Track existing exclusions that also appear in other so the second pass
+	// doesn't add them again.
+	existingExclusions := wildcard.GetExcludedSubjects()
+	exclusions := make([]T, 0, len(existingExclusions)+len(other.concrete))
+	matched := make(map[string]struct{}, min(len(existingExclusions), len(other.concrete)))
+	for _, exclusion := range existingExclusions {
+		if removing, ok := other.concrete[exclusion.GetSubjectId()]; ok {
+			// Either caveat can exclude the subject, so combine them with OR. If
+			// either has no caveat, the result is an unconditional exclusion:
+			//   {* - {tom[c1]}} - {tom[c2], sarah} => {* - {tom[c1 || c2], sarah}}
+			// Keep both source subjects so the constructor can combine their bookkeeping.
+			matched[exclusion.GetSubjectId()] = struct{}{}
+			exclusion = bss.constructor(exclusion.GetSubjectId(), shortcircuitedOr(exclusion.GetCaveatExpression(), removing.GetCaveatExpression()), nil, exclusion, removing)
+		}
+		exclusions = append(exclusions, exclusion)
 	}
+
+	for subjectID, removing := range other.concrete {
+		if _, ok := matched[subjectID]; !ok {
+			// Keep the removed subject's caveat and bookkeeping for a new exclusion.
+			exclusions = append(exclusions, removing)
+		}
+		if existing, ok := bss.concrete[subjectID]; ok {
+			// The set can also contain this subject as a concrete entry. Remove it
+			// too, or keep it only when its caveat holds and the removal's does not.
+			bss.setConcrete(subjectID, subtractConcreteFromConcrete(existing, removing, bss.constructor))
+		}
+	}
+
+	// Only the exclusions changed; keep the wildcard's own caveat and bookkeeping.
+	updated := bss.constructor(tuple.PublicWildcard, wildcard.GetCaveatExpression(), exclusions, wildcard)
+	bss.wildcard.setOrNil(&updated)
 }
 
 // MustIntersectionDifference performs an intersection between this set and the other set, modifying
@@ -257,8 +301,19 @@ func (bss BaseSubjectSet[T]) IsEmpty() bool {
 
 // AsSlice returns the contents of the subject set as a slice of found subjects.
 func (bss BaseSubjectSet[T]) AsSlice() []T {
-	values := slices.Collect(maps.Values(bss.concrete))
-	if wildcard, ok := bss.wildcard.get(); ok {
+	wildcard, hasWildcard := bss.wildcard.get()
+	count := len(bss.concrete)
+	if hasWildcard {
+		count++
+	}
+	if count == 0 {
+		return nil
+	}
+	values := make([]T, 0, count)
+	for _, concrete := range bss.concrete {
+		values = append(values, concrete)
+	}
+	if hasWildcard {
 		values = append(values, wildcard)
 	}
 	return values

internal/datasets/subjectset_subtract_test.go

@@ -0,0 +1,201 @@
+package datasets
+
+import (
+	"fmt"
+	"testing"
+
+	"google.golang.org/protobuf/proto"
+
+	"github.com/authzed/spicedb/internal/testutil"
+	core "github.com/authzed/spicedb/pkg/proto/core/v1"
+	v1 "github.com/authzed/spicedb/pkg/proto/dispatch/v1"
+)
+
+func TestBatchedSubtract(t *testing.T) {
+	tcs := []struct {
+		name             string
+		startingSubjects []*v1.FoundSubject
+		toSubtract       []*v1.FoundSubject
+		expected         []*v1.FoundSubject
+	}{
+		{
+			"subtract two subjects from a bare wildcard",
+			[]*v1.FoundSubject{wc()},
+			[]*v1.FoundSubject{sub("alice"), sub("bob")},
+			[]*v1.FoundSubject{wc("alice", "bob")},
+		},
+		{
+			"retain an existing exclusion when adding new exclusions",
+			[]*v1.FoundSubject{wc("carol")},
+			[]*v1.FoundSubject{sub("alice"), sub("bob")},
+			[]*v1.FoundSubject{wc("alice", "bob", "carol")},
+		},
+		{
+			"merge overlapping exclusions without duplicates",
+			[]*v1.FoundSubject{wc("alice", "carol")},
+			[]*v1.FoundSubject{sub("alice"), sub("bob")},
+			[]*v1.FoundSubject{wc("alice", "bob", "carol")},
+		},
+		{
+			"all removed subjects are already excluded",
+			[]*v1.FoundSubject{wc("alice", "bob")},
+			[]*v1.FoundSubject{sub("alice"), sub("bob")},
+			[]*v1.FoundSubject{wc("alice", "bob")},
+		},
+		{
+			"preserve the wildcard caveat and new exclusion caveats",
+			[]*v1.FoundSubject{cwc(caveatexpr("public"))},
+			[]*v1.FoundSubject{csub("alice", caveatexpr("banned")), csub("bob", caveatexpr("suspended"))},
+			[]*v1.FoundSubject{
+				cwc(caveatexpr("public"),
+					csub("alice", caveatexpr("banned")),
+					csub("bob", caveatexpr("suspended")),
+				),
+			},
+		},
+		{
+			"combine overlapping exclusion caveats with OR",
+			[]*v1.FoundSubject{cwc(nil, csub("alice", caveatexpr("banned")), csub("carol", caveatexpr("suspended")))},
+			[]*v1.FoundSubject{csub("alice", caveatexpr("suspended")), sub("bob")},
+			[]*v1.FoundSubject{
+				cwc(nil,
+					csub("alice", caveatOr(caveatexpr("banned"), caveatexpr("suspended"))),
+					sub("bob"),
+					csub("carol", caveatexpr("suspended")),
+				),
+			},
+		},
+		{
+			"unconditional removal replaces a conditional exclusion",
+			[]*v1.FoundSubject{cwc(nil, csub("alice", caveatexpr("banned")))},
+			[]*v1.FoundSubject{sub("alice"), sub("bob")},
+			[]*v1.FoundSubject{wc("alice", "bob")},
+		},
+		{
+			"conditional removal preserves an unconditional exclusion",
+			[]*v1.FoundSubject{wc("alice")},
+			[]*v1.FoundSubject{csub("alice", caveatexpr("banned")), sub("bob")},
+			[]*v1.FoundSubject{wc("alice", "bob")},
+		},
+		{
+			"remove a matching concrete subject and retain an unrelated subject",
+			[]*v1.FoundSubject{wc(), sub("alice"), sub("carol")},
+			[]*v1.FoundSubject{sub("alice"), sub("bob")},
+			[]*v1.FoundSubject{wc("alice", "bob"), sub("carol")},
+		},
+		{
+			"conditional removal updates both the concrete subject and wildcard",
+			[]*v1.FoundSubject{cwc(caveatexpr("public")), csub("alice", caveatexpr("member"))},
+			[]*v1.FoundSubject{csub("alice", caveatexpr("banned")), sub("bob")},
+			[]*v1.FoundSubject{
+				cwc(caveatexpr("public"), csub("alice", caveatexpr("banned")), sub("bob")),
+				csub("alice", caveatAnd(caveatexpr("member"), caveatInvert(caveatexpr("banned")))),
+			},
+		},
+		{
+			"subtract a wildcard alongside concrete subjects",
+			[]*v1.FoundSubject{wc("carol")},
+			[]*v1.FoundSubject{wc("alice"), sub("bob"), sub("dave")},
+			[]*v1.FoundSubject{sub("alice")},
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			set := NewSubjectSet()
+			for _, subject := range tc.startingSubjects {
+				set.MustAdd(subject)
+			}
+			toRemove := NewSubjectSet()
+			for _, subject := range tc.toSubtract {
+				toRemove.MustAdd(subject)
+			}
+
+			set.SubtractAll(toRemove)
+			testutil.RequireEquivalentSets(t, tc.expected, set.AsSlice())
+		})
+	}
+}
+
+func TestBatchedSubtractMatchesSequential(t *testing.T) {
+	expressions := []*core.CaveatExpression{nil, caveatexpr("a"), caveatexpr("b")}
+	for wi, wexpr := range expressions {
+		for ei, eexpr := range expressions {
+			for ri, rexpr := range expressions {
+				for _, count := range []int{0, 1, 2, 10} {
+					t.Run(fmt.Sprintf("%d/%d/%d/%d", wi, ei, ri, count), func(t *testing.T) {
+						initial := NewSubjectSet()
+						initial.MustAdd(cwc(wexpr, csub("u0", eexpr), csub("retained", eexpr)))
+						initial.MustAdd(csub("u1", caveatexpr("c")))
+						initial.MustAdd(sub("unrelated"))
+						removing := NewSubjectSet()
+						for i := range count {
+							removing.MustAdd(csub(fmt.Sprintf("u%d", i), rexpr))
+						}
+						original := cloneSubjects(initial.AsSlice())
+						reference := initial.Clone()
+						actual := initial.Clone()
+						beforeRemoving := cloneSubjects(removing.AsSlice())
+						for _, subject := range removing.AsSlice() {
+							reference.Subtract(subject)
+						}
+						actual.SubtractAll(removing)
+						testutil.RequireEquivalentSets(t, reference.AsSlice(), actual.AsSlice())
+						testutil.RequireEquivalentSets(t, original, initial.AsSlice())
+						testutil.RequireEquivalentSets(t, beforeRemoving, removing.AsSlice())
+					})
+				}
+			}
+		}
+	}
+	for _, wildcard := range []*v1.FoundSubject{wc(), cwc(caveatexpr("a"), csub("u0", caveatexpr("b")))} {
+		initial := NewSubjectSet()
+		initial.MustAdd(wildcard)
+		initial.MustAdd(csub("u1", caveatexpr("c")))
+		reference := initial.Clone()
+		actual := initial.Clone()
+		for _, subject := range reference.AsSlice() {
+			reference.Subtract(subject)
+		}
+		actual.SubtractAll(actual)
+		testutil.RequireEquivalentSets(t, reference.AsSlice(), actual.AsSlice())
+	}
+}
+
+func cloneSubjects(subjects []*v1.FoundSubject) []*v1.FoundSubject {
+	cloned := make([]*v1.FoundSubject, len(subjects))
+	for i, subject := range subjects {
+		cloned[i] = proto.Clone(subject).(*v1.FoundSubject)
+	}
+	return cloned
+}
+
+func BenchmarkSubjectSetSubtractAll(b *testing.B) {
+	for _, tc := range []struct {
+		name                                string
+		existingCount, removalCount, offset int
+	}{
+		{"wildcard/5000", 0, 5000, 0},
+		{"existing5000/disjoint10", 5000, 10, 5000},
+		{"existing5000/overlap10", 5000, 10, 0},
+		{"existing5000/overlap5000", 5000, 5000, 0},
+	} {
+		b.Run(tc.name, func(b *testing.B) {
+			excluded := make([]string, tc.existingCount)
+			for i := range excluded {
+				excluded[i] = fmt.Sprintf("u%d", i)
+			}
+			wildcard := wc(excluded...)
+			removing := NewSubjectSet()
+			for i := range tc.removalCount {
+				removing.MustAdd(sub(fmt.Sprintf("u%d", tc.offset+i)))
+			}
+			b.ReportAllocs()
+			for b.Loop() {
+				set := NewSubjectSet()
+				set.MustAdd(wildcard)
+				set.SubtractAll(removing)
+			}
+		})
+	}
+}

internal/developmentmembership/trackingsubjectset_test.go

@@ -383,3 +383,29 @@ func TestTrackingSubjectSetResourceTrackingWithWildcard(t *testing.T) {
 	require.True(t, ok)
 	require.Len(t, found.ParentResources(), 1)
 }
+
+func TestBatchedExclusionResources(t *testing.T) {
+	wildcardResource := tuple.ONR("document", "public", "view")
+	existingResource := tuple.ONR("document", "old", "banned")
+	removingResource := tuple.ONR("document", "new", "banned")
+	wildcard := NewFoundSubject(DS("user", "*", "..."), wildcardResource)
+	wildcard.excludedSubjects = []FoundSubject{NewFoundSubject(DS("user", "u0", "..."), existingResource)}
+	initial := MustNewTrackingSubjectSetWith(wildcard)
+	removing := MustNewTrackingSubjectSetWith(
+		NewFoundSubject(DS("user", "u0", "..."), removingResource),
+		NewFoundSubject(DS("user", "u1", "..."), removingResource),
+	)
+	initial.RemoveFrom(removing)
+	found, ok := initial.getSet(wildcard).Get("*")
+	require.True(t, ok)
+	require.ElementsMatch(t, []tuple.ObjectAndRelation{wildcardResource}, found.resources.AsSlice())
+	require.Len(t, found.excludedSubjects, 2)
+	for _, exclusion := range found.excludedSubjects {
+		expected := []tuple.ObjectAndRelation{removingResource}
+		if exclusion.GetSubjectId() == "u0" {
+			expected = append(expected, existingResource)
+		}
+		require.ElementsMatch(t, expected, exclusion.resources.AsSlice())
+	}
+	require.ElementsMatch(t, []tuple.ObjectAndRelation{existingResource}, wildcard.excludedSubjects[0].resources.AsSlice())
+}