feat: added AsyncResource factory function option (#173)

Author: imjuni

README.md

@@ -58,15 +58,16 @@ fastify.register(fastifyRequestContextPlugin, {
 });
 ```
 
-This plugin accepts options `hook` and `defaultStoreValues`. 
+This plugin accepts options `hook` and `defaultStoreValues`, `createAsyncResource`.
 
 * `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` / `defaultStoreValues(req: FastifyRequest)` sets initial values for the store (that can be later overwritten during request execution if needed). Can be set to either an object or a function that returns an object. The function will be sent the request object for the new context. This is an optional parameter.
+* `createAsyncResource` can specify a factory function that creates an extended `AsyncResource` object.
 
 From there you can set a context in another hook, route, or method that is within scope.
 
 Request context (with methods `get` and `set`) is exposed by library itself, but is also available as decorator on `fastify.requestContext` app instance as well as on `req` request instance.
- 
+
 For instance:
 
 ```js
@@ -77,7 +78,8 @@ const app = fastify({ logger: true })
 app.register(fastifyRequestContextPlugin, { 
   defaultStoreValues: {
     user: { id: 'system' } 
-  }
+  },
+  createAsyncResource: (req, context) => new MyCustomAsyncResource('custom-resource-type', req.id, context.user.id)
 });
 
 app.addHook('onRequest', (req, reply, done) => {

index.js

@@ -34,7 +34,10 @@ function fastifyRequestContext(fastify, opts, next) {
       : opts.defaultStoreValues
 
     asyncLocalStorage.run({ ...defaultStoreValues }, () => {
-      const asyncResource = new AsyncResource('fastify-request-context')
+      const asyncResource =
+        opts.createAsyncResource != null
+          ? opts.createAsyncResource(req, requestContext)
+          : new AsyncResource('fastify-request-context')
       req[asyncResourceSymbol] = asyncResource
       asyncResource.runInAsyncScope(done, req.raw)
     })

test-tap/requestContextPlugin.e2e.test.js

@@ -1,13 +1,17 @@
 'use strict'
 
+const fastify = require('fastify')
 const request = require('superagent')
 const {
   initAppPostWithPrevalidation,
   initAppPostWithAllPlugins,
   initAppGetWithDefaultStoreValues,
 } = require('../test/internal/appInitializer')
+const { fastifyRequestContext } = require('..')
 const { TestService } = require('../test/internal/testService')
 const t = require('tap')
+const { CustomResource, AsyncHookContainer } = require('../test/internal/watcherService')
+const { executionAsyncId } = require('async_hooks')
 const test = t.test
 
 let app
@@ -332,3 +336,34 @@ test('does not throw when accessing context object outside of context', (t) => {
     })
   })
 })
+
+test('passing a custom resource factory function when create as AsyncResource', (t) => {
+  t.plan(2)
+
+  const container = new AsyncHookContainer(['fastify-request-context', 'custom-resource-type'])
+
+  app = fastify({ logger: true })
+  app.register(fastifyRequestContext, {
+    defaultStoreValues: { user: { id: 'system' } },
+    createAsyncResource: () => {
+      return new CustomResource('custom-resource-type', '1111-2222-3333')
+    },
+  })
+
+  const route = (req) => {
+    const store = container.getStore(executionAsyncId())
+    t.equal(store.traceId, '1111-2222-3333')
+    return Promise.resolve({ userId: req.requestContext.get('user').id })
+  }
+
+  app.get('/', route)
+
+  return app.listen({ port: 0, host: '127.0.0.1' }).then(() => {
+    const { address, port } = app.server.address()
+    const url = `${address}:${port}`
+
+    return request('GET', url).then((response1) => {
+      t.equal(response1.body.userId, 'system')
+    })
+  })
+})

test/internal/watcherService.js

@@ -0,0 +1,118 @@
+'use strict'
+
+const { executionAsyncId, createHook, AsyncResource } = require('node:async_hooks')
+const { EventEmitter } = require('node:events')
+
+class CustomResource extends AsyncResource {
+  constructor(type, traceId) {
+    super(type)
+
+    this.traceId = traceId
+  }
+}
+
+class AsyncWatcher extends EventEmitter {
+  setupInitHook() {
+    // init is called during object construction. The resource may not have
+    // completed construction when this callback runs, therefore all fields of the
+    // resource referenced by "asyncId" may not have been populated.
+    this.init = (asyncId, type, triggerAsyncId, resource) => {
+      this.emit('INIT', {
+        asyncId,
+        type,
+        triggerAsyncId,
+        executionAsyncId: executionAsyncId(),
+        resource,
+      })
+    }
+    return this
+  }
+
+  setupDestroyHook() {
+    // Destroy is called when an AsyncWrap instance is destroyed.
+    this.destroy = (asyncId) => {
+      this.emit('DESTROY', {
+        asyncId,
+        executionAsyncId: executionAsyncId(),
+      })
+    }
+    return this
+  }
+
+  start() {
+    createHook({
+      init: this.init.bind(this),
+      destroy: this.destroy.bind(this),
+    }).enable()
+
+    return this
+  }
+}
+
+class AsyncHookContainer {
+  constructor(types) {
+    const checkedTypes = types
+
+    const idMap = new Map()
+    const resourceMap = new Map()
+    const watcher = new AsyncWatcher()
+    const check = (t) => {
+      try {
+        return checkedTypes.includes(t)
+      } catch (err) {
+        return false
+      }
+    }
+
+    watcher
+      .setupInitHook()
+      .setupDestroyHook()
+      .start()
+      .on('INIT', ({ asyncId, type, resource, triggerAsyncId }) => {
+        idMap.set(asyncId, triggerAsyncId)
+
+        if (check(type)) {
+          resourceMap.set(asyncId, resource)
+        }
+      })
+      .on('DESTROY', ({ asyncId }) => {
+        idMap.delete(asyncId)
+        resourceMap.delete(asyncId)
+      })
+
+    this.types = checkedTypes
+    this.idMap = idMap
+    this.resourceMap = resourceMap
+    this.watcher = watcher
+  }
+
+  getStore(asyncId) {
+    let resource = this.resourceMap.get(asyncId)
+
+    if (resource != null) {
+      return resource
+    }
+
+    let id = this.idMap.get(asyncId)
+    let sentinel = 0
+
+    while (id != null && sentinel < 100) {
+      resource = this.resourceMap.get(id)
+
+      if (resource != null) {
+        return resource
+      }
+
+      id = this.idMap.get(id)
+      sentinel += 1
+    }
+
+    return undefined
+  }
+}
+
+module.exports = {
+  AsyncWatcher,
+  AsyncHookContainer,
+  CustomResource,
+}

types/index.d.ts

@@ -1,3 +1,4 @@
+import { AsyncResource } from 'async_hooks'
 import { FastifyPluginCallback, FastifyRequest } from 'fastify'
 
 type FastifyRequestContext =
@@ -23,6 +24,11 @@ declare namespace fastifyRequestContext {
     set<K extends keyof RequestContextData>(key: K, value: RequestContextData[K]): void
   }
 
+  export type CreateAsyncResourceFactory<T extends AsyncResource = AsyncResource> = (
+    req: FastifyRequest,
+    context: RequestContext,
+  ) => T
+
   export type RequestContextDataFactory = (req: FastifyRequest) => RequestContextData
 
   export type Hook =
@@ -43,6 +49,7 @@ declare namespace fastifyRequestContext {
   export interface FastifyRequestContextOptions {
     defaultStoreValues?: RequestContextData | RequestContextDataFactory
     hook?: Hook
+    createAsyncResource?: CreateAsyncResourceFactory
   }
 
   export const requestContext: RequestContext