Add per request global context (#1789)

* Allow TS errors during development.

* Add branches for local (per-request) async storage.

* Use async storage.

* Remove extra space.

Co-authored-by: Tobbe Lundberg <tobbe@tlundberg.com>

* Write tests.

* Make tests slightly more robust.

* Add a helpful note about fixing the tests.

* Fix imports.

* On Fridays the Boolean logic is above my paygrade.

* Fix comments.

* Thanks @Tobbe

Co-authored-by: Tobbe Lundberg <tobbe@tlundberg.com>

Co-authored-by: Tobbe Lundberg <tobbe@tlundberg.com>

package.json

@@ -3,7 +3,7 @@
     "build": "lerna run build:js && ttsc --build --verbose",
     "build:types": "ttsc --build --verbose",
     "build:clean": "rimraf ./packages/**/dist",
-    "build:watch": "ttsc --build && lerna run build:watch --parallel",
+    "build:watch": "lerna run build:watch --parallel; ttsc --build",
     "test": "lerna run test --stream -- --colors --maxWorkers=4",
     "lint": "eslint -c .eslintrc.js packages",
     "lint:fix": "eslint -c .eslintrc.js --fix packages"

packages/api/src/functions/graphql.test.ts

@@ -1,13 +1,23 @@
 /* eslint-disable @typescript-eslint/ban-ts-comment */
-import { context } from '../globalContext'
+import { context, initPerRequestContext } from '../globalContext'
 
 import { createContextHandler } from './graphql'
 
-describe('graphql createContextHandler', () => {
-  it('merges the context correctly', async () => {
+describe('global context handlers', () => {
+  beforeAll(() => {
+    process.env.SAFE_GLOBAL_CONTEXT = '1'
+  })
+
+  afterAll(() => {
+    process.env.SAFE_GLOBAL_CONTEXT = '0'
+  })
+
+  it('merges the apollo resolver and global context correctly', async () => {
     const handler = createContextHandler({ a: 1 })
     // @ts-ignore
-    expect(await handler({ context: { b: 2 } })).toEqual({
+
+    const inlineContext = await handler({ context: { b: 2 } })
+    expect(inlineContext).toEqual({
       a: 1,
       b: 2,
       callbackWaitsForEmptyEventLoop: false,
@@ -68,3 +78,56 @@ describe('graphql createContextHandler', () => {
     })
   })
 })
+
+describe('per request context handlers', () => {
+  it('merges the apollo resolver and global context correctly', async () => {
+    const localAsyncStorage = initPerRequestContext()
+
+    localAsyncStorage.run(new Map(), async () => {
+      const handler = createContextHandler({ a: 1 })
+      // @ts-ignore
+      const inlineContext = await handler({ context: { b: 2 } })
+      expect(inlineContext).toEqual({
+        a: 1,
+        b: 2,
+        callbackWaitsForEmptyEventLoop: false,
+      })
+
+      expect(context).toEqual({
+        a: 1,
+        b: 2,
+        callbackWaitsForEmptyEventLoop: false,
+      })
+    })
+  })
+
+  it('maintains separate contexts for each request', (done) => {
+    const localAsyncStorage = initPerRequestContext()
+
+    // request 1 and request 2...
+    // request 1 is slow...
+    // request 2 is fast!
+    // they should have different contexts.
+
+    let request2Complete = false
+    localAsyncStorage.run(new Map(), async () => {
+      const handler = createContextHandler({ request: 1 })
+      // @ts-ignore
+      await handler({ context: { favoriteFood: 'cheese' } })
+      setTimeout(() => {
+        expect(context).toMatchObject({ request: 1, favoriteFood: 'cheese' })
+        // If this isn't true, then we might need to increase the timeouts
+        expect(request2Complete).toBeTruthy()
+        done()
+      }, 1)
+    })
+
+    localAsyncStorage.run(new Map(), async () => {
+      const handler = createContextHandler({ request: 2 })
+      // @ts-ignore
+      await handler({ context: { favoriteFood: 'cake' } })
+      request2Complete = true
+      expect(context).toMatchObject({ request: 2, favoriteFood: 'cake' })
+    })
+  })
+})

packages/api/src/functions/graphql.ts

@@ -5,8 +5,12 @@ import type { APIGatewayProxyEvent, Context as LambdaContext } from 'aws-lambda'
 
 import type { AuthContextPayload } from 'src/auth'
 import { getAuthenticationContext } from 'src/auth'
-import type { GlobalContext } from 'src/globalContext'
-import { setContext } from 'src/globalContext'
+import {
+  GlobalContext,
+  setContext,
+  initPerRequestContext,
+  usePerRequestContext,
+} from 'src/globalContext'
 
 export type GetCurrentUser = (
   decoded: AuthContextPayload[0],
@@ -52,7 +56,6 @@ export const createContextHandler = (
       // if userContext is a function, run that and return just the result
       customUserContext = await userContext({ event, context })
     }
-
     // Sets the **global** context object, which can be imported with:
     // import { context } from '@redwoodjs/api'
     return setContext({
@@ -127,11 +130,26 @@ export const createGraphQLHandler = ({
     context: LambdaContext,
     callback: any
   ): void => {
-    try {
-      handler(event, context, callback)
-    } catch (e) {
-      onException && onException()
-      throw e
+    if (usePerRequestContext()) {
+      // This must be used when you're self-hosting RedwoodJS.
+      const localAsyncStorage = initPerRequestContext()
+      localAsyncStorage.run(new Map(), () => {
+        try {
+          handler(event, context, callback)
+        } catch (e) {
+          onException && onException()
+          throw e
+        }
+      })
+    } else {
+      // This is OK for AWS (Netlify/Vercel) because each Lambda request
+      // is handled individually.
+      try {
+        handler(event, context, callback)
+      } catch (e) {
+        onException && onException()
+        throw e
+      }
     }
   }
 }

packages/api/src/globalContext.ts

@@ -1,5 +1,49 @@
+/* eslint-disable react-hooks/rules-of-hooks */
+
+// AWS Lambda run each request in a new process,
+// a process is not reused until a request is completed.
+//
+// Which means that each `global.context` is scoped to the lifetime of each request.
+// This makes it safe to use the global context for Redwood Functions.
+
+// However when not in AWS Lambda, NodeJS is single-threaded, you must use the
+// per-request global context, otherwise you risk a race-condition
+// where one request overwrites another's global context.
+//
+// Alternatively only use the local `context` in a graphql resolver.
+
+import { AsyncLocalStorage } from 'async_hooks'
+
 export interface GlobalContext {
-  [key: string]: any
+  [key: string]: unknown
+}
+
+let GLOBAL_CONTEXT: GlobalContext = {}
+let PER_REQUEST_CONTEXT:
+  | undefined
+  | AsyncLocalStorage<Map<string, GlobalContext>> = undefined
+
+export const usePerRequestContext = () =>
+  process.env.SAFE_GLOBAL_CONTEXT !== '1'
+
+export const initPerRequestContext = () => {
+  GLOBAL_CONTEXT = {}
+  PER_REQUEST_CONTEXT = new AsyncLocalStorage()
+  return PER_REQUEST_CONTEXT
+}
+
+export const createContextProxy = () => {
+  return new Proxy<GlobalContext>(GLOBAL_CONTEXT, {
+    get: (_target, property: string) => {
+      const store = PER_REQUEST_CONTEXT?.getStore()
+      if (!store) {
+        throw new Error(
+          'Async local storage is not initialized. Call `initGlobalContext` before attempting to read from the store.'
+        )
+      }
+      return store.get('context')?.[property]
+    },
+  })
 }
 
 export let context: GlobalContext = {}
@@ -8,6 +52,21 @@ export let context: GlobalContext = {}
  * Replace the existing global context.
  */
 export const setContext = (newContext: GlobalContext): GlobalContext => {
-  context = newContext
+  GLOBAL_CONTEXT = newContext
+
+  if (usePerRequestContext()) {
+    // re-init the proxy, so that calls to `console.log(context)` is the full object
+    // not the one initialized earlier.
+    context = createContextProxy()
+    const store = PER_REQUEST_CONTEXT?.getStore()
+    if (!store) {
+      throw new Error(
+        'Per request context is not initialized, please use `initPerRequestContext`'
+      )
+    }
+    store.set('context', GLOBAL_CONTEXT)
+  } else {
+    context = GLOBAL_CONTEXT
+  }
   return context
 }