fix: memoize collectSubfields and getStreamUsage across Executors (#4548)

and add tests!

Author: yaacovCR

src/execution/Executor.ts

@@ -3,8 +3,8 @@ import { invariant } from '../jsutils/invariant.js';
 import { isAsyncIterable } from '../jsutils/isAsyncIterable.js';
 import { isIterableObject } from '../jsutils/isIterableObject.js';
 import { isPromise } from '../jsutils/isPromise.js';
-import { memoize1 } from '../jsutils/memoize1.js';
 import { memoize2 } from '../jsutils/memoize2.js';
+import { memoize3 } from '../jsutils/memoize3.js';
 import type { ObjMap } from '../jsutils/ObjMap.js';
 import type { Path } from '../jsutils/Path.js';
 import { addPath, pathToArray } from '../jsutils/Path.js';
@@ -54,7 +54,7 @@ import {
   collectSubfields as _collectSubfields,
 } from './collectFields.js';
 import type { StreamUsage } from './getStreamUsage.js';
-import { getStreamUsage } from './getStreamUsage.js';
+import { getStreamUsage as _getStreamUsage } from './getStreamUsage.js';
 import type {
   DeferUsageSet,
   ExecutionPlan,
@@ -121,6 +121,37 @@ export interface ValidatedExecutionArgs {
   enableEarlyExecution: boolean;
 }
 
+/**
+ * A memoized collection of relevant subfields with regard to the return
+ * type. Memoizing ensures the subfields are not repeatedly calculated, which
+ * saves overhead when resolving lists of values.
+ */
+export const collectSubfields = memoize3(
+  (
+    validatedExecutionArgs: ValidatedExecutionArgs,
+    returnType: GraphQLObjectType,
+    fieldDetailsList: FieldDetailsList,
+  ) => {
+    const { schema, fragments, variableValues, hideSuggestions } =
+      validatedExecutionArgs;
+    return _collectSubfields(
+      schema,
+      fragments,
+      variableValues,
+      returnType,
+      fieldDetailsList,
+      hideSuggestions,
+    );
+  },
+);
+
+export const getStreamUsage = memoize2(
+  (
+    validatedExecutionArgs: ValidatedExecutionArgs,
+    fieldDetailsList: FieldDetailsList,
+  ) => _getStreamUsage(validatedExecutionArgs, fieldDetailsList),
+);
+
 /**
  * @internal
  */
@@ -384,18 +415,6 @@ export class Executor {
   tasks: Array<ExecutionGroup>;
   streams: Array<ItemStream>;
 
-  collectSubfields: (
-    returnType: GraphQLObjectType,
-    fieldDetailsList: FieldDetailsList,
-  ) => {
-    groupedFieldSet: GroupedFieldSet;
-    newDeferUsages: ReadonlyArray<DeferUsage>;
-  };
-
-  getStreamUsage: (
-    fieldDetailsList: FieldDetailsList,
-  ) => StreamUsage | undefined;
-
   constructor(
     validatedExecutionArgs: ValidatedExecutionArgs,
     deferUsageSet?: DeferUsageSet,
@@ -408,28 +427,6 @@ export class Executor {
     this.groups = [];
     this.tasks = [];
     this.streams = [];
-
-    /**
-     * A memoized collection of relevant subfields with regard to the return
-     * type. Memoizing ensures the subfields are not repeatedly calculated, which
-     * saves overhead when resolving lists of values.
-     */
-    this.collectSubfields = memoize2((returnType, fieldDetailsList) => {
-      const { schema, fragments, variableValues, hideSuggestions } =
-        this.validatedExecutionArgs;
-      return _collectSubfields(
-        schema,
-        fragments,
-        variableValues,
-        returnType,
-        fieldDetailsList,
-        hideSuggestions,
-      );
-    });
-
-    this.getStreamUsage = memoize1((fieldDetailsList) =>
-      getStreamUsage(this.validatedExecutionArgs, fieldDetailsList),
-    );
   }
 
   executeQueryOrMutationOrSubscriptionEvent(): PromiseOrValue<
@@ -1057,7 +1054,7 @@ export class Executor {
     const streamUsage =
       typeof path.key === 'number'
         ? undefined
-        : this.getStreamUsage(fieldDetailsList);
+        : getStreamUsage(this.validatedExecutionArgs, fieldDetailsList);
 
     let containsPromise = false;
     const completedResults: Array<unknown> = [];
@@ -1177,7 +1174,7 @@ export class Executor {
     const streamUsage =
       typeof path.key === 'number'
         ? undefined
-        : this.getStreamUsage(fieldDetailsList);
+        : getStreamUsage(this.validatedExecutionArgs, fieldDetailsList);
 
     // This is specified as a simple map, however we're optimizing the path
     // where the list contains no Promises by avoiding creating another Promise.
@@ -1549,7 +1546,8 @@ export class Executor {
     deliveryGroupMap: ReadonlyMap<DeferUsage, DeliveryGroup> | undefined,
   ): PromiseOrValue<ObjMap<unknown>> {
     // Collect sub-fields to execute to complete this value.
-    const { groupedFieldSet, newDeferUsages } = this.collectSubfields(
+    const { groupedFieldSet, newDeferUsages } = collectSubfields(
+      this.validatedExecutionArgs,
       returnType,
       fieldDetailsList,
     );

src/execution/__tests__/executor-test.ts

@@ -7,6 +7,7 @@ import { resolveOnNextTick } from '../../__testUtils__/resolveOnNextTick.js';
 import { inspect } from '../../jsutils/inspect.js';
 import { promiseWithResolvers } from '../../jsutils/promiseWithResolvers.js';
 
+import type { FieldNode } from '../../language/ast.js';
 import { Kind } from '../../language/kinds.js';
 import { parse } from '../../language/parser.js';
 
@@ -31,7 +32,9 @@ import {
 } from '../../type/scalars.js';
 import { GraphQLSchema } from '../../type/schema.js';
 
-import { execute, executeSync } from '../execute.js';
+import type { FieldDetailsList } from '../collectFields.js';
+import { execute, executeSync, validateExecutionArgs } from '../execute.js';
+import { collectSubfields, getStreamUsage } from '../Executor.js';
 
 describe('Execute: Handles basic execution tasks', () => {
   it('executes arbitrary code', async () => {
@@ -1450,4 +1453,94 @@ describe('Execute: Handles basic execution tasks', () => {
       ],
     });
   });
+
+  it('memoizes collectSubfields results', () => {
+    const deepType = new GraphQLObjectType({
+      name: 'DeepType',
+      fields: {
+        name: { type: GraphQLString },
+      },
+    });
+    const schema = new GraphQLSchema({
+      query: new GraphQLObjectType({
+        name: 'Query',
+        fields: {
+          deep: { type: deepType },
+        },
+      }),
+    });
+    const document = parse('{ deep { name } }');
+    const validatedExecutionArgs = validateExecutionArgs({
+      schema,
+      document,
+    });
+
+    assert('schema' in validatedExecutionArgs);
+
+    const operation = validatedExecutionArgs.operation;
+    const node = operation.selectionSet.selections[0] as FieldNode;
+
+    const fieldDetailsList: FieldDetailsList = [{ node }];
+
+    const first = collectSubfields(
+      validatedExecutionArgs,
+      deepType,
+      fieldDetailsList,
+    );
+
+    const second = collectSubfields(
+      validatedExecutionArgs,
+      deepType,
+      fieldDetailsList,
+    );
+
+    expect(second).to.equal(first);
+
+    const third = collectSubfields(validatedExecutionArgs, deepType, [
+      { node },
+    ]);
+
+    expect(third).to.not.equal(first);
+  });
+
+  it('memoizes getStreamUsage results', () => {
+    const itemType = new GraphQLObjectType({
+      name: 'Item',
+      fields: {
+        id: { type: GraphQLString },
+      },
+    });
+    const schema = new GraphQLSchema({
+      query: new GraphQLObjectType({
+        name: 'Query',
+        fields: {
+          items: { type: new GraphQLList(itemType) },
+        },
+      }),
+      directives: [GraphQLStreamDirective],
+    });
+    const document = parse('{ items @stream(initialCount: 1) { id } }');
+    const validatedExecutionArgs = validateExecutionArgs({
+      schema,
+      document,
+    });
+
+    assert('schema' in validatedExecutionArgs);
+
+    const operation = validatedExecutionArgs.operation;
+    const node = operation.selectionSet.selections[0] as FieldNode;
+
+    const fieldDetailsList = [{ node }];
+    const first = getStreamUsage(validatedExecutionArgs, fieldDetailsList);
+
+    expect(first).to.not.equal(undefined);
+
+    const second = getStreamUsage(validatedExecutionArgs, fieldDetailsList);
+
+    expect(second).to.equal(first);
+
+    const third = getStreamUsage(validatedExecutionArgs, [{ node }]);
+
+    expect(third).to.not.equal(first);
+  });
 });

src/jsutils/memoize1.ts

@@ -0,0 +1,34 @@
+/**
+ * Memoizes the provided three-argument function.
+ */
+export function memoize3<
+  A1 extends object,
+  A2 extends object,
+  A3 extends object,
+  R,
+>(fn: (a1: A1, a2: A2, a3: A3) => R): (a1: A1, a2: A2, a3: A3) => R {
+  let cache0: WeakMap<A1, WeakMap<A2, WeakMap<A3, R>>>;
+
+  return function memoized(a1, a2, a3) {
+    cache0 ??= new WeakMap();
+
+    let cache1 = cache0.get(a1);
+    if (cache1 === undefined) {
+      cache1 = new WeakMap();
+      cache0.set(a1, cache1);
+    }
+
+    let cache2 = cache1.get(a2);
+    if (cache2 === undefined) {
+      cache2 = new WeakMap();
+      cache1.set(a2, cache2);
+    }
+
+    let fnResult = cache2.get(a3);
+    if (fnResult === undefined) {
+      fnResult = fn(a1, a2, a3);
+      cache2.set(a3, fnResult);
+    }
+    return fnResult;
+  };
+}