Improve performance

Author: ardatan

packages/plugins/graphql-jit/src/index.ts

@@ -1,22 +1,29 @@
 /* eslint-disable no-console */
-import { Plugin } from '@envelop/core';
-import { DocumentNode, Source, ExecutionArgs, ExecutionResult } from 'graphql';
-import { compileQuery, isCompiledQuery, CompilerOptions, CompiledQuery } from 'graphql-jit';
+import { EnvelopError, Plugin, PromiseOrValue } from '@envelop/core';
+import { DocumentNode, Source, ExecutionArgs, ExecutionResult, print, GraphQLSchema, getOperationAST } from 'graphql';
+import { compileQuery, CompilerOptions, CompiledQuery } from 'graphql-jit';
 import lru from 'tiny-lru';
 
 const DEFAULT_MAX = 1000;
 const DEFAULT_TTL = 3600000;
 
-type JITCacheEntry = {
-  query: CompiledQuery['query'];
-  subscribe?: CompiledQuery['subscribe'];
-};
+type CompilationResult = ReturnType<typeof compileQuery>;
 
 export interface JITCache {
-  get(key: string): JITCacheEntry | undefined;
-  set(key: string, value: JITCacheEntry): void;
+  get(key: string): CompilationResult | undefined;
+  set(key: string, value: CompilationResult): void;
+}
+
+function isSource(source: any): source is Source {
+  return source.body != null;
+}
+
+function isCompiledQuery(compilationResult: any): compilationResult is CompiledQuery<unknown, unknown> {
+  return compilationResult.query != null;
 }
 
+const MUST_DEFINE_OPERATION_NAME_ERR = 'Must provide operation name if query contains multiple operations.';
+
 export const useGraphQlJit = (
   compilerOptions: Partial<CompilerOptions> = {},
   pluginOptions: {
@@ -36,74 +43,89 @@ export const useGraphQlJit = (
 ): Plugin => {
   const documentSourceMap = new WeakMap<DocumentNode, string>();
   const jitCache =
-    typeof pluginOptions.cache !== 'undefined' ? pluginOptions.cache : lru<JITCacheEntry>(DEFAULT_MAX, DEFAULT_TTL);
-
-  function compile(args: ExecutionArgs) {
-    const compilationResult = compileQuery(args.schema, args.document, args.operationName ?? undefined, compilerOptions);
-
-    if (isCompiledQuery(compilationResult)) {
-      return compilationResult;
-    } else {
-      if (pluginOptions?.onError) {
-        pluginOptions.onError(compilationResult);
-      } else {
-        console.error(compilationResult);
+    typeof pluginOptions.cache !== 'undefined' ? pluginOptions.cache : lru<CompilationResult>(DEFAULT_MAX, DEFAULT_TTL);
+
+  function getCacheKey(document: DocumentNode, operationName?: string) {
+    if (!operationName) {
+      const operationAST = getOperationAST(document);
+      if (!operationAST) {
+        throw new EnvelopError(MUST_DEFINE_OPERATION_NAME_ERR);
       }
-      return {
-        query: () => compilationResult,
-      };
+      operationName = operationAST.name?.value;
     }
+    const documentSource = getDocumentSource(document);
+    return `${documentSource}::${operationName}`;
   }
 
-  function getCacheEntry(args: ExecutionArgs): JITCacheEntry {
-    const documentSource = documentSourceMap.get(args.document);
-
+  function getDocumentSource(document: DocumentNode): string {
+    let documentSource = documentSourceMap.get(document);
     if (!documentSource) {
-      return compile(args);
+      documentSource = print(document);
+      documentSourceMap.set(document, documentSource);
     }
+    return documentSource;
+  }
 
-    let compiledQuery = jitCache.get(documentSource);
+  function getCompilationResult(schema: GraphQLSchema, document: DocumentNode, operationName?: string) {
+    const cacheKey = getCacheKey(document, operationName);
+
+    let compiledQuery = jitCache.get(cacheKey);
 
     if (!compiledQuery) {
-      compiledQuery = compile(args);
-      jitCache.set(documentSource, compiledQuery);
+      compiledQuery = compileQuery(schema, document, operationName, compilerOptions);
+      jitCache.set(cacheKey, compiledQuery);
     }
 
     return compiledQuery;
   }
 
-  function jitExecute(args: ExecutionArgs) {
-    const cacheEntry = getCacheEntry(args);
-
-    return cacheEntry.query(args.rootValue, args.contextValue, args.variableValues);
-  }
-
-  function jitSubscribe(args: ExecutionArgs) {
-    const cacheEntry = getCacheEntry(args);
-
-    return cacheEntry.subscribe
-      ? (cacheEntry.subscribe(args.rootValue, args.contextValue, args.variableValues) as any)
-      : cacheEntry.query(args.rootValue, args.contextValue, args.variableValues);
-  }
-
   return {
     onParse({ params: { source } }) {
-      const key = source instanceof Source ? source.body : source;
+      const key = isSource(source) ? source.body : source;
 
       return ({ result }) => {
         if (!result || result instanceof Error) return;
 
         documentSourceMap.set(result, key);
       };
     },
-    async onExecute({ args, setExecuteFn }) {
+    onValidate({ params, setResult }) {
+      try {
+        const compilationResult = getCompilationResult(params.schema, params.documentAST);
+        if (!isCompiledQuery(compilationResult) && compilationResult?.errors != null) {
+          setResult(compilationResult.errors);
+        }
+      } catch (e: any) {
+        // Validate doesn't work in case of multiple operations
+        if (e.message !== MUST_DEFINE_OPERATION_NAME_ERR) {
+          throw e;
+        }
+      }
+    },
+    async onExecute({ args, setExecuteFn, setResultAndStopExecution }) {
       if (!pluginOptions.enableIf || (pluginOptions.enableIf && (await pluginOptions.enableIf(args)))) {
-        setExecuteFn(jitExecute);
+        const compilationResult = getCompilationResult(args.schema, args.document, args.operationName ?? undefined);
+        if (isCompiledQuery(compilationResult)) {
+          setExecuteFn(function jitExecute(args): PromiseOrValue<ExecutionResult<any, any>> {
+            return compilationResult.query(args.rootValue, args.contextValue, args.variableValues);
+          });
+        } else if (compilationResult != null) {
+          setResultAndStopExecution(compilationResult as ExecutionResult);
+        }
       }
     },
-    async onSubscribe({ args, setSubscribeFn }) {
+    async onSubscribe({ args, setSubscribeFn, setResultAndStopExecution }) {
       if (!pluginOptions.enableIf || (pluginOptions.enableIf && (await pluginOptions.enableIf(args)))) {
-        setSubscribeFn(jitSubscribe);
+        const compilationResult = getCompilationResult(args.schema, args.document, args.operationName ?? undefined);
+        if (isCompiledQuery(compilationResult)) {
+          setSubscribeFn(function jitSubscribe(args: ExecutionArgs) {
+            return compilationResult.subscribe
+              ? compilationResult.subscribe(args.rootValue, args.contextValue, args.variableValues)
+              : compilationResult.query(args.rootValue, args.contextValue, args.variableValues);
+          } as any);
+        } else if (compilationResult != null) {
+          setResultAndStopExecution(compilationResult as ExecutionResult);
+        }
       }
     },
   };

packages/plugins/graphql-jit/test/graphql-jit.spec.ts

@@ -8,6 +8,7 @@ import { makeExecutableSchema } from '@graphql-tools/schema';
 import { execute, subscribe } from 'graphql';
 import { useGraphQlJit } from '../src';
 import lru from 'tiny-lru';
+import { GraphQLError } from 'graphql';
 
 describe('useGraphQlJit', () => {
   const schema = makeExecutableSchema({
@@ -163,4 +164,28 @@ describe('useGraphQlJit', () => {
     expect(cache.get).toHaveBeenCalled();
     expect(cache.set).toHaveBeenCalled();
   });
+
+  it('Should throw validation errors if compilation fails', async () => {
+    const plugin = useGraphQlJit(
+      {},
+      {
+        cache: {
+          get() {
+            return {
+              errors: [new GraphQLError('Some random error')],
+            };
+          },
+          set() {},
+        },
+      }
+    );
+    jest.spyOn(plugin, 'onValidate');
+    jest.spyOn(plugin, 'onExecute');
+    const testInstance = createTestkit([plugin], schema);
+
+    const result = await testInstance.execute(`query Foo { test }`);
+    expect(result['errors']).toBeTruthy();
+    expect(plugin.onValidate).toHaveBeenCalled();
+    expect(plugin.onExecute).not.toHaveBeenCalled();
+  });
 });