Fix context loss when async storage is instantiated early in the request lifecycle

.eslintrc.json

@@ -12,7 +12,7 @@
     "prettier"
   ],
   "parserOptions": {
-    "ecmaVersion": 2015,
+    "ecmaVersion": 2018,
     "sourceType": "module"
   },
   "rules": {

README.md

@@ -27,15 +27,17 @@ const { fastifyRequestContextPlugin } = require('fastify-request-context')
 const fastify = require('fastify');
 
 fastify.register(fastifyRequestContextPlugin, { 
+  hook: 'preValidation',
   defaultStoreValues: {
     user: { id: 'system' } 
   }
 });
 ``` 
 
-This plugin accepts option named `defaultStoreValues`. 
+This plugin accepts options `hook` and `defaultStoreValues`. 
 
-`defaultStoreValues` set initial values for the store (that can be later overwritten during request execution if needed). This is an optional parameter.
+* `hook` allows you to specify to which lifecycle hook should request context initialization be bound. Note that you need to initialize it on the earliest lifecycle stage that you intend to use it in, or earlier. Default value is `onRequest`.
+* `defaultStoreValues` sets initial values for the store (that can be later overwritten during request execution if needed). This is an optional parameter.
 
 From there you can set a context in another hook, route, or method that is within scope.
 

index.d.ts

@@ -1,12 +1,28 @@
-import Fastify, { FastifyPlugin, FastifyRequest } from 'fastify'
+import { FastifyPlugin, FastifyRequest } from 'fastify'
 
 export type RequestContext = {
   get: <T>(key: string) => T | undefined
   set: <T>(key: string, value: T) => void
 }
 
+export type Hook =
+  | 'onRequest'
+  | 'preParsing'
+  | 'preValidation'
+  | 'preHandler'
+  | 'preSerialization'
+  | 'onSend'
+  | 'onResponse'
+  | 'onTimeout'
+  | 'onError'
+  | 'onRoute'
+  | 'onRegister'
+  | 'onReady'
+  | 'onClose'
+
 export type RequestContextOptions = {
   defaultStoreValues?: Record<string, any>
+  hook?: Hook
 }
 
 declare module 'fastify' {

index.test-d.ts

@@ -13,6 +13,10 @@ expectAssignable<RequestContextOptions>({})
 expectAssignable<RequestContextOptions>({
   defaultStoreValues: { a: 'dummy' },
 })
+expectAssignable<RequestContextOptions>({
+  hook: 'preValidation',
+  defaultStoreValues: { a: 'dummy' },
+})
 
 expectType<RequestContext>(app.requestContext)
 

lib/requestContextPlugin.js

@@ -1,5 +1,6 @@
 const fp = require('fastify-plugin')
 const { als } = require('asynchronous-local-storage')
+const { AsyncResource } = require('async_hooks')
 
 const requestContext = {
   get: als.get,
@@ -9,12 +10,27 @@ const requestContext = {
 function plugin(fastify, opts, next) {
   fastify.decorate('requestContext', requestContext)
   fastify.decorateRequest('requestContext', requestContext)
+  const hook = opts.hook || 'onRequest'
 
-  fastify.addHook('onRequest', (req, res, done) => {
+  fastify.addHook(hook, (req, res, done) => {
     als.runWith(() => {
-      done()
+      const asyncResource = new AsyncResource('fastify-request-context')
+      req.asyncResource = asyncResource
+      asyncResource.runInAsyncScope(done, req.raw)
     }, opts.defaultStoreValues)
   })
+
+  // Both of onRequest and preParsing are executed after the als.runWith call within the "proper" async context (AsyncResource implicitly created by ALS).
+  // However, preValidation, preHandler and the route handler are executed as a part of req.emit('end') call which happens
+  // in a different async context, as req/res may emit events in a different context.
+  // Related to https://github.com/nodejs/node/issues/34430 and https://github.com/nodejs/node/issues/33723
+  if (hook === 'onRequest' || hook === 'preParsing') {
+    fastify.addHook('preValidation', (req, res, done) => {
+      const asyncResource = req.asyncResource
+      asyncResource.runInAsyncScope(done, req.raw)
+    })
+  }
+
   next()
 }
 

package.json

@@ -33,6 +33,7 @@
     "eslint-plugin-prettier": "^3.1.4",
     "jest": "^26.4.0",
     "prettier": "^2.0.5",
+    "superagent": "^6.0.0",
     "tsd": "^0.13.1",
     "typescript": "3.9.7"
   },

test/internal/appInitializer.js

@@ -0,0 +1,91 @@
+const fastify = require('fastify')
+const { fastifyRequestContextPlugin } = require('../../lib/requestContextPlugin')
+
+function initAppGet(endpoint) {
+  const app = fastify({ logger: true })
+  app.register(fastifyRequestContextPlugin)
+
+  app.get('/', endpoint)
+  return app
+}
+
+function initAppPost(endpoint) {
+  const app = fastify({ logger: true })
+  app.register(fastifyRequestContextPlugin)
+
+  app.post('/', endpoint)
+
+  return app
+}
+
+function initAppPostWithPrevalidation(endpoint) {
+  const app = fastify({ logger: true })
+  app.register(fastifyRequestContextPlugin, { hook: 'preValidation' })
+
+  const preValidationFn = (req, reply, done) => {
+    const requestId = Number.parseInt(req.body.requestId)
+    req.requestContext.set('testKey', `testValue${requestId}`)
+    done()
+  }
+
+  app.route({
+    url: '/',
+    method: ['GET', 'POST'],
+    preValidation: preValidationFn,
+    handler: endpoint,
+  })
+
+  return app
+}
+
+function initAppPostWithAllPlugins(endpoint, requestHook) {
+  const app = fastify({ logger: true })
+  app.register(fastifyRequestContextPlugin, { hook: requestHook })
+
+  app.addHook('onRequest', (req, reply, done) => {
+    req.requestContext.set('onRequest', 'dummy')
+    done()
+  })
+
+  app.addHook('preParsing', (req, reply, payload, done) => {
+    req.requestContext.set('preParsing', 'dummy')
+    done(null, payload)
+  })
+
+  app.addHook('preValidation', (req, reply, done) => {
+    const requestId = Number.parseInt(req.body.requestId)
+    req.requestContext.set('preValidation', requestId)
+    req.requestContext.set('testKey', `testValue${requestId}`)
+    done()
+  })
+
+  app.addHook('preHandler', (req, reply, done) => {
+    const requestId = Number.parseInt(req.body.requestId)
+    req.requestContext.set('preHandler', requestId)
+    done()
+  })
+
+  app.addHook('preSerialization', (req, reply, payload, done) => {
+    const onRequestValue = req.requestContext.get('onRequest')
+    const preValidationValue = req.requestContext.get('preValidation')
+    done(null, {
+      ...payload,
+      preSerialization1: onRequestValue,
+      preSerialization2: preValidationValue,
+    })
+  })
+  app.route({
+    url: '/',
+    method: ['GET', 'POST'],
+    handler: endpoint,
+  })
+
+  return app
+}
+
+module.exports = {
+  initAppPostWithAllPlugins,
+  initAppPostWithPrevalidation,
+  initAppPost,
+  initAppGet,
+}

test/requestContextPlugin.e2e.spec.js

@@ -0,0 +1,192 @@
+const request = require('superagent')
+const {
+  initAppPostWithPrevalidation,
+  initAppPostWithAllPlugins,
+} = require('./internal/appInitializer')
+const { TestService } = require('./internal/testService')
+
+describe('requestContextPlugin E2E', () => {
+  let app
+  afterEach(() => {
+    return app.close()
+  })
+
+  it('correctly preserves values set in prevalidation phase within single POST request', () => {
+    expect.assertions(2)
+
+    let testService
+    let responseCounter = 0
+    return new Promise((resolveResponsePromise) => {
+      const promiseRequest2 = new Promise((resolveRequest2Promise) => {
+        const promiseRequest1 = new Promise((resolveRequest1Promise) => {
+          const route = (req) => {
+            const requestId = req.requestContext.get('testKey')
+
+            function prepareReply() {
+              return testService.processRequest(requestId.replace('testValue', '')).then(() => {
+                const storedValue = req.requestContext.get('testKey')
+                return Promise.resolve({ storedValue })
+              })
+            }
+
+            // We don't want to read values until both requests wrote their values to see if there is a racing condition
+            if (requestId === 'testValue1') {
+              resolveRequest1Promise()
+              return promiseRequest2.then(prepareReply)
+            }
+
+            if (requestId === 'testValue2') {
+              resolveRequest2Promise()
+              return promiseRequest1.then(prepareReply)
+            }
+
+            throw new Error(`Unexpected requestId: ${requestId}`)
+          }
+
+          app = initAppPostWithPrevalidation(route)
+          app.listen(0).then(() => {
+            testService = new TestService(app)
+            const { address, port } = app.server.address()
+            const url = `${address}:${port}`
+            const response1Promise = request('POST', url)
+              .send({ requestId: 1 })
+              .then((response) => {
+                expect(response.body.storedValue).toBe('testValue1')
+                responseCounter++
+                if (responseCounter === 2) {
+                  resolveResponsePromise()
+                }
+              })
+
+            const response2Promise = request('POST', url)
+              .send({ requestId: 2 })
+              .then((response) => {
+                expect(response.body.storedValue).toBe('testValue2')
+                responseCounter++
+                if (responseCounter === 2) {
+                  resolveResponsePromise()
+                }
+              })
+
+            return Promise.all([response1Promise, response2Promise])
+          })
+        })
+
+        return promiseRequest1
+      })
+
+      return promiseRequest2
+    })
+  })
+
+  it('correctly preserves values set in multiple phases within single POST request', () => {
+    expect.assertions(10)
+
+    let testService
+    let responseCounter = 0
+    return new Promise((resolveResponsePromise) => {
+      const promiseRequest2 = new Promise((resolveRequest2Promise) => {
+        const promiseRequest1 = new Promise((resolveRequest1Promise) => {
+          const route = (req) => {
+            const onRequestValue = req.requestContext.get('onRequest')
+            const preParsingValue = req.requestContext.get('preParsing')
+            const preValidationValue = req.requestContext.get('preValidation')
+            const preHandlerValue = req.requestContext.get('preHandler')
+
+            expect(onRequestValue).toBe(undefined)
+            expect(preParsingValue).toBe(undefined)
+            expect(preValidationValue).toEqual(expect.any(Number))
+            expect(preHandlerValue).toEqual(expect.any(Number))
+
+            const requestId = `testValue${preHandlerValue}`
+
+            function prepareReply() {
+              return testService.processRequest(requestId.replace('testValue', '')).then(() => {
+                const storedValue = req.requestContext.get('preValidation')
+                return Promise.resolve({ storedValue: `testValue${storedValue}` })
+              })
+            }
+
+            // We don't want to read values until both requests wrote their values to see if there is a racing condition
+            if (requestId === 'testValue1') {
+              resolveRequest1Promise()
+              return promiseRequest2.then(prepareReply)
+            }
+
+            if (requestId === 'testValue2') {
+              resolveRequest2Promise()
+              return promiseRequest1.then(prepareReply)
+            }
+
+            throw new Error(`Unexpected requestId: ${requestId}`)
+          }
+
+          app = initAppPostWithAllPlugins(route, 'preValidation')
+
+          app.listen(0).then(() => {
+            testService = new TestService(app)
+            const { address, port } = app.server.address()
+            const url = `${address}:${port}`
+            const response1Promise = request('POST', url)
+              .send({ requestId: 1 })
+              .then((response) => {
+                expect(response.body.storedValue).toBe('testValue1')
+                responseCounter++
+                if (responseCounter === 2) {
+                  resolveResponsePromise()
+                }
+              })
+
+            const response2Promise = request('POST', url)
+              .send({ requestId: 2 })
+              .then((response) => {
+                expect(response.body.storedValue).toBe('testValue2')
+                responseCounter++
+                if (responseCounter === 2) {
+                  resolveResponsePromise()
+                }
+              })
+
+            return Promise.all([response1Promise, response2Promise])
+          })
+        })
+
+        return promiseRequest1
+      })
+
+      return promiseRequest2
+    })
+  })
+
+  it('does not lose request context after body parsing', () => {
+    expect.assertions(7)
+    const route = (req) => {
+      const onRequestValue = req.requestContext.get('onRequest')
+      const preParsingValue = req.requestContext.get('preParsing')
+      const preValidationValue = req.requestContext.get('preValidation')
+      const preHandlerValue = req.requestContext.get('preHandler')
+
+      expect(onRequestValue).toBe('dummy')
+      expect(preParsingValue).toBe('dummy')
+      expect(preValidationValue).toEqual(expect.any(Number))
+      expect(preHandlerValue).toEqual(expect.any(Number))
+
+      const requestId = `testValue${preHandlerValue}`
+      return Promise.resolve({ storedValue: requestId })
+    }
+
+    app = initAppPostWithAllPlugins(route, 'onRequest')
+
+    return app.listen(0).then(() => {
+      const { address, port } = app.server.address()
+      const url = `${address}:${port}`
+      return request('POST', url)
+        .send({ requestId: 1 })
+        .then((response) => {
+          expect(response.body.storedValue).toBe('testValue1')
+          expect(response.body.preSerialization1).toBe('dummy')
+          expect(response.body.preSerialization2).toBe(1)
+        })
+    })
+  })
+})