OpenTelemetry: trace API routes in page router (#62120)

Co-authored-by: JJ Kasper <jj@jjsweb.site>

packages/next/src/build/templates/pages-edge-api.ts

@@ -4,6 +4,7 @@ import '../../server/web/globals'
 
 import { adapter } from '../../server/web/adapter'
 import { IncrementalCache } from '../../server/lib/incremental-cache'
+import { wrapApiHandler } from '../../server/api-utils'
 
 // Import the userland code.
 import handler from 'VAR_USERLAND'
@@ -23,6 +24,6 @@ export default function (
     ...opts,
     IncrementalCache,
     page: 'VAR_DEFINITION_PATHNAME',
-    handler,
+    handler: wrapApiHandler(page, handler),
   })
 }

packages/next/src/build/webpack/loaders/next-edge-function-loader.ts

@@ -40,6 +40,7 @@ const nextEdgeFunctionLoader: webpack.LoaderDefinitionFunction<EdgeFunctionLoade
         import 'next/dist/esm/server/web/globals'
         import { adapter } from 'next/dist/esm/server/web/adapter'
         import { IncrementalCache } from 'next/dist/esm/server/lib/incremental-cache'
+        import { wrapApiHandler } from 'next/dist/esm/server/api-utils'
 
         import handler from ${stringifiedPagePath}
 
@@ -52,7 +53,7 @@ const nextEdgeFunctionLoader: webpack.LoaderDefinitionFunction<EdgeFunctionLoade
               ...opts,
               IncrementalCache,
               page: ${JSON.stringify(page)},
-              handler,
+              handler: wrapApiHandler(${JSON.stringify(page)}, handler),
           })
         }
     `

packages/next/src/server/api-utils/index.ts

@@ -8,6 +8,8 @@ import {
   PRERENDER_REVALIDATE_HEADER,
   PRERENDER_REVALIDATE_ONLY_GENERATED_HEADER,
 } from '../../lib/constants'
+import { getTracer } from '../lib/trace/tracer'
+import { NodeSpan } from '../lib/trace/constants'
 
 export type NextApiRequestCookies = Partial<{ [key: string]: string }>
 export type NextApiRequestQuery = Partial<{ [key: string]: string | string[] }>
@@ -18,6 +20,23 @@ export type __ApiPreviewProps = {
   previewModeSigningKey: string
 }
 
+export function wrapApiHandler<T extends (...args: any[]) => any>(
+  page: string,
+  handler: T
+): T {
+  return ((...args) => {
+    getTracer().getRootSpanAttributes()?.set('next.route', page)
+    // Call API route method
+    return getTracer().trace(
+      NodeSpan.runHandler,
+      {
+        spanName: `executing api route (pages) ${page}`,
+      },
+      () => handler(...args)
+    )
+  }) as T
+}
+
 /**
  *
  * @param res response object

packages/next/src/server/api-utils/node/api-resolver.ts

@@ -23,8 +23,6 @@ import {
   RESPONSE_LIMIT_DEFAULT,
 } from './../index'
 import { getCookieParser } from './../get-cookie-parser'
-import { getTracer } from '../../lib/trace/tracer'
-import { NodeSpan } from '../../lib/trace/constants'
 import {
   PRERENDER_REVALIDATE_HEADER,
   PRERENDER_REVALIDATE_ONLY_GENERATED_HEADER,
@@ -416,15 +414,7 @@ export async function apiResolver(
       res.once('pipe', () => (wasPiped = true))
     }
 
-    getTracer().getRootSpanAttributes()?.set('next.route', page)
-    // Call API route method
-    const apiRouteResult = await getTracer().trace(
-      NodeSpan.runHandler,
-      {
-        spanName: `executing api route (pages) ${page}`,
-      },
-      () => resolver(req, res)
-    )
+    const apiRouteResult = await resolver(req, res)
 
     if (process.env.NODE_ENV !== 'production') {
       if (typeof apiRouteResult !== 'undefined') {

packages/next/src/server/future/route-modules/pages-api/module.ts

@@ -2,7 +2,7 @@ import type { IncomingMessage, ServerResponse } from 'http'
 import type { PagesAPIRouteDefinition } from '../../route-definitions/pages-api-route-definition'
 import type { PageConfig } from '../../../../../types'
 import type { ParsedUrlQuery } from 'querystring'
-import type { __ApiPreviewProps } from '../../../api-utils'
+import { wrapApiHandler, type __ApiPreviewProps } from '../../../api-utils'
 import type { RouteModuleOptions } from '../route-module'
 
 import { RouteModule, type RouteModuleHandleContext } from '../route-module'
@@ -104,6 +104,8 @@ export class PagesAPIRouteModule extends RouteModule<
   PagesAPIRouteDefinition,
   PagesAPIUserlandModule
 > {
+  private apiResolverWrapped: typeof apiResolver
+
   constructor(options: PagesAPIRouteModuleOptions) {
     super(options)
 
@@ -112,6 +114,11 @@ export class PagesAPIRouteModule extends RouteModule<
         `Page ${options.definition.page} does not export a default function.`
       )
     }
+
+    this.apiResolverWrapped = wrapApiHandler(
+      options.definition.page,
+      apiResolver
+    )
   }
 
   /**
@@ -125,7 +132,8 @@ export class PagesAPIRouteModule extends RouteModule<
     res: ServerResponse,
     context: PagesAPIRouteHandlerContext
   ): Promise<void> {
-    await apiResolver(
+    const { apiResolverWrapped } = this
+    await apiResolverWrapped(
       req,
       res,
       context.query,

test/e2e/opentelemetry/app/api/app/[param]/data/edge/route.ts

@@ -0,0 +1,6 @@
+export async function GET() {
+  return new Response(JSON.stringify({ test: 'data-edge' }))
+}
+
+export const runtime = 'edge'
+export const dynamic = 'force-dynamic'

test/e2e/opentelemetry/app/app/[param]/rsc-fetch/edge/page.tsx

@@ -0,0 +1,13 @@
+// We want to trace this fetch in runtime
+export const dynamic = 'force-dynamic'
+
+export const runtime = 'edge'
+
+export async function generateMetadata() {
+  return {}
+}
+
+export default async function Page() {
+  const data = await fetch('https://vercel.com')
+  return <pre>RESONSE: {data.status}</pre>
+}

test/e2e/opentelemetry/collector.ts

@@ -0,0 +1,62 @@
+import { Server as HttpServer } from 'node:http'
+import { SavedSpan } from './constants'
+
+export interface Collector {
+  getSpans: () => SavedSpan[]
+  shutdown: () => Promise<void>
+}
+
+export async function connectCollector({
+  port,
+}: {
+  port: number
+}): Promise<Collector> {
+  const spans: SavedSpan[] = []
+
+  const server = new HttpServer(async (req, res) => {
+    if (req.method !== 'POST') {
+      res.writeHead(405)
+      res.end()
+      return
+    }
+
+    const body = await new Promise<Buffer>((resolve, reject) => {
+      const acc: Buffer[] = []
+      req.on('data', (chunk: Buffer) => {
+        acc.push(chunk)
+      })
+      req.on('end', () => {
+        resolve(Buffer.concat(acc))
+      })
+      req.on('error', reject)
+    })
+
+    const newSpans = JSON.parse(body.toString('utf-8')) as SavedSpan[]
+    spans.push(...newSpans)
+    res.statusCode = 202
+    res.end()
+  })
+
+  await new Promise<void>((resolve, reject) => {
+    server.listen(port, (err?: Error) => {
+      if (err) {
+        reject(err)
+      } else {
+        resolve()
+      }
+    })
+  })
+
+  return {
+    getSpans() {
+      return spans
+    },
+    shutdown() {
+      return new Promise<void>((resolve, reject) =>
+        server.close((err) => (err ? reject(err) : resolve()))
+      ).catch((err) => {
+        console.warn('WARN: collector server disconnect failure:', err)
+      })
+    },
+  }
+}

test/e2e/opentelemetry/constants.ts

@@ -1,8 +1,7 @@
 import { SpanKind } from '@opentelemetry/api'
 
-export const traceFile = 'otel-trace.txt'
-
 export type SavedSpan = {
+  runtime?: string
   traceId?: string
   parentId?: string
   traceState?: any

test/e2e/opentelemetry/instrumentation-polyfill.ts

@@ -0,0 +1,3 @@
+if (globalThis.performance === undefined) {
+  globalThis.performance = { timeOrigin: 0, now: () => Date.now() } as any
+}

test/e2e/opentelemetry/instrumentation-test.ts

@@ -0,0 +1,93 @@
+import './instrumentation-polyfill'
+
+import { Resource } from '@opentelemetry/resources'
+import { SemanticResourceAttributes } from '@opentelemetry/semantic-conventions'
+import { BasicTracerProvider } from '@opentelemetry/sdk-trace-base'
+import { AsyncLocalStorageContextManager } from '@opentelemetry/context-async-hooks'
+import {
+  SimpleSpanProcessor,
+  SpanExporter,
+  ReadableSpan,
+} from '@opentelemetry/sdk-trace-base'
+import {
+  ExportResult,
+  ExportResultCode,
+  hrTimeToMicroseconds,
+} from '@opentelemetry/core'
+
+import { SavedSpan } from './constants'
+
+const serializeSpan = (span: ReadableSpan): SavedSpan => ({
+  runtime: process.env.NEXT_RUNTIME,
+  traceId: span.spanContext().traceId,
+  parentId: span.parentSpanId,
+  traceState: span.spanContext().traceState?.serialize(),
+  name: span.name,
+  id: span.spanContext().spanId,
+  kind: span.kind,
+  timestamp: hrTimeToMicroseconds(span.startTime),
+  duration: hrTimeToMicroseconds(span.duration),
+  attributes: span.attributes,
+  status: span.status,
+  events: span.events,
+  links: span.links,
+})
+
+class TestExporter implements SpanExporter {
+  constructor(private port: number) {}
+
+  async export(
+    spans: ReadableSpan[],
+    resultCallback: (result: ExportResult) => void
+  ) {
+    try {
+      const response = await fetch(`http://localhost:${this.port}`, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+        },
+        body: JSON.stringify(spans.map(serializeSpan)),
+      })
+      try {
+        await response.arrayBuffer()
+      } catch (e) {
+        // ignore.
+      }
+      if (response.status >= 400) {
+        console.warn('WARN: TestExporter: response status:', response.status)
+        resultCallback({
+          code: ExportResultCode.FAILED,
+          error: new Error(`http status ${response.status}`),
+        })
+      }
+    } catch (e) {
+      console.warn('WARN: TestExporterP: error:', e)
+      resultCallback({ code: ExportResultCode.FAILED, error: e })
+    }
+
+    resultCallback({ code: ExportResultCode.SUCCESS })
+  }
+  shutdown(): Promise<void> {
+    return Promise.resolve()
+  }
+}
+
+export const register = () => {
+  const contextManager = new AsyncLocalStorageContextManager()
+  contextManager.enable()
+
+  const provider = new BasicTracerProvider({
+    resource: new Resource({
+      [SemanticResourceAttributes.SERVICE_NAME]: 'test-next-app',
+    }),
+  })
+
+  if (!process.env.TEST_OTEL_COLLECTOR_PORT) {
+    throw new Error('TEST_OTEL_COLLECTOR_PORT is not set')
+  }
+  const port = parseInt(process.env.TEST_OTEL_COLLECTOR_PORT)
+  provider.addSpanProcessor(new SimpleSpanProcessor(new TestExporter(port)))
+
+  // Make sure to register you provider
+  provider.register({ contextManager })
+}