perf(mysql): force per-query-shape indexes for relationship queries

CHANGELOG.md

@@ -28,6 +28,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 - Cache: switch to [otter](https://maypok86.github.io/otter/) as the primary cache implementation (https://github.com/authzed/spicedb/pull/3112)
 - Server handles: `GRPCDialContext` as a handle on the server used deprecated gRPC methods. We modernized it and renamed it to `NewClient` (https://github.com/authzed/spicedb/pull/3147)
+- MySQL: relationship reads now force the index that matches each query shape (and its sort order), rather than relying on the MySQL optimizer to pick one. This avoids cases where the optimizer chose a suboptimal index for `Check`, `LookupResources`, `LookupSubjects`, and `ReadRelationships`/reverse-relationship queries. (https://github.com/authzed/spicedb/pull/3173)
 
 ### Fixed
 - The watching schema cache (`--enable-experimental-watchable-schema-cache`) no longer enters permanent fallback on transient watch errors. A new supervisor restarts the watch cycle with bounded exponential backoff and only treats caller-driven cancellation or unsupported-watch as terminal (https://github.com/authzed/spicedb/pull/3134)

internal/datastore/benchmark/mysql_rels_bench_test.go

@@ -0,0 +1,181 @@
+package benchmark
+
+import (
+	"context"
+	"strconv"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/authzed/spicedb/internal/datastore/mysql"
+	"github.com/authzed/spicedb/internal/testfixtures"
+	testdatastore "github.com/authzed/spicedb/internal/testserver/datastore"
+	"github.com/authzed/spicedb/internal/testserver/datastore/config"
+	dsconfig "github.com/authzed/spicedb/pkg/cmd/datastore"
+	"github.com/authzed/spicedb/pkg/datastore"
+	"github.com/authzed/spicedb/pkg/datastore/options"
+	"github.com/authzed/spicedb/pkg/datastore/queryshape"
+	"github.com/authzed/spicedb/pkg/tuple"
+)
+
+const (
+	mysqlBenchNumDocuments = 5000
+	mysqlBenchUsersPerDoc  = 10
+)
+
+// BenchmarkMySQLRelQueries benchmarks the relationship-query shapes affected by the
+// MySQL index-hint / pagination work. Requires Docker (spins up mysql:8.4).
+//
+// Run with:
+//
+//	go test ./internal/datastore/benchmark/ -bench BenchmarkMySQLRelQueries \
+//	    -benchmem -run '^$' -timeout 30m
+func BenchmarkMySQLRelQueries(b *testing.B) {
+	b.StopTimer()
+	ctx := b.Context()
+
+	engine := testdatastore.RunDatastoreEngine(b, mysql.Engine)
+	ds := engine.NewDatastore(b, config.DatastoreConfigInitFunc(
+		b,
+		dsconfig.WithRevisionQuantization(5*time.Second),
+		dsconfig.WithGCWindow(2*time.Hour),
+		dsconfig.WithGCInterval(1*time.Hour),
+		dsconfig.WithWatchBufferLength(1000),
+		dsconfig.WithWriteAcquisitionTimeout(5*time.Second),
+	))
+	b.Cleanup(func() { _ = ds.Close() })
+
+	ds, _ = testfixtures.StandardDatastoreWithSchema(b, ds)
+
+	// Write a meaningful amount of data. Subject user "0" ends up a viewer of every
+	// document, which exercises the reverse (MatchingResourcesForSubject) shape well.
+	for docNum := range mysqlBenchNumDocuments {
+		_, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
+			updates := make([]tuple.RelationshipUpdate, 0, mysqlBenchUsersPerDoc)
+			for userNum := range mysqlBenchUsersPerDoc {
+				updates = append(updates, tuple.Create(docViewer(strconv.Itoa(docNum), strconv.Itoa(userNum))))
+			}
+			return rwt.WriteRelationships(ctx, updates)
+		})
+		require.NoError(b, err)
+	}
+
+	time.Sleep(1 * time.Second)
+
+	headRevResult, err := ds.HeadRevision(ctx)
+	require.NoError(b, err)
+	headRev := headRevResult.Revision
+
+	drain := func(b *testing.B, iter datastore.RelationshipIterator) int {
+		count := 0
+		for _, err := range iter {
+			require.NoError(b, err)
+			count++
+		}
+		return count
+	}
+
+	b.StartTimer()
+
+	// Forward: full resource + subject (direct check).
+	b.Run("ForwardDirectCheck", func(b *testing.B) {
+		for range b.N {
+			iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+				OptionalResourceType:     testfixtures.DocumentNS.Name,
+				OptionalResourceIds:      []string{"100"},
+				OptionalResourceRelation: "viewer",
+				OptionalSubjectsSelectors: []datastore.SubjectsSelector{{
+					OptionalSubjectType: testfixtures.UserNS.Name,
+					OptionalSubjectIds:  []string{"0"},
+				}},
+			}, options.WithQueryShape(queryshape.CheckPermissionSelectDirectSubjects))
+			require.NoError(b, err)
+			require.Equal(b, 1, drain(b, iter))
+		}
+	})
+
+	// Forward: all subjects for one resource.
+	b.Run("ForwardAllSubjects", func(b *testing.B) {
+		for range b.N {
+			iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+				OptionalResourceType:     testfixtures.DocumentNS.Name,
+				OptionalResourceIds:      []string{"100"},
+				OptionalResourceRelation: "viewer",
+			}, options.WithQueryShape(queryshape.AllSubjectsForResources))
+			require.NoError(b, err)
+			require.Equal(b, mysqlBenchUsersPerDoc, drain(b, iter))
+		}
+	})
+
+	// Reverse: all resources for a specific subject (Finding 1 main target).
+	b.Run("ReverseMatchingResources", func(b *testing.B) {
+		for range b.N {
+			iter, err := ds.SnapshotReader(headRev).ReverseQueryRelationships(ctx, datastore.SubjectsFilter{
+				SubjectType:        testfixtures.UserNS.Name,
+				OptionalSubjectIds: []string{"0"},
+			},
+				options.WithResRelation(&options.ResourceRelation{
+					Namespace: testfixtures.DocumentNS.Name,
+					Relation:  "viewer",
+				}),
+				options.WithSortForReverse(options.BySubject),
+				options.WithQueryShapeForReverse(queryshape.MatchingResourcesForSubject),
+			)
+			require.NoError(b, err)
+			require.Equal(b, mysqlBenchNumDocuments, drain(b, iter))
+		}
+	})
+
+	// Enumerate by type + relation, no object_id (Finding 3 candidate workload).
+	b.Run("EnumerateByTypeAndRelation", func(b *testing.B) {
+		for range b.N {
+			iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+				OptionalResourceType:     testfixtures.DocumentNS.Name,
+				OptionalResourceRelation: "viewer",
+			}, options.WithQueryShape(queryshape.Varying))
+			require.NoError(b, err)
+			require.Equal(b, mysqlBenchNumDocuments*mysqlBenchUsersPerDoc, drain(b, iter))
+		}
+	})
+
+	// Cursor-paginated read (Finding 4 target).
+	b.Run("PaginatedRead", func(b *testing.B) {
+		const pageSize = uint64(500)
+		for range b.N {
+			var after options.Cursor
+			total := 0
+			for {
+				limit := pageSize
+				opts := []options.QueryOptionsOption{
+					options.WithSort(options.ByResource),
+					options.WithLimit(&limit),
+					options.WithQueryShape(queryshape.Varying),
+				}
+				if after != nil {
+					opts = append(opts, options.WithAfter(after))
+				}
+				iter, err := ds.SnapshotReader(headRev).QueryRelationships(ctx, datastore.RelationshipsFilter{
+					OptionalResourceType:     testfixtures.DocumentNS.Name,
+					OptionalResourceRelation: "viewer",
+				}, opts...)
+				require.NoError(b, err)
+
+				pageCount := 0
+				var last tuple.Relationship
+				for rel, err := range iter {
+					require.NoError(b, err)
+					last = rel
+					pageCount++
+				}
+				total += pageCount
+				if uint64(pageCount) < pageSize {
+					break
+				}
+				lastCopy := last
+				after = &lastCopy
+			}
+			require.Equal(b, mysqlBenchNumDocuments*mysqlBenchUsersPerDoc, total)
+		}
+	})
+}

internal/datastore/mysql/indexes.go

@@ -2,54 +2,54 @@
 
 import (
 	"github.com/authzed/spicedb/internal/datastore/common"
+	"github.com/authzed/spicedb/pkg/datastore/options"
 	"github.com/authzed/spicedb/pkg/datastore/queryshape"
 )
 
-// IndexUniqueLivingRelationships is the UNIQUE constraint index on
-// living relationships.
+// IndexUniqueLivingRelationships is the UNIQUE constraint index on living relationships.
+// Because relation_tuple is clustered on the synthetic `id` primary key, this resource-keyed
+// index plays the role that the primary key plays in CRDB/Postgres for forward query shapes.
 var IndexUniqueLivingRelationships = common.IndexDefinition{
 	Name:       `uq_relation_tuple_living`,
 	ColumnsSQL: `UNIQUE (namespace, object_id, relation, userset_namespace, userset_object_id, userset_relation, deleted_transaction)`,
 	Shapes: []queryshape.Shape{
 		queryshape.CheckPermissionSelectDirectSubjects,
+		queryshape.CheckPermissionSelectIndirectSubjects,
+		queryshape.AllSubjectsForResources,
+		queryshape.FindResourceOfType,
 	},
+	PreferredSortOrder: options.ByResource,
 }
 
-// IndexUniqueRelationships is the UNIQUE constraint index on all relationships.
+// IndexUniqueRelationships is the UNIQUE constraint index on all relationships (including
+// historical ones). It enforces uniqueness; it is not used as a read index for any query shape.
 var IndexUniqueRelationships = common.IndexDefinition{
 	Name:       `uq_relation_tuple_namespace`,
 	ColumnsSQL: `UNIQUE (namespace, object_id, relation, userset_namespace, userset_object_id, userset_relation, created_transaction, deleted_transaction)`,
-	Shapes: []queryshape.Shape{
-		queryshape.CheckPermissionSelectDirectSubjects,
-		queryshape.CheckPermissionSelectIndirectSubjects,
-	},
 }
 
-// IndexRelationshipBySubject is the index on the relationship table for
-// looking up relationships by subject.
+// IndexRelationshipBySubject is the index on the relationship table for looking up
+// relationships by subject (the reverse lookup).
 var IndexRelationshipBySubject = common.IndexDefinition{
 	Name:       `ix_relation_tuple_by_subject`,
 	ColumnsSQL: `INDEX ix_relation_tuple_by_subject (userset_object_id, userset_namespace, userset_relation, namespace, relation)`,
 	Shapes: []queryshape.Shape{
-		// TODO: this index isn't great to use for this query shape, so we probably need to fix
-		// the indexes
-		queryshape.CheckPermissionSelectDirectSubjects,
-
-		queryshape.CheckPermissionSelectIndirectSubjects,
+		queryshape.MatchingResourcesForSubject,
 	},
+	PreferredSortOrder: options.BySubject,
 }
 
-// IndexRelationshipBySubjectRelation is the index on the relationship table for
-// looking up relationships by subject and relation.
+// IndexRelationshipBySubjectRelation is the index on the relationship table for looking up
+// relationships by subject type and relation.
 var IndexRelationshipBySubjectRelation = common.IndexDefinition{
 	Name:       `ix_relation_tuple_by_subject_relation`,
 	ColumnsSQL: `INDEX ix_relation_tuple_by_subject_relation (userset_namespace, userset_relation, namespace, relation)`,
 	Shapes: []queryshape.Shape{
-		// TODO: this index isn't great to use for these query shape, so we probably need to fix
-		// the indexes
-		queryshape.CheckPermissionSelectDirectSubjects,
-		queryshape.CheckPermissionSelectIndirectSubjects,
+		queryshape.FindResourceAndSubjectWithRelations,
+		queryshape.FindSubjectOfTypeAndRelation,
+		queryshape.FindResourceRelationForSubjectRelation,
 	},
+	PreferredSortOrder: options.BySubject,
 }
 
 // IndexRelationshipByDeletedTransaction is the index on the relationship table for
@@ -59,16 +59,14 @@
 	ColumnsSQL: `INDEX ix_relation_tuple_by_deleted_transaction (deleted_transaction)`,
 }
 
-// IndexRelationTupleWatch is the index on the relationship table for
-// watching relationships.
+// IndexRelationTupleWatch is the index on the relationship table for watching relationships.
 var IndexRelationTupleWatch = common.IndexDefinition{
 	Name:       `ix_relation_tuple_watch`,
 	ColumnsSQL: `INDEX ix_relation_tuple_watch (created_transaction, deleted_transaction DESC)`,
 }
 
-// IndexRelationTupleExpired is the index on the relationship table for
-// expired relationships. This is used for garbage collection of expired
-// relationships.
+// IndexRelationTupleExpired is the index on the relationship table for expired relationships.
+// This is used for garbage collection of expired relationships.
 var IndexRelationTupleExpired = common.IndexDefinition{
 	Name:       `ix_relation_tuple_expired`,
 	ColumnsSQL: `INDEX ix_relation_tuple_expired (expiration)`,
@@ -83,3 +81,51 @@
 	IndexRelationTupleWatch,
 	IndexRelationTupleExpired,
 }
+
+// NoIndexingHint is the absence of an indexing hint.
+var NoIndexingHint common.IndexingHint = nil
+
+// IndexingHintForQueryShape returns the index hint to force for a given query shape, if any.
+func IndexingHintForQueryShape(qs queryshape.Shape) common.IndexingHint {
+	switch qs {
+	case queryshape.CheckPermissionSelectDirectSubjects,
+		queryshape.CheckPermissionSelectIndirectSubjects,
+		queryshape.AllSubjectsForResources,
+		queryshape.FindResourceOfType:
+		return forcedIndex{IndexUniqueLivingRelationships}
+
+	case queryshape.MatchingResourcesForSubject:
+		return forcedIndex{IndexRelationshipBySubject}
+
+	case queryshape.FindResourceAndSubjectWithRelations,
+		queryshape.FindSubjectOfTypeAndRelation,
+		queryshape.FindResourceRelationForSubjectRelation:
+		return forcedIndex{IndexRelationshipBySubjectRelation}
+
+	default:
+		return NoIndexingHint
+	}
+}
+
+// forcedIndex is an index hint that forces use of a specific index via MySQL's FORCE INDEX clause.
+type forcedIndex struct {
+	index common.IndexDefinition
+}
+
+func (f forcedIndex) FromSQLSuffix() (string, error) {
+	return "FORCE INDEX (" + f.index.Name + ")", nil
+}
+
+func (f forcedIndex) FromTable(existingTableName string) (string, error) {
+	return existingTableName, nil
+}
+
+func (f forcedIndex) SQLPrefix() (string, error) {
+	return "", nil
+}
+
+func (f forcedIndex) SortOrder() options.SortOrder {
+	return f.index.PreferredSortOrder
+}
+
+var _ common.IndexingHint = forcedIndex{}

internal/datastore/mysql/indexes_test.go

@@ -0,0 +1,46 @@
+package mysql
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/authzed/spicedb/pkg/datastore/options"
+	"github.com/authzed/spicedb/pkg/datastore/queryshape"
+)
+
+func TestIndexingHintForQueryShape(t *testing.T) {
+	cases := map[queryshape.Shape]struct {
+		indexName string
+		sort      options.SortOrder
+	}{
+		queryshape.CheckPermissionSelectDirectSubjects:    {"uq_relation_tuple_living", options.ByResource},
+		queryshape.CheckPermissionSelectIndirectSubjects:  {"uq_relation_tuple_living", options.ByResource},
+		queryshape.AllSubjectsForResources:                {"uq_relation_tuple_living", options.ByResource},
+		queryshape.FindResourceOfType:                     {"uq_relation_tuple_living", options.ByResource},
+		queryshape.MatchingResourcesForSubject:            {"ix_relation_tuple_by_subject", options.BySubject},
+		queryshape.FindResourceAndSubjectWithRelations:    {"ix_relation_tuple_by_subject_relation", options.BySubject},
+		queryshape.FindSubjectOfTypeAndRelation:           {"ix_relation_tuple_by_subject_relation", options.BySubject},
+		queryshape.FindResourceRelationForSubjectRelation: {"ix_relation_tuple_by_subject_relation", options.BySubject},
+	}
+
+	for shape, expected := range cases {
+		t.Run(string(shape), func(t *testing.T) {
+			hint := IndexingHintForQueryShape(shape)
+			require.NotNil(t, hint)
+
+			suffix, err := hint.FromSQLSuffix()
+			require.NoError(t, err)
+			require.Equal(t, "FORCE INDEX ("+expected.indexName+")", suffix)
+
+			prefix, err := hint.SQLPrefix()
+			require.NoError(t, err)
+			require.Empty(t, prefix)
+
+			require.Equal(t, expected.sort, hint.SortOrder())
+		})
+	}
+
+	require.Nil(t, IndexingHintForQueryShape(queryshape.Varying))
+	require.Nil(t, IndexingHintForQueryShape(queryshape.Unspecified))
+}