OpenTelemetry: propagate context to sandbox (#58791)

The sandboxed request processing do not share the same async context with the `BaseServer` and thus the context should be propagated independently.

Notably, the `headers` API is different in `BaseServer` (`Record<string, string | string[]>`) vs in the sandbox (`Headers`), so additionally, we have to use the right `getter`.

Co-authored-by: Leah <8845940+ForsakenHarmony@users.noreply.github.com>

packages/next-swc/crates/next-core/src/next_import_map.rs

@@ -511,14 +511,16 @@ async fn insert_next_server_special_aliases(
         NextRuntime::Edge => request_to_import_mapping(context_dir, request),
         NextRuntime::NodeJs => external_request_to_import_mapping(request),
     };
+
+    import_map.insert_exact_alias(
+        "@opentelemetry/api",
+        // TODO(WEB-625) this actually need to prefer the local version of
+        // @opentelemetry/api
+        external_if_node(project_path, "next/dist/compiled/@opentelemetry/api"),
+    );
+
     match ty {
         ServerContextType::Pages { pages_dir } | ServerContextType::PagesApi { pages_dir } => {
-            import_map.insert_exact_alias(
-                "@opentelemetry/api",
-                // TODO(WEB-625) this actually need to prefer the local version of
-                // @opentelemetry/api
-                external_if_node(pages_dir, "next/dist/compiled/@opentelemetry/api/index.js"),
-            );
             insert_alias_to_alternatives(
                 import_map,
                 format!("{VIRTUAL_PACKAGE_NAME}/pages/_app"),
@@ -549,15 +551,6 @@ async fn insert_next_server_special_aliases(
         ServerContextType::AppSSR { app_dir }
         | ServerContextType::AppRSC { app_dir, .. }
         | ServerContextType::AppRoute { app_dir } => {
-            import_map.insert_exact_alias(
-                "@opentelemetry/api",
-                // TODO(WEB-625) this actually need to prefer the local version of
-                // @opentelemetry/api
-                request_to_import_mapping(
-                    app_dir,
-                    "next/dist/compiled/@opentelemetry/api/index.js",
-                ),
-            );
             import_map.insert_exact_alias(
                 "styled-jsx",
                 request_to_import_mapping(get_next_package(app_dir), "styled-jsx"),

packages/next/src/server/base-server.ts

@@ -784,7 +784,7 @@ export default abstract class Server<ServerOptions extends Options = Options> {
     const method = req.method.toUpperCase()
 
     const tracer = getTracer()
-    return tracer.withPropagatedContext(req, () => {
+    return tracer.withPropagatedContext(req.headers, () => {
       return tracer.trace(
         BaseServerSpan.handleRequest,
         {

packages/next/src/server/lib/trace/tracer.ts

@@ -1,4 +1,3 @@
-import type { BaseNextRequest } from '../../base-http'
 import type { SpanTypes } from './constants'
 import { NextVanillaSpanAllowlist } from './constants'
 
@@ -8,6 +7,7 @@ import type {
   SpanOptions,
   Tracer,
   AttributeValue,
+  TextMapGetter,
 } from 'next/dist/compiled/@opentelemetry/api'
 
 let api: typeof import('next/dist/compiled/@opentelemetry/api')
@@ -173,13 +173,17 @@ class NextTracerImpl implements NextTracer {
     return trace.getSpan(context?.active())
   }
 
-  public withPropagatedContext<T>(req: BaseNextRequest, fn: () => T): T {
+  public withPropagatedContext<T, C>(
+    carrier: C,
+    fn: () => T,
+    getter?: TextMapGetter<C>
+  ): T {
     const activeContext = context.active()
     if (trace.getSpanContext(activeContext)) {
       // Active span is already set, too late to propagate.
       return fn()
     }
-    const remoteContext = propagation.extract(activeContext, req.headers)
+    const remoteContext = propagation.extract(activeContext, carrier, getter)
     return context.with(remoteContext, fn)
   }
 

packages/next/src/server/web/adapter.ts

@@ -16,6 +16,8 @@ import { NEXT_QUERY_PARAM_PREFIX } from '../../lib/constants'
 import { ensureInstrumentationRegistered } from './globals'
 import { RequestAsyncStorageWrapper } from '../async-storage/request-async-storage-wrapper'
 import { requestAsyncStorage } from '../../client/components/request-async-storage.external'
+import { getTracer } from '../lib/trace/tracer'
+import type { TextMapGetter } from 'next/dist/compiled/@opentelemetry/api'
 
 class NextRequestHint extends NextRequest {
   sourcePage: string
@@ -43,6 +45,11 @@ class NextRequestHint extends NextRequest {
   }
 }
 
+const headersGetter: TextMapGetter<Headers> = {
+  keys: (headers) => Array.from(headers.keys()),
+  get: (headers, key) => headers.get(key) ?? undefined,
+}
+
 export type AdapterOptions = {
   handler: NextMiddleware
   page: string
@@ -176,31 +183,37 @@ export async function adapter(
   let response
   let cookiesFromResponse
 
-  // we only care to make async storage available for middleware
-  const isMiddleware =
-    params.page === '/middleware' || params.page === '/src/middleware'
-  if (isMiddleware) {
-    response = await RequestAsyncStorageWrapper.wrap(
-      requestAsyncStorage,
-      {
-        req: request,
-        renderOpts: {
-          onUpdateCookies: (cookies) => {
-            cookiesFromResponse = cookies
+  const tracer = getTracer()
+  response = await tracer.withPropagatedContext(
+    request.headers,
+    () => {
+      // we only care to make async storage available for middleware
+      const isMiddleware =
+        params.page === '/middleware' || params.page === '/src/middleware'
+      if (isMiddleware) {
+        return RequestAsyncStorageWrapper.wrap(
+          requestAsyncStorage,
+          {
+            req: request,
+            renderOpts: {
+              onUpdateCookies: (cookies) => {
+                cookiesFromResponse = cookies
+              },
+              // @ts-expect-error: TODO: investigate why previewProps isn't on RenderOpts
+              previewProps: prerenderManifest?.preview || {
+                previewModeId: 'development-id',
+                previewModeEncryptionKey: '',
+                previewModeSigningKey: '',
+              },
+            },
           },
-          // @ts-expect-error: TODO: investigate why previewProps isn't on RenderOpts
-          previewProps: prerenderManifest?.preview || {
-            previewModeId: 'development-id',
-            previewModeEncryptionKey: '',
-            previewModeSigningKey: '',
-          },
-        },
-      },
-      () => params.handler(request, event)
-    )
-  } else {
-    response = await params.handler(request, event)
-  }
+          () => params.handler(request, event)
+        )
+      }
+      return params.handler(request, event)
+    },
+    headersGetter
+  )
 
   // check if response is a Response object
   if (response && !(response instanceof Response)) {