Use type information to optimize TOUCH operations in the PG datastore

When a TOUCH is performed over a relationship for a relation+subject type that cannot be caveated, we can skip the UPDATE query used to check for the caveat changing, which is a significant performance boost

internal/datastore/postgres/readwrite.go

@@ -7,7 +7,9 @@ import (
 	"strings"
 
 	"github.com/authzed/spicedb/pkg/datastore/options"
+	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	"github.com/authzed/spicedb/pkg/spiceerrors"
+	"github.com/authzed/spicedb/pkg/typesystem"
 
 	sq "github.com/Masterminds/squirrel"
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
@@ -88,14 +90,88 @@ func appendForInsertion(builder sq.InsertBuilder, tpl *core.RelationTuple) sq.In
 	return builder.Values(valuesToWrite...)
 }
 
+func (rwt *pgReadWriteTXN) collectSimplifiedTouchTypes(ctx context.Context, mutations []*core.RelationTupleUpdate) (*mapz.Set[string], error) {
+	// Collect the list of namespaces used for resources for relationships being TOUCHed.
+	touchedResourceNamespaces := mapz.NewSet[string]()
+	for _, mut := range mutations {
+		if mut.Operation == core.RelationTupleUpdate_TOUCH {
+			touchedResourceNamespaces.Add(mut.Tuple.ResourceAndRelation.Namespace)
+		}
+	}
+
+	// Load the namespaces for any resources that are being TOUCHed and check if the relation being touched
+	// *can* have a caveat. If not, mark the relation as supported simplified TOUCH operations.
+	relationSupportSimplifiedTouch := mapz.NewSet[string]()
+	if touchedResourceNamespaces.IsEmpty() {
+		return relationSupportSimplifiedTouch, nil
+	}
+
+	namespaces, err := rwt.LookupNamespacesWithNames(ctx, touchedResourceNamespaces.AsSlice())
+	if err != nil {
+		return nil, fmt.Errorf(errUnableToWriteRelationships, err)
+	}
+
+	if len(namespaces) == 0 {
+		return relationSupportSimplifiedTouch, nil
+	}
+
+	nsDefByName := make(map[string]*core.NamespaceDefinition, len(namespaces))
+	for _, ns := range namespaces {
+		nsDefByName[ns.Definition.Name] = ns.Definition
+	}
+
+	for _, mut := range mutations {
+		tpl := mut.Tuple
+
+		if mut.Operation != core.RelationTupleUpdate_TOUCH {
+			continue
+		}
+
+		nsDef, ok := nsDefByName[tpl.ResourceAndRelation.Namespace]
+		if !ok {
+			continue
+		}
+
+		vts, err := typesystem.NewNamespaceTypeSystem(nsDef, typesystem.ResolverForDatastoreReader(rwt))
+		if err != nil {
+			return nil, fmt.Errorf(errUnableToWriteRelationships, err)
+		}
+
+		notAllowed, err := vts.RelationDoesNotAllowCaveatsForSubject(tpl.ResourceAndRelation.Relation, tpl.Subject.Namespace)
+		if err != nil {
+			// Ignore errors and just fallback to the less efficient path.
+			continue
+		}
+
+		if notAllowed {
+			relationSupportSimplifiedTouch.Add(nsDef.Name + "#" + tpl.ResourceAndRelation.Relation + "@" + tpl.Subject.Namespace)
+			continue
+		}
+	}
+
+	return relationSupportSimplifiedTouch, nil
+}
+
 func (rwt *pgReadWriteTXN) WriteRelationships(ctx context.Context, mutations []*core.RelationTupleUpdate) error {
 	touchMutationsByNonCaveat := make(map[string]*core.RelationTupleUpdate, len(mutations))
 	hasCreateInserts := false
 
 	createInserts := writeTuple
 	touchInserts := writeTuple
+
 	deleteClauses := sq.Or{}
 
+	// Determine the set of relation+subject types for whom a "simplified" TOUCH operation can be used. A
+	// simplified TOUCH operation is one in which the relationship does not support caveats for the subject
+	// type. In such cases, the "DELETE" operation is unnecessary because the relationship does not support
+	// caveats for the subject type, and thus the relationship can be TOUCHed without needing to check for
+	// the existence of a relationship with a different caveat name and/or context which might need to be
+	// replaced.
+	relationSupportSimplifiedTouch, err := rwt.collectSimplifiedTouchTypes(ctx, mutations)
+	if err != nil {
+		return err
+	}
+
 	// Parse the updates, building inserts for CREATE/TOUCH and deletes for DELETE.
 	for _, mut := range mutations {
 		tpl := mut.Tuple
@@ -106,8 +182,8 @@ func (rwt *pgReadWriteTXN) WriteRelationships(ctx context.Context, mutations []*
 			hasCreateInserts = true
 
 		case core.RelationTupleUpdate_TOUCH:
-			touchMutationsByNonCaveat[tuple.StringWithoutCaveat(tpl)] = mut
 			touchInserts = appendForInsertion(touchInserts, tpl)
+			touchMutationsByNonCaveat[tuple.StringWithoutCaveat(tpl)] = mut
 
 		case core.RelationTupleUpdate_DELETE:
 			deleteClauses = append(deleteClauses, exactRelationshipClause(tpl))
@@ -184,18 +260,24 @@ func (rwt *pgReadWriteTXN) WriteRelationships(ctx context.Context, mutations []*
 		}
 		rows.Close()
 
-		// For each remaining TOUCH mutation, add a DELETE operation for the row iff the caveat and/or
-		// context has changed. For ones in which the caveat name and/or context did not change, there is
-		// no need to replace the row, as it is already present.
+		// For each remaining TOUCH mutation, add a "DELETE" operation for the row such that if the caveat and/or
+		// context has changed, the row will be deleted. For ones in which the caveat name and/or context did cause
+		// the deletion (because of a change), the row will be re-inserted with the new caveat name and/or context.
 		for _, mut := range touchMutationsByNonCaveat {
+			// If the relation support a simplified TOUCH operation, then skip the DELETE operation, as it is unnecessary
+			// because the relation does not support a caveat for a subject of this type.
+			if relationSupportSimplifiedTouch.Has(mut.Tuple.ResourceAndRelation.Namespace + "#" + mut.Tuple.ResourceAndRelation.Relation + "@" + mut.Tuple.Subject.Namespace) {
+				continue
+			}
+
 			deleteClauses = append(deleteClauses, exactRelationshipDifferentCaveatClause(mut.Tuple))
 		}
 	}
 
-	// Execute the DELETE operation for any DELETE mutations or TOUCH mutations that matched existing
-	// relationships and whose caveat name or context is different in some manner. We use RETURNING
-	// to determine which TOUCHed relationships were deleted by virtue of their caveat name and/or
-	// context being changed.
+	// Execute the "DELETE" operation (an UPDATE with setting the deletion ID to the current transaction ID)
+	// for any DELETE mutations or TOUCH mutations that matched existing relationships and whose caveat name
+	// or context is different in some manner. We use RETURNING to determine which TOUCHed relationships were
+	// deleted by virtue of their caveat name and/or context being changed.
 	if len(deleteClauses) == 0 {
 		// Nothing more to do.
 		return nil

pkg/datastore/test/datastore.go

@@ -97,7 +97,7 @@ func AllWithExceptions(t *testing.T, tester DatastoreTester, except Categories)
 	t.Run("TestDeleteAlreadyDeleted", func(t *testing.T) { DeleteAlreadyDeletedTest(t, tester) })
 	t.Run("TestWriteDeleteWrite", func(t *testing.T) { WriteDeleteWriteTest(t, tester) })
 	t.Run("TestCreateAlreadyExisting", func(t *testing.T) { CreateAlreadyExistingTest(t, tester) })
-	t.Run("TestTouchAlreadyExisting", func(t *testing.T) { TouchAlreadyExistingTest(t, tester) })
+	t.Run("TestTouchAlreadyExistingWithoutCaveat", func(t *testing.T) { TouchAlreadyExistingTest(t, tester) })
 	t.Run("TestCreateDeleteTouch", func(t *testing.T) { CreateDeleteTouchTest(t, tester) })
 	t.Run("TestDeleteOneThousandIndividualInOneCall", func(t *testing.T) { DeleteOneThousandIndividualInOneCallTest(t, tester) })
 	t.Run("TestCreateTouchDeleteTouch", func(t *testing.T) { CreateTouchDeleteTouchTest(t, tester) })
@@ -107,6 +107,8 @@ func AllWithExceptions(t *testing.T, tester DatastoreTester, except Categories)
 	t.Run("TestDeleteWithLimit", func(t *testing.T) { DeleteWithLimitTest(t, tester) })
 	t.Run("TestQueryRelationshipsWithVariousFilters", func(t *testing.T) { QueryRelationshipsWithVariousFiltersTest(t, tester) })
 	t.Run("TestDeleteRelationshipsWithVariousFilters", func(t *testing.T) { DeleteRelationshipsWithVariousFiltersTest(t, tester) })
+	t.Run("TestTouchTypedAlreadyExistingWithoutCaveat", func(t *testing.T) { TypedTouchAlreadyExistingTest(t, tester) })
+	t.Run("TestTouchTypedAlreadyExistingWithCaveat", func(t *testing.T) { TypedTouchAlreadyExistingWithCaveatTest(t, tester) })
 
 	t.Run("TestMultipleReadsInRWT", func(t *testing.T) { MultipleReadsInRWTTest(t, tester) })
 	t.Run("TestConcurrentWriteSerialization", func(t *testing.T) { ConcurrentWriteSerializationTest(t, tester) })

pkg/datastore/test/tuples.go

@@ -1319,6 +1319,50 @@ func QueryRelationshipsWithVariousFiltersTest(t *testing.T, tester DatastoreTest
 	}
 }
 
+// TypedTouchAlreadyExistingTest tests touching a relationship twice, when valid type information is provided.
+func TypedTouchAlreadyExistingTest(t *testing.T, tester DatastoreTester) {
+	require := require.New(t)
+
+	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
+	require.NoError(err)
+
+	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
+	ctx := context.Background()
+
+	tpl1 := tuple.Parse("document:foo#viewer@user:tom")
+
+	_, err = common.WriteTuples(ctx, ds, core.RelationTupleUpdate_TOUCH, tpl1)
+	require.NoError(err)
+	ensureTuples(ctx, require, ds, tpl1)
+
+	_, err = common.WriteTuples(ctx, ds, core.RelationTupleUpdate_TOUCH, tpl1)
+	require.NoError(err)
+	ensureTuples(ctx, require, ds, tpl1)
+}
+
+// TypedTouchAlreadyExistingWithCaveatTest tests touching a relationship twice, when valid type information is provided.
+func TypedTouchAlreadyExistingWithCaveatTest(t *testing.T, tester DatastoreTester) {
+	require := require.New(t)
+
+	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
+	require.NoError(err)
+
+	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
+	ctx := context.Background()
+
+	ctpl1 := tuple.Parse("document:foo#caveated_viewer@user:tom[test:{\"foo\":\"bar\"}]")
+
+	_, err = common.WriteTuples(ctx, ds, core.RelationTupleUpdate_TOUCH, ctpl1)
+	require.NoError(err)
+	ensureTuples(ctx, require, ds, ctpl1)
+
+	ctpl1Updated := tuple.Parse("document:foo#caveated_viewer@user:tom[test:{\"foo\":\"baz\"}]")
+
+	_, err = common.WriteTuples(ctx, ds, core.RelationTupleUpdate_TOUCH, ctpl1Updated)
+	require.NoError(err)
+	ensureTuples(ctx, require, ds, ctpl1Updated)
+}
+
 // CreateTouchDeleteTouchTest tests writing a relationship, touching it, deleting it, and then touching it.
 func CreateTouchDeleteTouchTest(t *testing.T, tester DatastoreTester) {
 	require := require.New(t)

pkg/typesystem/typesystem.go

@@ -594,6 +594,32 @@ func (nts *TypeSystem) typeSystemForNamespace(ctx context.Context, namespaceName
 	return NewNamespaceTypeSystem(nsDef, nts.resolver)
 }
 
+// RelationDoesNotAllowCaveatsForSubject returns true if and only if it can be conclusively determined that
+// the given subject type does not accept any caveats on the given relation. If the relation does not have type information,
+// returns an error.
+func (nts *TypeSystem) RelationDoesNotAllowCaveatsForSubject(relationName string, subjectTypeName string) (bool, error) {
+	relation, ok := nts.relationMap[relationName]
+	if !ok {
+		return false, NewRelationNotFoundErr(nts.nsDef.Name, relationName)
+	}
+
+	typeInfo := relation.GetTypeInformation()
+	if typeInfo == nil {
+		return false, NewTypeErrorWithSource(
+			fmt.Errorf("relation `%s` does not have type information", relationName),
+			relation, relationName,
+		)
+	}
+
+	for _, allowedRelation := range typeInfo.GetAllowedDirectRelations() {
+		if allowedRelation.GetNamespace() == subjectTypeName && allowedRelation.GetRequiredCaveat() != nil && allowedRelation.GetRequiredCaveat().CaveatName != "" {
+			return false, nil
+		}
+	}
+
+	return true, nil
+}
+
 // ValidatedNamespaceTypeSystem is validated type system for a namespace.
 type ValidatedNamespaceTypeSystem struct {
 	*TypeSystem

pkg/typesystem/typesystem_test.go

@@ -468,6 +468,16 @@ func TestTypeSystemAccessors(t *testing.T) {
 						require.True(t, vts.IsPermission("edit"))
 					})
 
+					t.Run("RelationDoesNotAllowCaveatsForSubject", func(t *testing.T) {
+						ok, err := vts.RelationDoesNotAllowCaveatsForSubject("viewer", "user")
+						require.NoError(t, err)
+						require.True(t, ok)
+
+						ok, err = vts.RelationDoesNotAllowCaveatsForSubject("editor", "user")
+						require.NoError(t, err)
+						require.True(t, ok)
+					})
+
 					t.Run("IsAllowedPublicNamespace", func(t *testing.T) {
 						require.Equal(t, PublicSubjectNotAllowed, noError(vts.IsAllowedPublicNamespace("editor", "user")))
 						require.Equal(t, PublicSubjectNotAllowed, noError(vts.IsAllowedPublicNamespace("viewer", "user")))
@@ -544,6 +554,16 @@ func TestTypeSystemAccessors(t *testing.T) {
 						require.True(t, vts.IsPermission("view"))
 					})
 
+					t.Run("RelationDoesNotAllowCaveatsForSubject", func(t *testing.T) {
+						ok, err := vts.RelationDoesNotAllowCaveatsForSubject("viewer", "user")
+						require.NoError(t, err)
+						require.True(t, ok)
+
+						ok, err = vts.RelationDoesNotAllowCaveatsForSubject("editor", "user")
+						require.NoError(t, err)
+						require.True(t, ok)
+					})
+
 					t.Run("IsAllowedPublicNamespace", func(t *testing.T) {
 						require.Equal(t, PublicSubjectNotAllowed, noError(vts.IsAllowedPublicNamespace("editor", "user")))
 						require.Equal(t, PublicSubjectAllowed, noError(vts.IsAllowedPublicNamespace("viewer", "user")))
@@ -603,6 +623,16 @@ func TestTypeSystemAccessors(t *testing.T) {
 						require.False(t, vts.IsPermission("member"))
 					})
 
+					t.Run("RelationDoesNotAllowCaveatsForSubject", func(t *testing.T) {
+						ok, err := vts.RelationDoesNotAllowCaveatsForSubject("member", "user")
+						require.NoError(t, err)
+						require.True(t, ok)
+
+						ok, err = vts.RelationDoesNotAllowCaveatsForSubject("member", "group")
+						require.NoError(t, err)
+						require.True(t, ok)
+					})
+
 					t.Run("IsAllowedPublicNamespace", func(t *testing.T) {
 						require.Equal(t, PublicSubjectNotAllowed, noError(vts.IsAllowedPublicNamespace("member", "user")))
 					})
@@ -664,6 +694,20 @@ func TestTypeSystemAccessors(t *testing.T) {
 						require.False(t, vts.IsPermission("onlycaveated"))
 					})
 
+					t.Run("RelationDoesNotAllowCaveatsForSubject", func(t *testing.T) {
+						ok, err := vts.RelationDoesNotAllowCaveatsForSubject("viewer", "user")
+						require.NoError(t, err)
+						require.False(t, ok)
+
+						ok, err = vts.RelationDoesNotAllowCaveatsForSubject("editor", "user")
+						require.NoError(t, err)
+						require.True(t, ok)
+
+						ok, err = vts.RelationDoesNotAllowCaveatsForSubject("onlycaveated", "user")
+						require.NoError(t, err)
+						require.False(t, ok)
+					})
+
 					t.Run("IsAllowedPublicNamespace", func(t *testing.T) {
 						require.Equal(t, PublicSubjectNotAllowed, noError(vts.IsAllowedPublicNamespace("editor", "user")))
 						require.Equal(t, PublicSubjectNotAllowed, noError(vts.IsAllowedPublicNamespace("viewer", "user")))