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

.eslintrc.json

@@ -12,7 +12,7 @@
     "prettier"
   ],
   "parserOptions": {
-    "ecmaVersion": 2015,
+    "ecmaVersion": 2017,
     "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,17 +1,37 @@
 const fp = require('fastify-plugin')
 const { als } = require('asynchronous-local-storage')
+const { AsyncResource } = require('async_hooks')
 
 const requestContext = {
   get: als.get,
   set: als.set,
 }
 
+/**
+ * Monkey patches `.emit()` method of the given emitter, so
+ * that all event listeners are run in scope of the provided
+ * async resource.
+ */
+const wrapEmitter = (emitter, asyncResource) => {
+  const original = emitter.emit
+  emitter.emit = function (type, ...args) {
+    return asyncResource.runInAsyncScope(original, emitter, type, ...args)
+  }
+}
+
+const wrapHttpEmitters = (req, res) => {
+  const asyncResource = new AsyncResource('fastify-request-context')
+  wrapEmitter(req, asyncResource)
+  wrapEmitter(res, asyncResource)
+}
+
 function plugin(fastify, opts, next) {
   fastify.decorate('requestContext', requestContext)
   fastify.decorateRequest('requestContext', requestContext)
 
-  fastify.addHook('onRequest', (req, res, done) => {
+  fastify.addHook(opts.hook || 'onRequest', (req, res, done) => {
     als.runWith(() => {
+      wrapHttpEmitters(req.raw, res.raw)
       done()
     }, opts.defaultStoreValues)
   })

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,82 @@
+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.route({
+    url: '/',
+    method: ['GET', 'POST'],
+    handler: endpoint,
+  })
+
+  return app
+}
+
+module.exports = {
+  initAppPostWithAllPlugins,
+  initAppPostWithPrevalidation,
+  initAppPost,
+  initAppGet,
+}

test/requestContextPlugin.e2e.spec.js

@@ -0,0 +1,190 @@
+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(5)
+    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')
+        })
+    })
+  })
+})