From d3a5378dd8fbd5f6541b7f4a99f21717c593eefa Mon Sep 17 00:00:00 2001
From: Joseph Schorr <josephschorr@users.noreply.github.com>
Date: Thu, 8 Feb 2024 15:09:44 -0500
Subject: [PATCH] Reduce memory usage of WriteSchema

We now only annotate the namespaces that have been changed or added, as annotation and canonicalization uses a significant amount of memory
---
 internal/services/shared/schema.go      | 35 ++++++++++++++++---------
 internal/services/shared/schema_test.go | 12 ++++++++-
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/internal/services/shared/schema.go b/internal/services/shared/schema.go
index 42ed4d353e..92a84b64b4 100644
--- a/internal/services/shared/schema.go
+++ b/internal/services/shared/schema.go
@@ -12,16 +12,18 @@ import (
 	"github.com/authzed/spicedb/pkg/genutil/mapz"
 	core "github.com/authzed/spicedb/pkg/proto/core/v1"
 	"github.com/authzed/spicedb/pkg/schemadsl/compiler"
+	"github.com/authzed/spicedb/pkg/spiceerrors"
 	"github.com/authzed/spicedb/pkg/tuple"
 	"github.com/authzed/spicedb/pkg/typesystem"
 )
 
 // ValidatedSchemaChanges is a set of validated schema changes that can be applied to the datastore.
 type ValidatedSchemaChanges struct {
-	compiled          *compiler.CompiledSchema
-	newCaveatDefNames *mapz.Set[string]
-	newObjectDefNames *mapz.Set[string]
-	additiveOnly      bool
+	compiled             *compiler.CompiledSchema
+	validatedTypeSystems map[string]*typesystem.ValidatedNamespaceTypeSystem
+	newCaveatDefNames    *mapz.Set[string]
+	newObjectDefNames    *mapz.Set[string]
+	additiveOnly         bool
 }
 
 // ValidateSchemaChanges validates the schema found in the compiled schema and returns a
@@ -39,6 +41,8 @@ func ValidateSchemaChanges(ctx context.Context, compiled *compiler.CompiledSchem
 
 	// 2) Validate the namespaces defined.
 	newObjectDefNames := mapz.NewSet[string]()
+	validatedTypeSystems := make(map[string]*typesystem.ValidatedNamespaceTypeSystem, len(compiled.ObjectDefinitions))
+
 	for _, nsdef := range compiled.ObjectDefinitions {
 		ts, err := typesystem.NewNamespaceTypeSystem(nsdef,
 			typesystem.ResolverForPredefinedDefinitions(typesystem.PredefinedElements{
@@ -54,18 +58,16 @@ func ValidateSchemaChanges(ctx context.Context, compiled *compiler.CompiledSchem
 			return nil, err
 		}
 
-		if err := namespace.AnnotateNamespace(vts); err != nil {
-			return nil, err
-		}
-
+		validatedTypeSystems[nsdef.Name] = vts
 		newObjectDefNames.Insert(nsdef.Name)
 	}
 
 	return &ValidatedSchemaChanges{
-		compiled:          compiled,
-		newCaveatDefNames: newCaveatDefNames,
-		newObjectDefNames: newObjectDefNames,
-		additiveOnly:      additiveOnly,
+		compiled:             compiled,
+		validatedTypeSystems: validatedTypeSystems,
+		newCaveatDefNames:    newCaveatDefNames,
+		newObjectDefNames:    newObjectDefNames,
+		additiveOnly:         additiveOnly,
 	}, nil
 }
 
@@ -156,6 +158,15 @@ func ApplySchemaChangesOverExisting(
 
 		if len(diff.Deltas()) > 0 {
 			objectDefsWithChanges = append(objectDefsWithChanges, nsdef)
+
+			vts, ok := validated.validatedTypeSystems[nsdef.Name]
+			if !ok {
+				return nil, spiceerrors.MustBugf("validated type system not found for namespace `%s`", nsdef.Name)
+			}
+
+			if err := namespace.AnnotateNamespace(vts); err != nil {
+				return nil, err
+			}
 		}
 	}
 
diff --git a/internal/services/shared/schema_test.go b/internal/services/shared/schema_test.go
index 209e9c373d..78012466f4 100644
--- a/internal/services/shared/schema_test.go
+++ b/internal/services/shared/schema_test.go
@@ -38,7 +38,10 @@ func TestApplySchemaChanges(t *testing.T) {
 		SchemaString: `
 			definition user {}
 
-			definition organization {}
+			definition organization {
+				relation member: user
+				permission admin = member
+			}
 
 			caveat catchTwentyTwo(value int) {
 			  value == 22
@@ -58,6 +61,13 @@ func TestApplySchemaChanges(t *testing.T) {
 		require.Equal(applied.RemovedObjectDefNames, []string{"document"})
 		require.Equal(applied.NewCaveatDefNames, []string{"catchTwentyTwo"})
 		require.Equal(applied.RemovedCaveatDefNames, []string{"hasFortyTwo"})
+
+		orgDef, err := rwt.LookupNamespacesWithNames(ctx, []string{"organization"})
+		require.NoError(err)
+		require.Len(orgDef, 1)
+
+		require.NotEmpty(orgDef[0].Definition.Relation[0].CanonicalCacheKey)
+		require.NotEmpty(orgDef[0].Definition.Relation[1].CanonicalCacheKey)
 		return nil
 	})
 	require.NoError(err)