diff --git a/README.md b/README.md index 8252b4d..8c8c0ef 100644 --- a/README.md +++ b/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) => { diff --git a/index.js b/index.js index ba6be91..eeb6485 100644 --- a/index.js +++ b/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) }) diff --git a/test-tap/requestContextPlugin.e2e.test.js b/test-tap/requestContextPlugin.e2e.test.js index 9853715..5da3ce2 100644 --- a/test-tap/requestContextPlugin.e2e.test.js +++ b/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') + }) + }) +}) diff --git a/test/internal/watcherService.js b/test/internal/watcherService.js new file mode 100644 index 0000000..b66f0ff --- /dev/null +++ b/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, +} diff --git a/types/index.d.ts b/types/index.d.ts index ecd9cce..d25c399 100644 --- a/types/index.d.ts +++ b/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(key: K, value: RequestContextData[K]): void } + export type CreateAsyncResourceFactory = ( + 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