Make SchemaUpgrader faster

.changeset/rotten-clocks-mate.md

@@ -0,0 +1,5 @@
+---
+"@apollo/federation-internals": patch
+---
+
+Save time in SchemaUpgrader by pre-computing which subgraphs contain each type

internals-js/src/schemaUpgrader.ts

@@ -11,6 +11,7 @@ import {
   Directive,
   Extension,
   FieldDefinition,
+  InterfaceType,
   isCompositeType,
   isInterfaceType,
   isObjectType,
@@ -230,15 +231,36 @@ export function upgradeSubgraphsIfNecessary(inputs: Subgraphs): UpgradeResult {
   const subgraphs = new Subgraphs();
   let errors: GraphQLError[] = [];
   const subgraphsUsingInterfaceObject = [];
+
+  // build a data structure to help us do computation only once
+  const objectTypeMap = new Map<string, Map<string, [ObjectType | InterfaceType, FederationMetadata]>>();
+  for (const subgraph of inputs.values()) {
+    for (const t of subgraph.schema.objectTypes()) {
+      let entry = objectTypeMap.get(t.name);
+      if (!entry) {
+        entry = new Map();
+        objectTypeMap.set(t.name, entry);
+      }
+      entry.set(subgraph.name, [t, subgraph.metadata()]);
+    }
+    for (const t of subgraph.schema.interfaceTypes()) {
+      let entry = objectTypeMap.get(t.name);
+      if (!entry) {
+        entry = new Map();
+        objectTypeMap.set(t.name, entry);
+      }
+      entry.set(subgraph.name, [t, subgraph.metadata()]);
+    }
+  }
+
   for (const subgraph of inputs.values()) {
     if (subgraph.isFed2Subgraph()) {
       subgraphs.add(subgraph);
       if (subgraph.metadata().interfaceObjectDirective().applications().size > 0) {
         subgraphsUsingInterfaceObject.push(subgraph.name);
       }
     } else {
-      const otherSubgraphs = inputs.values().filter((s) => s.name !== subgraph.name);
-      const res = new SchemaUpgrader(subgraph, otherSubgraphs).upgrade();
+      const res = new SchemaUpgrader(subgraph, inputs.values(), objectTypeMap).upgrade();
       if (res.errors) {
         errors = errors.concat(res.errors);
       } else {
@@ -291,16 +313,6 @@ function isRootTypeExtension(type: NamedType): boolean {
     && (type.hasAppliedDirective(metadata.extendsDirective()) || (type.hasExtensionElements() && !type.hasNonExtensionElements()));
 }
 
-function resolvesField(subgraph: Subgraph, field: FieldDefinition<ObjectType>): boolean  {
-  const metadata = subgraph.metadata();
-  const t = subgraph.schema.type(field.parent.name);
-  if (!t || !isObjectType(t)) {
-    return false;
-  }
-  const f = t.field(field.name);
-  return !!f && (!metadata.isFieldExternal(f) || metadata.isFieldPartiallyExternal(f));
-}
-
 function getField(schema: Schema, typeName: string, fieldName: string): FieldDefinition<CompositeType> | undefined {
   const type = schema.type(typeName);
   return type && isCompositeType(type) ? type.field(fieldName) : undefined;
@@ -313,7 +325,7 @@ class SchemaUpgrader {
   private readonly metadata: FederationMetadata;
   private readonly errors: GraphQLError[] = [];
 
-  constructor(private readonly originalSubgraph: Subgraph, private readonly otherSubgraphs: Subgraph[]) {
+  constructor(private readonly originalSubgraph: Subgraph, private readonly allSubgraphs: readonly Subgraph[], private readonly objectTypeMap: Map<string, Map<string, [ObjectType | InterfaceType, FederationMetadata]>>) {
     // Note that as we clone the original schema, the 'sourceAST' values in the elements of the new schema will be those of the original schema
     // and those won't be updated as we modify the schema to make it fed2-enabled. This is _important_ for us here as this is what ensures that
     // later merge errors "AST" nodes ends up pointing to the original schema, the one that make sense to the user.
@@ -380,8 +392,9 @@ class SchemaUpgrader {
     }
 
     const extensionAST = firstOf<Extension<any>>(type.extensions().values())?.sourceAST;
-    for (const subgraph of this.otherSubgraphs) {
-      const otherType = subgraph.schema.type(type.name);
+    const typeInOtherSubgraphs = Array.from(this.objectTypeMap.get(type.name)!.entries()).filter(([subgraphName, _]) => subgraphName !== this.subgraph.name);
+    for (let i = 0; i < typeInOtherSubgraphs.length; i += 1) {
+      const otherType = typeInOtherSubgraphs[i][1][0];
       if (otherType && otherType.hasNonExtensionElements()) {
         return;
       }
@@ -589,13 +602,12 @@ class SchemaUpgrader {
         // case territory in the first place, so this is probably good enough (that is, there is
         // customer schema for which what we do here matter but not that I know of for which it's
         // not good enough).
-        for (const other of this.otherSubgraphs) {
-          const typeInOther = other.schema.type(type.name);
-          if (!typeInOther) {
-            continue;
-          }
-          assert(isCompositeType(typeInOther), () => `Type ${type} is of kind ${type.kind} in ${this.subgraph.name} but ${typeInOther.kind} in ${other.name}`);
-          const keysInOther = typeInOther.appliedDirectivesOf(other.metadata().keyDirective());
+        const typeInOtherSubgraphs = Array.from(this.objectTypeMap.get(type.name)!.entries()).filter(([subgraphName, _]) => subgraphName !== this.subgraph.name);
+
+        for (const [otherSubgraphName, v] of typeInOtherSubgraphs) {
+          const [typeInOther, metadata] = v;
+          assert(isCompositeType(typeInOther), () => `Type ${type} is of kind ${type.kind} in ${this.subgraph.name} but ${typeInOther.kind} in ${otherSubgraphName}`);
+          const keysInOther = typeInOther.appliedDirectivesOf(metadata.keyDirective());
           if (keysInOther.length === 0) {
             continue;
           }
@@ -710,17 +722,26 @@ class SchemaUpgrader {
           if (originalMetadata.isFieldShareable(field)) {
             continue;
           }
-          const otherResolvingSubgraphs = this.otherSubgraphs.filter((s) => resolvesField(s, field));
-          if (otherResolvingSubgraphs.length > 0 && !field.hasAppliedDirective(shareableDirective)) {
+
+          const entries = Array.from(this.objectTypeMap.get(type.name)!.entries());
+          const typeInOtherSubgraphs = entries.filter(([subgraphName, v]) => {
+            if (subgraphName === this.subgraph.name) {
+              return false;
+            }
+            const f = v[0].field(field.name);
+            return !!f && (!v[1].isFieldExternal(f) || v[1].isFieldPartiallyExternal(f));
+          });
+
+          if (typeInOtherSubgraphs.length > 0 && !field.hasAppliedDirective(shareableDirective)) {
             field.applyDirective(shareableDirective);
-            this.addChange(new ShareableFieldAddition(field.coordinate, otherResolvingSubgraphs.map((s) => s.name)));
+            this.addChange(new ShareableFieldAddition(field.coordinate, typeInOtherSubgraphs.map(([s]) => s)));
           }
         }
       } else {
-        const otherDeclaringSubgraphs = this.otherSubgraphs.filter((s) => s.schema.type(type.name));
-        if (otherDeclaringSubgraphs.length > 0 && !type.hasAppliedDirective(shareableDirective)) {
+        const typeInOtherSubgraphs = Array.from(this.objectTypeMap.get(type.name)!.entries()).filter(([subgraphName, _]) => subgraphName !== this.subgraph.name);
+        if (typeInOtherSubgraphs.length > 0 && !type.hasAppliedDirective(shareableDirective)) {
           type.applyDirective(shareableDirective);
-          this.addChange(new ShareableTypeAddition(type.coordinate, otherDeclaringSubgraphs.map((s) => s.name)));
+          this.addChange(new ShareableTypeAddition(type.coordinate, typeInOtherSubgraphs.map(([s]) => s)));
         }
       }
     }
@@ -739,8 +760,8 @@ class SchemaUpgrader {
         continue;
       }
       if (this.external(element)) {
-        const tagIsUsedInOtherDefinition = this.otherSubgraphs
-          .map((s) => getField(s.schema, element.parent.name, element.name))
+        const tagIsUsedInOtherDefinition = this.allSubgraphs
+          .map((s) => s.name === this.originalSubgraph.name ? undefined : getField(s.schema, element.parent.name, element.name))
           .filter((f) => !(f && f.hasAppliedDirective('external')))
           .some((f) => f && f.appliedDirectivesOf('tag').some((d) => valueEquals(application.arguments(), d.arguments())));