[Cache Components] Allow span creation while prerendering (vercel#82350)

When creating spans with OTel random values must be created to follow
the OTel spec and is expected in practical implementations. While it is
a loose concession Next.js will allow spans to be created while
prerendering that have ids generated with random values. and by reading
the current time from the system clock.

This is pragmatic but potentially risky because if you are not careful
and for instance wrap a "use cache" function in a span it is possible
you will accidentally pass the span into the cache function causing it
to miss on every invocation due to the random nature of one of it's
arguments.

To defend against this we do check whether `startActiveSpan` is being
passed a Cache Function. However it is easy to circumvent by wrapping it
in an intermediate function or by using startSpan and context.with
yourself.

In the long run we should consider tainting the Span objects created in
these scopes so they cannot be sent through to a cache. I will consider
a follow up PR to land this if we determine taint is likely to be
stabilized in React.

The pratical implementation of this change is to patch the tracer
provider so that it always provides a patched tracer which exits the
workUnitStorage on span start/creation and re-enters the scope on inner
function exectuion.

Author: gnoff

packages/next/errors.json

@@ -772,5 +772,6 @@
   "771": "\\`%s\\` was called during a runtime prerender. Next.js should be preventing %s from being included in server components statically, but did not in this case.",
   "772": "FetchStrategy.PPRRuntime should never be used when `experimental.clientSegmentCache` is disabled",
   "773": "Missing workStore in createPrerenderParamsForClientSegment",
-  "774": "Route %s used %s outside of a Server Component. This is not allowed."
+  "774": "Route %s used %s outside of a Server Component. This is not allowed.",
+  "775": "Node.js instrumentation extensions should not be loaded in the Edge runtime."
 }

packages/next/src/server/lib/router-utils/instrumentation-globals.external.ts

@@ -6,6 +6,7 @@ import type {
   InstrumentationOnRequestError,
 } from '../../instrumentation/types'
 import { interopDefault } from '../../../lib/interop-default'
+import { afterRegistration as extendInstrumentationAfterRegistration } from './instrumentation-node-extensions'
 
 let cachedInstrumentationModule: InstrumentationModule
 
@@ -55,6 +56,7 @@ async function registerInstrumentation(projectDir: string, distDir: string) {
   if (instrumentation?.register) {
     try {
       await instrumentation.register()
+      extendInstrumentationAfterRegistration()
     } catch (err: any) {
       err.message = `An error occurred while loading instrumentation hook: ${err.message}`
       throw err

packages/next/src/server/lib/router-utils/instrumentation-node-extensions.ts

@@ -0,0 +1,122 @@
+/**
+ * This extension augments opentelemetry after registration if applicable.
+ * This extension must only be loaded in Node environments.
+ */
+
+import type { Tracer } from '@opentelemetry/api'
+import {
+  type WorkUnitStore,
+  workUnitAsyncStorage,
+} from '../../app-render/work-unit-async-storage.external'
+import { InvariantError } from '../../../shared/lib/invariant-error'
+import { isUseCacheFunction } from '../../../lib/client-and-server-references'
+
+export function afterRegistration(): void {
+  if (process.env.NEXT_RUNTIME === 'edge') {
+    throw new InvariantError(
+      'Node.js instrumentation extensions should not be loaded in the Edge runtime.'
+    )
+  }
+
+  extendTracerProviderForCacheComponents()
+}
+
+// In theory we only want to enable this extension when cacheComponents is enabled
+// however there are certain servers that might load instrumentation before nextConfig is available
+// and so gating it on the config might lead to skipping this extension even when it is necessary.
+// When cacheComponents is disabled this extension should be a no-op so we enable it universally.
+// Additionally, soon, cacheComponents will be enabled always so this just pulls the extension forward in time
+function extendTracerProviderForCacheComponents(): void {
+  let api: typeof import('next/dist/compiled/@opentelemetry/api')
+
+  // we want to allow users to use their own version of @opentelemetry/api if they
+  // want to, so we try to require it first, and if it fails we fall back to the
+  // version that is bundled with Next.js
+  // this is because @opentelemetry/api has to be synced with the version of
+  // @opentelemetry/tracing that is used, and we don't want to force users to use
+  // the version that is bundled with Next.js.
+  // the API is ~stable, so this should be fine
+  try {
+    api = require('@opentelemetry/api') as typeof import('@opentelemetry/api')
+  } catch (err) {
+    api =
+      require('next/dist/compiled/@opentelemetry/api') as typeof import('next/dist/compiled/@opentelemetry/api')
+  }
+
+  const provider = api.trace.getTracerProvider()
+
+  // When Cache Components is enabled we need to instrument the tracer
+  // to exit the workUnitAsyncStorage context when generating spans.
+  const originalGetTracer = provider.getTracer.bind(provider)
+  provider.getTracer = (...args) => {
+    const tracer = originalGetTracer.apply(provider, args)
+    if (WeakTracers.has(tracer)) {
+      return tracer
+    }
+    const originalStartSpan = tracer.startSpan
+    tracer.startSpan = (...startSpanArgs) => {
+      return workUnitAsyncStorage.exit(() =>
+        originalStartSpan.apply(tracer, startSpanArgs)
+      )
+    }
+
+    const originalStartActiveSpan = tracer.startActiveSpan
+    // @ts-ignore TS doesn't recognize the overloads correctly
+    tracer.startActiveSpan = (...startActiveSpanArgs: any[]) => {
+      const workUnitStore = workUnitAsyncStorage.getStore()
+      if (!workUnitStore) {
+        // @ts-ignore TS doesn't recognize the overloads correctly
+        return originalStartActiveSpan.apply(tracer, startActiveSpanArgs)
+      }
+
+      let fnIdx: number = 0
+      if (
+        startActiveSpanArgs.length === 2 &&
+        typeof startActiveSpanArgs[1] === 'function'
+      ) {
+        fnIdx = 1
+      } else if (
+        startActiveSpanArgs.length === 3 &&
+        typeof startActiveSpanArgs[2] === 'function'
+      ) {
+        fnIdx = 2
+      } else if (
+        startActiveSpanArgs.length > 3 &&
+        typeof startActiveSpanArgs[3] === 'function'
+      ) {
+        fnIdx = 3
+      }
+
+      if (fnIdx) {
+        const originalFn = startActiveSpanArgs[fnIdx]
+        if (isUseCacheFunction(originalFn)) {
+          console.error(
+            'A Cache Function (`use cache`) was passed to startActiveSpan which means it will receive a Span argument with a possibly random ID on every invocation leading to cache misses. Provide a wrapping function around the Cache Function that does not forward the Span argument to avoid this issue.'
+          )
+        }
+        startActiveSpanArgs[fnIdx] = withWorkUnitContext(
+          workUnitStore,
+          originalFn
+        )
+      }
+
+      return workUnitAsyncStorage.exit(() => {
+        // @ts-ignore TS doesn't recognize the overloads correctly
+        return originalStartActiveSpan.apply(tracer, startActiveSpanArgs)
+      })
+    }
+
+    WeakTracers.add(tracer)
+    return tracer
+  }
+}
+
+const WeakTracers = new WeakSet<Tracer>()
+
+function withWorkUnitContext(
+  workUnitStore: WorkUnitStore,
+  fn: (...args: any[]) => any
+) {
+  return (...args: any[]) =>
+    workUnitAsyncStorage.run(workUnitStore, fn, ...args)
+}

test/e2e/app-dir/cache-components-allow-otel-spans/.gitignore

@@ -0,0 +1,2 @@
+node_modules
+pnpm-lock.yaml

test/e2e/app-dir/cache-components-allow-otel-spans/app/[slug]/cache/page.tsx

@@ -0,0 +1,36 @@
+import {
+  CachedInnerTraceManualSpan,
+  InnerTraceManualSpan,
+  CachedTracedComponentManualSpan,
+  TracedComponentManualSpan,
+  CachedInnerTraceActiveSpan,
+  InnerTraceActiveSpan,
+  CachedTracedComponentActiveSpan,
+  TracedComponentActiveSpan,
+} from '../../traced-work'
+
+export function generateStaticParams() {
+  return [{ slug: 'prerendered' }]
+}
+
+export default async function Page({
+  params,
+}: {
+  params: Promise<{ slug: string }>
+}) {
+  'use cache'
+  return (
+    <>
+      <h1>{(await params).slug}</h1>
+      <div>We are inside a "use cache" scope</div>
+      <CachedInnerTraceManualSpan />
+      <InnerTraceManualSpan />
+      <CachedTracedComponentManualSpan />
+      <TracedComponentManualSpan />
+      <CachedInnerTraceActiveSpan />
+      <InnerTraceActiveSpan />
+      <CachedTracedComponentActiveSpan />
+      <TracedComponentActiveSpan />
+    </>
+  )
+}

test/e2e/app-dir/cache-components-allow-otel-spans/app/[slug]/fallback/layout.tsx

@@ -0,0 +1,9 @@
+import { Suspense } from 'react'
+
+export default async function Layout({
+  children,
+}: {
+  children: React.ReactNode
+}) {
+  return <Suspense fallback={<div>Loading...</div>}>{children}</Suspense>
+}

test/e2e/app-dir/cache-components-allow-otel-spans/app/[slug]/fallback/page.tsx

@@ -0,0 +1,35 @@
+import {
+  CachedInnerTraceManualSpan,
+  InnerTraceManualSpan,
+  CachedTracedComponentManualSpan,
+  TracedComponentManualSpan,
+  CachedInnerTraceActiveSpan,
+  InnerTraceActiveSpan,
+  CachedTracedComponentActiveSpan,
+  TracedComponentActiveSpan,
+} from '../../traced-work'
+
+export function generateStaticParams() {
+  return [{ slug: 'prerendered' }]
+}
+
+export default async function Page({
+  params,
+}: {
+  params: Promise<{ slug: string }>
+}) {
+  return (
+    <>
+      <h1>{(await params).slug}</h1>
+      <div>We are inside a "use server" scope</div>
+      <CachedInnerTraceManualSpan />
+      <InnerTraceManualSpan />
+      <CachedTracedComponentManualSpan />
+      <TracedComponentManualSpan />
+      <CachedInnerTraceActiveSpan />
+      <InnerTraceActiveSpan />
+      <CachedTracedComponentActiveSpan />
+      <TracedComponentActiveSpan />
+    </>
+  )
+}

test/e2e/app-dir/cache-components-allow-otel-spans/app/[slug]/server/page.tsx

@@ -0,0 +1,35 @@
+import {
+  CachedInnerTraceManualSpan,
+  InnerTraceManualSpan,
+  CachedTracedComponentManualSpan,
+  TracedComponentManualSpan,
+  CachedInnerTraceActiveSpan,
+  InnerTraceActiveSpan,
+  CachedTracedComponentActiveSpan,
+  TracedComponentActiveSpan,
+} from '../../traced-work'
+
+export function generateStaticParams() {
+  return [{ slug: 'prerendered' }]
+}
+
+export default async function Page({
+  params,
+}: {
+  params: Promise<{ slug: string }>
+}) {
+  return (
+    <>
+      <h1>{(await params).slug}</h1>
+      <div>We are inside a "use server" scope</div>
+      <CachedInnerTraceManualSpan />
+      <InnerTraceManualSpan />
+      <CachedTracedComponentManualSpan />
+      <TracedComponentManualSpan />
+      <CachedInnerTraceActiveSpan />
+      <InnerTraceActiveSpan />
+      <CachedTracedComponentActiveSpan />
+      <TracedComponentActiveSpan />
+    </>
+  )
+}

test/e2e/app-dir/cache-components-allow-otel-spans/app/favicon.ico

@@ -0,0 +1,7 @@
+export default function Root({ children }: { children: React.ReactNode }) {
+  return (
+    <html>
+      <body>{children}</body>
+    </html>
+  )
+}