Bind context to async execution avoiding race-conditions

.changeset/shiny-turkeys-dream.md

@@ -0,0 +1,5 @@
+---
+'graphql-modules': patch
+---
+
+Bind context to async execution avoiding race-conditions

packages/graphql-modules/src/application/apollo.ts

@@ -2,7 +2,7 @@ import { wrapSchema } from '@graphql-tools/wrap';
 import { DocumentNode, execute, GraphQLSchema } from 'graphql';
 import { uniqueId } from '../shared/utils';
 import { InternalAppContext } from './application';
-import { ExecutionContextBuilder } from './context';
+import { ExecutionContextBuilder, ExecutionContextEnv } from './context';
 import { Application } from './types';
 
 const CONTEXT_ID = Symbol.for('context-id');
@@ -60,11 +60,12 @@ export function apolloSchemaCreator({
     > = {};
     const subscription = createSubscription();
 
-    function getSession(ctx: any) {
+    function getSession(
+      ctx: any,
+      { context, ɵdestroy: destroy }: ExecutionContextEnv
+    ) {
       if (!ctx[CONTEXT_ID]) {
         ctx[CONTEXT_ID] = uniqueId((id) => !sessions[id]);
-        const { context, ɵdestroy: destroy } = contextBuilder(ctx);
-
         sessions[ctx[CONTEXT_ID]] = {
           count: 0,
           session: {
@@ -99,24 +100,27 @@ export function apolloSchemaCreator({
             operationName: input.operationName,
           });
         }
-        // Create an execution context
-        const { context, destroy } = getSession(input.context!);
 
-        // It's important to wrap the executeFn within a promise
-        // so we can easily control the end of execution (with finally)
-        return Promise.resolve()
-          .then(
-            () =>
-              execute({
-                schema,
-                document: input.document,
-                contextValue: context,
-                variableValues: input.variables as any,
-                rootValue: input.rootValue,
-                operationName: input.operationName,
-              }) as any
-          )
-          .finally(destroy);
+        // Create an execution context and run within it
+        return contextBuilder(input.context!).runWithContext((env) => {
+          const { context, destroy } = getSession(input.context!, env);
+
+          // It's important to wrap the executeFn within a promise
+          // so we can easily control the end of execution (with finally)
+          return Promise.resolve()
+            .then(
+              () =>
+                execute({
+                  schema,
+                  document: input.document,
+                  contextValue: context,
+                  variableValues: input.variables as any,
+                  rootValue: input.rootValue,
+                  operationName: input.operationName,
+                }) as any
+            )
+            .finally(destroy);
+        });
       },
     });
   };

packages/graphql-modules/src/application/context.ts

@@ -1,3 +1,4 @@
+import { AsyncLocalStorage } from 'async_hooks';
 import { Injector, ReflectiveInjector } from '../di';
 import { ResolvedProvider } from '../di/resolution';
 import { ID } from '../shared/types';
@@ -6,11 +7,22 @@ import type { InternalAppContext, ModulesMap } from './application';
 import { attachGlobalProvidersMap } from './di';
 import { CONTEXT } from './tokens';
 
+const alc = new AsyncLocalStorage<{
+  getApplicationContext(): GraphQLModules.AppContext;
+  getModuleContext(moduleId: string): GraphQLModules.ModuleContext;
+}>();
+
 export type ExecutionContextBuilder<
   TContext extends {
     [key: string]: any;
   } = {},
-> = (context: TContext) => {
+> = (context: TContext) => ExecutionContextEnv & {
+  runWithContext<TReturn = any>(
+    cb: (env: ExecutionContextEnv) => TReturn
+  ): TReturn;
+};
+
+export type ExecutionContextEnv = {
   context: InternalAppContext;
   ɵdestroy(): void;
   ɵinjector: Injector;
@@ -67,12 +79,15 @@ export function createContextBuilder({
     });
 
     appInjector.setExecutionContextGetter(function executionContextGetter() {
-      return appContext;
+      return alc.getStore()?.getApplicationContext() || appContext;
     } as any);
 
     function createModuleExecutionContextGetter(moduleId: string) {
       return function moduleExecutionContextGetter() {
-        return getModuleContext(moduleId, context);
+        return (
+          alc.getStore()?.getModuleContext(moduleId) ||
+          getModuleContext(moduleId, context)
+        );
       };
     }
 
@@ -164,7 +179,7 @@ export function createContextBuilder({
       },
     });
 
-    return {
+    const env: ExecutionContextEnv = {
       ɵdestroy: once(() => {
         providersToDestroy.forEach(([injector, keyId]) => {
           // If provider was instantiated
@@ -178,6 +193,23 @@ export function createContextBuilder({
       ɵinjector: operationAppInjector,
       context: sharedContext,
     };
+
+    return {
+      ...env,
+      runWithContext(cb) {
+        return alc.run(
+          {
+            getApplicationContext() {
+              return appContext;
+            },
+            getModuleContext(moduleId) {
+              return getModuleContext(moduleId, context);
+            },
+          },
+          () => cb(env)
+        );
+      },
+    };
   };
 
   return contextBuilder;

packages/graphql-modules/src/application/execution.ts

@@ -10,6 +10,7 @@ import { Application } from './types';
 import { ExecutionContextBuilder } from './context';
 import { Maybe } from '../shared/types';
 import { isNotSchema } from '../shared/utils';
+import { InternalAppContext } from './application';
 
 export function executionCreator({
   contextBuilder,
@@ -30,38 +31,47 @@ export function executionCreator({
       fieldResolver?: Maybe<GraphQLFieldResolver<any, any>>,
       typeResolver?: Maybe<GraphQLTypeResolver<any, any>>
     ) => {
-      // Create an execution context
-      const { context, ɵdestroy: destroy } =
-        options?.controller ??
-        contextBuilder(
-          isNotSchema<ExecutionArgs>(argsOrSchema)
-            ? argsOrSchema.contextValue
-            : contextValue
-        );
+      function perform({
+        context,
+        ɵdestroy: destroy,
+      }: {
+        context: InternalAppContext;
+        ɵdestroy: () => void;
+      }) {
+        const executionArgs: ExecutionArgs = isNotSchema<ExecutionArgs>(
+          argsOrSchema
+        )
+          ? {
+              ...argsOrSchema,
+              contextValue: context,
+            }
+          : {
+              schema: argsOrSchema,
+              document: document!,
+              rootValue,
+              contextValue: context,
+              variableValues,
+              operationName,
+              fieldResolver,
+              typeResolver,
+            };
 
-      const executionArgs: ExecutionArgs = isNotSchema<ExecutionArgs>(
-        argsOrSchema
-      )
-        ? {
-            ...argsOrSchema,
-            contextValue: context,
-          }
-        : {
-            schema: argsOrSchema,
-            document: document!,
-            rootValue,
-            contextValue: context,
-            variableValues,
-            operationName,
-            fieldResolver,
-            typeResolver,
-          };
+        // It's important to wrap the executeFn within a promise
+        // so we can easily control the end of execution (with finally)
+        return Promise.resolve()
+          .then(() => executeFn(executionArgs))
+          .finally(destroy);
+      }
 
-      // It's important to wrap the executeFn within a promise
-      // so we can easily control the end of execution (with finally)
-      return Promise.resolve()
-        .then(() => executeFn(executionArgs))
-        .finally(destroy);
+      if (options?.controller) {
+        return perform(options.controller);
+      }
+
+      return contextBuilder(
+        isNotSchema<ExecutionArgs>(argsOrSchema)
+          ? argsOrSchema.contextValue
+          : contextValue
+      ).runWithContext(perform);
     };
   };
 

packages/graphql-modules/src/application/subscription.ts

@@ -13,6 +13,7 @@ import {
 } from '../shared/utils';
 import { ExecutionContextBuilder } from './context';
 import { Application } from './types';
+import { InternalAppContext } from './application';
 
 export function subscriptionCreator({
   contextBuilder,
@@ -33,51 +34,59 @@ export function subscriptionCreator({
       fieldResolver?: Maybe<GraphQLFieldResolver<any, any>>,
       subscribeFieldResolver?: Maybe<GraphQLFieldResolver<any, any>>
     ) => {
-      // Create an subscription context
-      const { context, ɵdestroy: destroy } =
-        options?.controller ??
-        contextBuilder(
+      function perform({
+        context,
+        ɵdestroy: destroy,
+      }: {
+        context: InternalAppContext;
+        ɵdestroy: () => void;
+      }) {
+        const subscriptionArgs: SubscriptionArgs =
           isNotSchema<SubscriptionArgs>(argsOrSchema)
-            ? argsOrSchema.contextValue
-            : contextValue
-        );
+            ? {
+                ...argsOrSchema,
+                contextValue: context,
+              }
+            : {
+                schema: argsOrSchema,
+                document: document!,
+                rootValue,
+                contextValue: context,
+                variableValues,
+                operationName,
+                fieldResolver,
+                subscribeFieldResolver,
+              };
 
-      const subscriptionArgs: SubscriptionArgs = isNotSchema<SubscriptionArgs>(
-        argsOrSchema
-      )
-        ? {
-            ...argsOrSchema,
-            contextValue: context,
-          }
-        : {
-            schema: argsOrSchema,
-            document: document!,
-            rootValue,
-            contextValue: context,
-            variableValues,
-            operationName,
-            fieldResolver,
-            subscribeFieldResolver,
-          };
+        let isIterable = false;
 
-      let isIterable = false;
+        // It's important to wrap the subscribeFn within a promise
+        // so we can easily control the end of subscription (with finally)
+        return Promise.resolve()
+          .then(() => subscribeFn(subscriptionArgs))
+          .then((sub) => {
+            if (isAsyncIterable(sub)) {
+              isIterable = true;
+              return tapAsyncIterator(sub, destroy);
+            }
+            return sub;
+          })
+          .finally(() => {
+            if (!isIterable) {
+              destroy();
+            }
+          });
+      }
 
-      // It's important to wrap the subscribeFn within a promise
-      // so we can easily control the end of subscription (with finally)
-      return Promise.resolve()
-        .then(() => subscribeFn(subscriptionArgs))
-        .then((sub) => {
-          if (isAsyncIterable(sub)) {
-            isIterable = true;
-            return tapAsyncIterator(sub, destroy);
-          }
-          return sub;
-        })
-        .finally(() => {
-          if (!isIterable) {
-            destroy();
-          }
-        });
+      if (options?.controller) {
+        return perform(options.controller);
+      }
+
+      return contextBuilder(
+        isNotSchema<SubscriptionArgs>(argsOrSchema)
+          ? argsOrSchema.contextValue
+          : contextValue
+      ).runWithContext(perform);
     };
   };