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,7 @@
 const fp = require('fastify-plugin')
 const { als } = require('asynchronous-local-storage')
+const { AsyncResource } = require('async_hooks')
+const asyncResourceSymbol = Symbol('asyncResource')
 
 const requestContext = {
   get: als.get,
@@ -9,12 +11,28 @@ const requestContext = {
 function plugin(fastify, opts, next) {
   fastify.decorate('requestContext', requestContext)
   fastify.decorateRequest('requestContext', requestContext)
+  fastify.decorateRequest(asyncResourceSymbol, null)
+  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[asyncResourceSymbol] = 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[asyncResourceSymbol]
+      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,
+}