fix!: move span out of H3EventContext (#13)

Author: huang-julien

src/augment.ts

@@ -1,10 +1,17 @@
-import type { Span } from "@opentelemetry/api";
+import type { Context, Span } from "@opentelemetry/api";
 import type { H3Event } from "h3"
 import { Presets } from "./types";
 
 declare module 'h3' {
-    interface H3EventContext {
-        span: Span
+    interface H3Event {
+        otel: {
+            span: Span
+            /**
+             * @internal
+             */
+            __endTime: number|undefined
+            ctx: Context
+        }
     }
 }
 

src/runtime/plugin.ts

@@ -4,19 +4,17 @@ import { ATTR_URL_PATH, ATTR_URL_FULL, ATTR_HTTP_REQUEST_METHOD, ATTR_HTTP_RESPO
 import type { NitroAppPlugin, NitroRuntimeHooks } from "nitropack";
 import { getResponseStatus, getRequestProtocol, getRequestURL, getHeaders } from "h3"
 
-const context = api.context, trace = api.trace
-
+const context = api.context, trace = api.trace 
 export default <NitroAppPlugin>((nitro) => {
-
     nitro.hooks.hook('request', async (event) => {
         const tracer = trace.getTracer('nitro-opentelemetry')
         const requestURL = getRequestURL(event)
-        const currentContext = context.active()
+        const currentContext = context.active() 
 
         // Extract the parent context from the headers if it exists
-        const parentCtx = api.propagation.extract(currentContext, getHeaders(event));
-        
-        const span = tracer.startSpan(await getSpanName(event), {
+        // If a span is already set in the context, use it as the parent
+        const parentCtx = trace.getSpan(currentContext) ? currentContext : api.propagation.extract(currentContext, getHeaders(event));
+         const span = tracer.startSpan(await getSpanName(event), {
             attributes: {
                 [ATTR_URL_PATH]: event.context.matchedRoute?.path || event.path,
                 [ATTR_URL_FULL]: event.path,
@@ -27,27 +25,30 @@ export default <NitroAppPlugin>((nitro) => {
             },
             kind: api.SpanKind.SERVER
         }, parentCtx)
-        trace.setSpan(context.active(), span)
-        event.context.span = span
-        event.context.__otel = {}
+        const ctx =  trace.setSpan(context.active(), span) 
+        event.otel = {
+            span,
+            __endTime: undefined
+            ,ctx
+        } 
     })
 
     nitro.hooks.hook('beforeResponse', (event) => {
-        event.context.__otel.endTime = Date.now()
+        event.otel.__endTime = Date.now()
     })
 
     nitro.hooks.hook('afterResponse', async (event) => {
-        event.context.span.setAttribute(ATTR_HTTP_RESPONSE_STATUS_CODE, getResponseStatus(event))
-        await nitro.hooks.callHook('otel:span:end', { event, span: event.context.span })
-        event.context.span.end(event.context.__otel.endTime) 
+        event.otel.span.setAttribute(ATTR_HTTP_RESPONSE_STATUS_CODE, getResponseStatus(event))
+        await nitro.hooks.callHook('otel:span:end', { event, span: event.otel.span })
+        event.otel.span.end(event.otel.__endTime) 
     })
 
     nitro.hooks.hook('error', async (error, { event }) => {
         if (event) {
-            event.context.span.recordException(error)
-            event.context.span.setAttribute(ATTR_HTTP_RESPONSE_STATUS_CODE, getResponseStatus(event))
-            await nitro.hooks.callHook('otel:span:end', { event, span: event.context.span })
-            event.context.span.end()
+            event.otel.span.recordException(error)
+            event.otel.span.setAttribute(ATTR_HTTP_RESPONSE_STATUS_CODE, getResponseStatus(event))
+            await nitro.hooks.callHook('otel:span:end', { event, span: event.otel.span })
+            event.otel.span.end()
         } else {
             const span = trace.getSpan(api.ROOT_CONTEXT)
             span?.recordException(error)

src/runtime/utils.ts

@@ -1,10 +1,10 @@
 import * as api from "@opentelemetry/api" 
 import { defineEventHandler } from "h3";
 
-const context = api.context, trace = api.trace
+const context = api.context
 
 export function defineTracedEventHandler(handler: ReturnType<typeof defineEventHandler>) { 
     return defineEventHandler((event) => { 
-       return  context.with(trace.setSpan(context.active(), event.context.span),  handler, undefined,  event) 
+       return  context.with(event.otel.ctx, handler, undefined,  event) 
     })
 }

test/fixtures/basic/init.ts

@@ -3,9 +3,21 @@ import * as opentelemetry from '@opentelemetry/sdk-node';
 import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
 import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-proto';
 import { UndiciInstrumentation } from '@opentelemetry/instrumentation-undici';
- 
+import { NodeTracerProvider } from '@opentelemetry/sdk-trace-node'
+import { AsyncLocalStorageContextManager } from '@opentelemetry/context-async-hooks';
+
+const contextManager = new AsyncLocalStorageContextManager();
+// Create and configure NodeTracerProvider
+const provider = new NodeTracerProvider();
+
+// Initialize the provider
+provider.register({
+  contextManager,
+});
+
+
 const sdk = new opentelemetry.NodeSDK({
-  traceExporter: new OTLPTraceExporter(), 
+  traceExporter: new OTLPTraceExporter(),
   instrumentations: [getNodeAutoInstrumentations(), new UndiciInstrumentation()],
 });
 sdk.start();

test/fixtures/basic/routes/another-endpoint.ts

@@ -1,7 +1,8 @@
 export default defineTracedEventHandler((e) => {
     // @ts-expect-error internal API
-    const parentSpanId = e.context.span.parentSpanId
+    const parentSpanId = e.otel.span.parentSpanId
     return {
-        traceId: e.context.span.spanContext().traceId,parentSpanId,
+        traceId: e.otel.span.spanContext().traceId,
+        parentSpanId,
     }
 })

test/fixtures/basic/routes/dynamic/[slug].ts

@@ -1,5 +1,5 @@
 
 export default defineEventHandler((event) => {
     // @ts-expect-error - internal property
-    return event.context.span.name 
+    return event.otel.span.name 
 })

test/fixtures/basic/routes/event-fetch.ts

@@ -0,0 +1,13 @@
+export default defineTracedEventHandler(async (e) => {
+    const { traceId, parentSpanId } = await e.$fetch('/another-endpoint')
+    return {
+        traceId: e.otel.span.spanContext().traceId,
+        spanId: e.otel.span.spanContext().spanId,
+        anotherEndpoint: {
+            traceId,
+            parentSpanId
+        },
+        // @ts-expect-error internal API
+        parentSpanId: e.otel.span.parentSpanId
+    }
+})

test/fixtures/basic/routes/index.ts

@@ -1,8 +1,8 @@
 export default defineTracedEventHandler(async (e) => {
-    const { traceId, parentSpanId } = await $fetch('/another-endpoint')
+    const { traceId, parentSpanId } = await globalThis.$fetch('/another-endpoint')
     return {
-        traceId: e.context.span.spanContext().traceId,
-        spanId: e.context.span.spanContext().spanId,
+        traceId: e.otel.span.spanContext().traceId,
+        spanId: e.otel.span.spanContext().spanId,
         anotherEndpoint: {
             traceId,
             parentSpanId

test/fixtures/basic/routes/wait-ms.ts

@@ -0,0 +1,15 @@
+
+export default defineTracedEventHandler(async (e) => {
+    const ms = Number(getQuery(e).ms)
+    
+    await new Promise((resolve) => setTimeout(resolve, ms))
+    const { traceId, parentSpanId } = await globalThis.$fetch('/another-endpoint')
+    return {
+        traceId: e.otel.span.spanContext().traceId,
+        spanId: e.otel.span.spanContext().spanId,
+        anotherEndpoint: {
+            traceId,
+            parentSpanId
+        }
+    }
+})

test/index.test.ts

@@ -1,31 +1,72 @@
 import { describe, expect, it } from 'vitest'
 import { $fetchRaw } from 'nitro-test-utils/e2e'
- const dummyTrace = '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01'
+const dummyTrace = '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01'
+const dummyTrace2 = '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01'
 
 
 describe('traces', async () => {
     it('should trace requests', async () => {
-        const { data } = await $fetchRaw('')
+        const { data } = await $fetchRaw('/')
 
         expect(data.traceId).toBeDefined()
     })
 
     it('should keep trace context', async () => {
-        
-            const { data } = await $fetchRaw('', {
+        const { data } = await $fetchRaw('/', {
+            headers: {
+                traceparent: dummyTrace
+            }
+        })
+        // assert that the traceId is the same
+        expect(data.traceId).toBe(dummyTrace.split('-')[1])
+
+        // assert that localFetch is correctly traced
+        expect(data.anotherEndpoint.traceId).toBe(data.traceId)
+        expect(data.anotherEndpoint.parentSpanId).toBe(data.spanId)
+
+    })
+
+    it('event.$fetch', async () => {
+        const { data } = await $fetchRaw('event-fetch', {
+            headers: {
+                traceparent: dummyTrace
+            }
+        })
+        // assert that the traceId is the same
+        expect(data.traceId).toBe(dummyTrace.split('-')[1])
+
+        // assert that localFetch is correctly traced
+        expect(data.anotherEndpoint.traceId).toBe(data.traceId)
+        expect(data.anotherEndpoint.parentSpanId).toBe(data.spanId)
+    })
+
+    it('expect no XRSP', async () => {
+      const [{data}, {data: data2}] =   await Promise.all([
+            $fetchRaw('/wait-ms?ms=350', {
                 headers: {
-                    traceparent:dummyTrace
+                    traceparent: dummyTrace
                 }
-            }) 
-            // assert that the traceId is the same
-            expect(data.traceId).toBe(dummyTrace.split('-')[1])
-
-            // assert that localFetch is correctly traced
-            expect(data.anotherEndpoint.traceId).toBe(data.traceId)
-            expect(data.anotherEndpoint.parentSpanId).toBe(data.spanId)
-       
-    })
+            }),
+            $fetchRaw('/wait-ms?ms=100', {
+                headers: {
+                    traceparent: dummyTrace2
+                }
+            })
+        ])
 
+        
+        // assert that the traceId is the same
+        expect(data.traceId).toBe(dummyTrace.split('-')[1])
+
+        // assert that localFetch is correctly traced
+        expect(data.anotherEndpoint.traceId).toBe(data.traceId)
+        expect(data.anotherEndpoint.parentSpanId).toBe(data.spanId)
+
+        // assert that localFetch is correctly traced
+        expect(data2.traceId).toBe(dummyTrace2.split('-')[1])
+        expect(data2.anotherEndpoint.traceId).toBe(data2.traceId)
+        expect(data2.anotherEndpoint.parentSpanId).toBe(data2.spanId)
+    })
 })
 
 describe('names', () => {