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 (or multiple hooks) should request context initialization be bound. Make sure that you are not trying to access the request context earlier in the lifecycle than you initialize it. Default value is `preValidation`.
+* `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

@@ -5,8 +5,22 @@ export type RequestContext = {
   set: <T>(key: string, value: T) => void
 }
 
+export type Hook =
+  | 'preValidation'
+  | 'preHandler'
+  | 'preSerialization'
+  | 'onSend'
+  | 'onResponse'
+  | 'onTimeout'
+  | 'onError'
+  | 'onRoute'
+  | 'onRegister'
+  | 'onReady'
+  | 'onClose'
+
 export type RequestContextOptions = {
   defaultStoreValues?: Record<string, any>
+  hook?: 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

@@ -10,7 +10,7 @@ function plugin(fastify, opts, next) {
   fastify.decorate('requestContext', requestContext)
   fastify.decorateRequest('requestContext', requestContext)
 
-  fastify.addHook('onRequest', (req, res, done) => {
+  fastify.addHook(opts.hook || 'preValidation', (req, res, done) => {
     als.runWith(() => {
       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/requestContextPlugin.e2e.spec.js

@@ -0,0 +1,215 @@
+const fastify = require('fastify')
+const request = require('superagent')
+const { fastifyRequestContextPlugin } = require('../lib/requestContextPlugin')
+const { TestService } = require('./internal/testService')
+
+async 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,
+  })
+
+  await app.listen(0)
+  return app
+}
+
+async 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,
+  })
+
+  await app.listen(0)
+  return app
+}
+
+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}`)
+          }
+
+          initAppPostWithPrevalidation(route).then((_app) => {
+            app = _app
+            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}`)
+          }
+
+          initAppPostWithAllPlugins(route, 'preValidation').then((_app) => {
+            app = _app
+            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
+    })
+  })
+})

test/requestContextPlugin.spec.js

@@ -20,6 +20,26 @@ function initAppPost(endpoint) {
   return app.ready()
 }
 
+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.ready()
+}
+
 describe('requestContextPlugin', () => {
   let app
   afterEach(() => {
@@ -171,4 +191,76 @@ describe('requestContextPlugin', () => {
       return promiseRequest2
     })
   })
+
+  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, reply) => {
+            const requestId = req.requestContext.get('testKey')
+
+            function prepareReply() {
+              testService.processRequest(requestId.replace('testValue', '')).then(() => {
+                const storedValue = req.requestContext.get('testKey')
+                reply.status(204).send({
+                  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)
+            }
+          }
+
+          initAppPostWithPrevalidation(route).then((_app) => {
+            app = _app
+            testService = new TestService(app)
+            const response1Promise = app
+              .inject()
+              .post('/')
+              .body({ requestId: 1 })
+              .end()
+              .then((response) => {
+                expect(response.json().storedValue).toBe('testValue1')
+                responseCounter++
+                if (responseCounter === 2) {
+                  resolveResponsePromise()
+                }
+              })
+
+            const response2Promise = app
+              .inject()
+              .post('/')
+              .body({ requestId: 2 })
+              .end()
+              .then((response) => {
+                expect(response.json().storedValue).toBe('testValue2')
+                responseCounter++
+                if (responseCounter === 2) {
+                  resolveResponsePromise()
+                }
+              })
+
+            return Promise.all([response1Promise, response2Promise])
+          })
+        })
+
+        return promiseRequest1
+      })
+
+      return promiseRequest2
+    })
+  })
 })