feat: add support for otel context

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@athenna/http",
-  "version": "5.58.0",
+  "version": "5.59.0",
   "description": "The Athenna Http server. Built on top of fastify.",
   "license": "MIT",
   "author": "João Lenon <lenon@athenna.io>",
@@ -92,6 +92,8 @@
     "@fastify/static": "^8.3.0",
     "@fastify/swagger": "^9.7.0",
     "@fastify/swagger-ui": "^5.2.5",
+    "@opentelemetry/api": "^1.9.1",
+    "@opentelemetry/context-async-hooks": "^2.7.0",
     "@typescript-eslint/eslint-plugin": "^8.57.0",
     "@typescript-eslint/parser": "^8.57.0",
     "autocannon": "^7.15.0",

src/handlers/FastifyHandler.ts

@@ -7,33 +7,29 @@
  * file that was distributed with this source code.
  */
 
-import { Is } from '@athenna/common'
+import { Config } from '@athenna/config'
+import { Is, Module } from '@athenna/common'
 import { Request } from '#src/context/Request'
 import { Response } from '#src/context/Response'
+import type { Context as OtelContext } from '@opentelemetry/api'
 import type { RequestHandler } from '#src/types/contexts/Context'
 import type { ErrorHandler } from '#src/types/contexts/ErrorContext'
 import type { InterceptHandler, TerminateHandler } from '#src/types'
 import { NotFoundException } from '#src/exceptions/NotFoundException'
 import type { FastifyReply, FastifyRequest, RouteHandlerMethod } from 'fastify'
 
+const otelApi = await Module.safeImport('@opentelemetry/api')
+
 export class FastifyHandler {
   /**
    * Parse the fastify request handler and the preHandler hook to an Athenna
    * request handler.
    */
   public static request(handler: RequestHandler): RouteHandlerMethod {
     return async (req: FastifyRequest, res: FastifyReply) => {
-      if (!req.data) {
-        req.data = {}
-      }
-
-      const ctx: any = {}
-
-      ctx.data = req.data
-      ctx.request = new Request(req)
-      ctx.response = new Response(res, ctx.request)
+      const ctx = this.createContext(req, res)
 
-      await handler(ctx)
+      await this.runWithOtelContext(req, ctx, () => handler(ctx))
     }
   }
 
@@ -49,24 +45,15 @@ export class FastifyHandler {
    */
   public static intercept(handler: InterceptHandler) {
     return async (req: FastifyRequest, res: FastifyReply, payload: any) => {
-      if (!req.data) {
-        req.data = {}
-      }
-
       if (Is.Json(payload)) {
         payload = JSON.parse(payload)
       }
 
       res.body = payload
 
-      const ctx: any = {}
+      const ctx = this.createContext(req, res, { status: res.statusCode })
 
-      ctx.data = req.data
-      ctx.request = new Request(req)
-      ctx.response = new Response(res, ctx.request)
-      ctx.status = ctx.response.statusCode
-
-      payload = await handler(ctx)
+      payload = await this.runWithOtelContext(req, ctx, () => handler(ctx))
 
       res.body = payload
 
@@ -83,19 +70,12 @@ export class FastifyHandler {
    */
   public static terminate(handler: TerminateHandler) {
     return async (req: FastifyRequest, res: FastifyReply) => {
-      if (!req.data) {
-        req.data = {}
-      }
+      const ctx = this.createContext(req, res, {
+        status: res.statusCode,
+        responseTime: res.elapsedTime
+      })
 
-      const ctx: any = {}
-
-      ctx.data = req.data
-      ctx.request = new Request(req)
-      ctx.response = new Response(res, ctx.request)
-      ctx.status = ctx.response.statusCode
-      ctx.responseTime = ctx.response.elapsedTime
-
-      await handler(ctx)
+      await this.runWithOtelContext(req, ctx, () => handler(ctx))
     }
   }
 
@@ -104,18 +84,9 @@ export class FastifyHandler {
    */
   public static error(handler: ErrorHandler) {
     return async (error: any, req: FastifyRequest, res: FastifyReply) => {
-      if (!req.data) {
-        req.data = {}
-      }
-
-      const ctx: any = {}
+      const ctx = this.createContext(req, res, { error })
 
-      ctx.data = req.data
-      ctx.request = new Request(req)
-      ctx.response = new Response(res, ctx.request)
-      ctx.error = error
-
-      await handler(ctx)
+      await this.runWithOtelContext(req, ctx, () => handler(ctx))
     }
   }
 
@@ -124,20 +95,68 @@ export class FastifyHandler {
    */
   public static notFoundError(handler: ErrorHandler) {
     return async (req: FastifyRequest, res: FastifyReply) => {
-      if (!req.data) {
-        req.data = {}
+      const ctx = this.createContext(req, res, {
+        error: new NotFoundException(`Route ${req.method}:${req.url} not found`)
+      })
+
+      await this.runWithOtelContext(req, ctx, () => handler(ctx))
+    }
+  }
+
+  private static createContext(
+    req: FastifyRequest,
+    res: FastifyReply,
+    extra = {}
+  ) {
+    if (!req.data) {
+      req.data = {}
+    }
+
+    const request = new Request(req)
+
+    return {
+      data: req.data,
+      request,
+      response: new Response(res, request),
+      ...extra
+    } as any
+  }
+
+  private static isOtelContextEnabled() {
+    return Config.is('http.otel.contextEnabled', true)
+  }
+
+  private static getOrCreateOtelContext(req: FastifyRequest, ctx: any) {
+    if (req.otelContext) {
+      return req.otelContext as OtelContext
+    }
+
+    let otelContext = otelApi.context.active()
+
+    for (const binding of Config.get('http.otel.contextBindings', [])) {
+      const value = binding.resolve(ctx)
+
+      if (Is.Undefined(value) && !binding.includeIfUndefined) {
+        continue
       }
 
-      const ctx: any = {}
+      otelContext = otelContext.setValue(binding.key, value)
+    }
+
+    req.otelContext = otelContext
 
-      ctx.data = req.data
-      ctx.request = new Request(req)
-      ctx.response = new Response(res, ctx.request)
-      ctx.error = new NotFoundException(
-        `Route ${req.method}:${req.url} not found`
-      )
+    return otelContext as OtelContext
+  }
 
-      await handler(ctx)
+  private static runWithOtelContext(
+    req: FastifyRequest,
+    ctx: any,
+    callback: () => any
+  ) {
+    if (!this.isOtelContextEnabled() || !otelApi) {
+      return callback()
     }
+
+    return otelApi.context.with(this.getOrCreateOtelContext(req, ctx), callback)
   }
 }

src/index.ts

@@ -10,6 +10,7 @@
 declare module 'fastify' {
   interface FastifyRequest {
     data: any
+    otelContext?: import('@opentelemetry/api').Context
     zodParsed?: {
       body?: any
       headers?: any

tests/unit/server/ServerTest.ts

@@ -7,19 +7,24 @@
  * file that was distributed with this source code.
  */
 
+import { Config } from '@athenna/config'
 import { Server, HttpServerProvider } from '#src'
-import { Test, AfterEach, BeforeEach, type Context } from '@athenna/test'
+import { context, createContextKey } from '@opentelemetry/api'
+import { Test, AfterEach, BeforeEach, type Context, Cleanup } from '@athenna/test'
+import { AsyncLocalStorageContextManager } from '@opentelemetry/context-async-hooks'
 
 export default class ServerTest {
   @BeforeEach()
   public async beforeEach() {
     ioc.reconstruct()
 
+    context.setGlobalContextManager(new AsyncLocalStorageContextManager().enable())
     new HttpServerProvider().register()
   }
 
   @AfterEach()
   public async afterEach() {
+    context.disable()
     await new HttpServerProvider().shutdown()
   }
 
@@ -193,4 +198,99 @@ export default class ServerTest {
     assert.isTrue(terminated)
     assert.deepEqual(response.json(), { hello: 'world' })
   }
+
+  @Test()
+  public async shouldKeepOriginalRequestBodyAvailableAfterInterceptingTheResponse({ assert }: Context) {
+    let requestBody: any
+    let responseBody: any
+
+    Server.intercept(ctx => {
+      return { ...ctx.response.body, intercepted: true }
+    })
+      .terminate(ctx => {
+        requestBody = ctx.request.body
+        responseBody = ctx.response.body
+      })
+      .post({
+        url: '/test',
+        handler: async ctx => ctx.response.send({ hello: 'world' })
+      })
+
+    const response = await Server.request().post('/test').body({ foo: 'bar' })
+
+    assert.deepEqual(response.json(), { hello: 'world', intercepted: true })
+    assert.deepEqual(requestBody, { foo: 'bar' })
+    assert.deepEqual(responseBody, { hello: 'world', intercepted: true })
+  }
+
+  @Test()
+  @Cleanup(() => Config.set('http.otel.contextEnabled', false))
+  @Cleanup(() => Config.set('http.otel.contextBindings', []))
+  public async shouldBindConfiguredOtelContextValuesAcrossTheRequestLifecycle({ assert }: Context) {
+    const methodKey = createContextKey('http.method')
+    const stageKey = createContextKey('request.stage')
+    let terminateValues: any = {}
+
+    Config.set('http.otel.contextEnabled', true)
+    Config.set('http.otel.contextBindings', [
+      { key: methodKey, resolve: ctx => ctx.request.method },
+      { key: stageKey, resolve: ctx => ctx.data.stage }
+    ])
+
+    Server.terminate(ctx => {
+      terminateValues = {
+        method: context.active().getValue(methodKey),
+        stage: context.active().getValue(stageKey),
+        status: ctx.status
+      }
+    }).post({
+      url: '/test',
+      data: { stage: 'route-default' },
+      handler: async ctx =>
+        ctx.response.send({
+          method: context.active().getValue(methodKey),
+          stage: context.active().getValue(stageKey)
+        })
+    })
+
+    const response = await Server.request({ path: '/test', method: 'POST' })
+
+    assert.deepEqual(response.json(), {
+      method: 'POST',
+      stage: 'route-default'
+    })
+    assert.deepEqual(terminateValues, {
+      method: 'POST',
+      stage: 'route-default',
+      status: 200
+    })
+  }
+
+  @Test()
+  @Cleanup(() => Config.set('http.otel.contextEnabled', false))
+  @Cleanup(() => Config.set('http.otel.contextBindings', []))
+  public async shouldReuseConfiguredOtelContextValuesInsideErrorHandlers({ assert }: Context) {
+    const routeKey = createContextKey('request.route')
+
+    Config.set('http.otel.contextEnabled', true)
+    Config.set('http.otel.contextBindings', [{ key: routeKey, resolve: ctx => ctx.request.baseUrl }])
+
+    Server.setErrorHandler(async ctx => {
+      await ctx.response.status(500).send({
+        route: context.active().getValue(routeKey)
+      })
+    })
+
+    Server.get({
+      url: '/boom',
+      handler: async () => {
+        throw new Error('boom')
+      }
+    })
+
+    const response = await Server.request().get('/boom')
+
+    assert.equal(response.statusCode, 500)
+    assert.deepEqual(response.json(), { route: '/boom' })
+  }
 }