Improve jaeger traces (#28168)

Co-authored-by: Harsh Karande <harshskarande@gmail.com>

.vscode/launch.json

@@ -27,12 +27,12 @@
       "outFiles": ["${workspaceFolder}/packages/next/dist/**/*"]
     },
     {
-      "name": "Launch app build trace",
+      "name": "Launch app build trace jaeger",
       "type": "node",
       "request": "launch",
       "cwd": "${workspaceFolder}",
       "runtimeExecutable": "yarn",
-      "runtimeArgs": ["run", "trace-debug", "build", "test/integration/basic"],
+      "runtimeArgs": ["run", "clean-trace-jaeger"],
       "skipFiles": ["<node_internals>/**"],
       "port": 9229,
       "outFiles": ["${workspaceFolder}/packages/next/dist/**/*"]

package.json

@@ -38,6 +38,7 @@
     "lint-staged": "lint-staged",
     "next-with-deps": "./scripts/next-with-deps.sh",
     "next": "node --trace-deprecation --enable-source-maps packages/next/dist/bin/next",
+    "clean-trace-jaeger": "rm -rf test/integration/basic/.next && TRACE_TARGET=JAEGER node --trace-deprecation --enable-source-maps packages/next/dist/bin/next build test/integration/basic",
     "debug": "node --inspect packages/next/dist/bin/next"
   },
   "pre-commit": "lint-staged",

packages/next/build/index.ts

@@ -73,7 +73,7 @@ import { generateBuildId } from './generate-build-id'
 import { isWriteable } from './is-writeable'
 import * as Log from './output/log'
 import createSpinner from './spinner'
-import { trace, setGlobal } from '../telemetry/trace'
+import { trace, flushAllTraces, setGlobal } from '../telemetry/trace'
 import {
   collectPages,
   detectConflictingPaths,
@@ -121,7 +121,7 @@ export default async function build(
 ): Promise<void> {
   const nextBuildSpan = trace('next-build')
 
-  return nextBuildSpan.traceAsyncFn(async () => {
+  const buildResult = await nextBuildSpan.traceAsyncFn(async () => {
     // attempt to load global env values so they are available in next.config.js
     const { loadedEnvFiles } = nextBuildSpan
       .traceChild('load-dotenv')
@@ -522,7 +522,8 @@ export default async function build(
         ignore: [] as string[],
       }))
 
-    const configs = await nextBuildSpan
+    const runWebpackSpan = nextBuildSpan.traceChild('run-webpack-compiler')
+    const configs = await runWebpackSpan
       .traceChild('generate-webpack-config')
       .traceAsyncFn(() =>
         Promise.all([
@@ -535,6 +536,7 @@ export default async function build(
             pagesDir,
             entrypoints: entrypoints.client,
             rewrites,
+            runWebpackSpan,
           }),
           getBaseWebpackConfig(dir, {
             buildId,
@@ -545,6 +547,7 @@ export default async function build(
             pagesDir,
             entrypoints: entrypoints.server,
             rewrites,
+            runWebpackSpan,
           }),
         ])
       )
@@ -566,24 +569,22 @@ export default async function build(
 
     let result: CompilerResult = { warnings: [], errors: [] }
     // We run client and server compilation separately to optimize for memory usage
-    await nextBuildSpan
-      .traceChild('run-webpack-compiler')
-      .traceAsyncFn(async () => {
-        const clientResult = await runCompiler(clientConfig)
-        // Fail build if clientResult contains errors
-        if (clientResult.errors.length > 0) {
-          result = {
-            warnings: [...clientResult.warnings],
-            errors: [...clientResult.errors],
-          }
-        } else {
-          const serverResult = await runCompiler(configs[1])
-          result = {
-            warnings: [...clientResult.warnings, ...serverResult.warnings],
-            errors: [...clientResult.errors, ...serverResult.errors],
-          }
+    await runWebpackSpan.traceAsyncFn(async () => {
+      const clientResult = await runCompiler(clientConfig)
+      // Fail build if clientResult contains errors
+      if (clientResult.errors.length > 0) {
+        result = {
+          warnings: [...clientResult.warnings],
+          errors: [...clientResult.errors],
         }
-      })
+      } else {
+        const serverResult = await runCompiler(configs[1])
+        result = {
+          warnings: [...clientResult.warnings, ...serverResult.warnings],
+          errors: [...clientResult.errors, ...serverResult.errors],
+        }
+      }
+    })
 
     const webpackBuildEnd = process.hrtime(webpackBuildStart)
     if (buildSpinner) {
@@ -1734,6 +1735,11 @@ export default async function build(
       .traceChild('telemetry-flush')
       .traceAsyncFn(() => telemetry.flush())
   })
+
+  // Ensure all traces are flushed before finishing the command
+  await flushAllTraces()
+
+  return buildResult
 }
 
 export type ClientSsgManifest = Set<string>

packages/next/build/webpack-config.ts

@@ -54,6 +54,7 @@ import WebpackConformancePlugin, {
 import { WellKnownErrorsPlugin } from './webpack/plugins/wellknown-errors-plugin'
 import { regexLikeCss } from './webpack/config/blocks/css'
 import { CopyFilePlugin } from './webpack/plugins/copy-file-plugin'
+import type { Span } from '../telemetry/trace'
 
 type ExcludesFalse = <T>(x: T | false) => x is T
 
@@ -228,6 +229,7 @@ export default async function getBaseWebpackConfig(
     entrypoints,
     rewrites,
     isDevFallback = false,
+    runWebpackSpan,
   }: {
     buildId: string
     config: NextConfigComplete
@@ -239,6 +241,7 @@ export default async function getBaseWebpackConfig(
     entrypoints: WebpackEntrypoints
     rewrites: CustomRoutes['rewrites']
     isDevFallback?: boolean
+    runWebpackSpan?: Span
   }
 ): Promise<webpack.Configuration> {
   const hasRewrites =
@@ -1306,7 +1309,7 @@ export default async function getBaseWebpackConfig(
         new BuildStatsPlugin({
           distDir,
         }),
-      new ProfilingPlugin(),
+      new ProfilingPlugin({ runWebpackSpan }),
       config.optimizeFonts &&
         !dev &&
         isServer &&

packages/next/build/webpack/plugins/profiling-plugin.ts

@@ -1,5 +1,5 @@
 import { webpack, isWebpack5 } from 'next/dist/compiled/webpack/webpack'
-import { trace, stackPush, stackPop, Span } from '../../../telemetry/trace'
+import { trace, Span } from '../../../telemetry/trace'
 
 const pluginName = 'ProfilingPlugin'
 export const spans = new WeakMap<any, Span>()
@@ -15,7 +15,11 @@ function getNormalModuleLoaderHook(compilation: any) {
 
 export class ProfilingPlugin {
   compiler: any
+  runWebpackSpan: Span | undefined
 
+  constructor({ runWebpackSpan }: { runWebpackSpan: Span | undefined }) {
+    this.runWebpackSpan = runWebpackSpan
+  }
   apply(compiler: any) {
     this.traceTopLevelHooks(compiler)
     this.traceCompilationHooks(compiler)
@@ -26,12 +30,19 @@ export class ProfilingPlugin {
     spanName: string,
     startHook: any,
     stopHook: any,
-    attrs?: any,
-    onSetSpan?: (span: Span) => void
+    {
+      parentSpan,
+      attrs,
+      onSetSpan,
+    }: {
+      parentSpan?: () => Span | undefined
+      attrs?: any
+      onSetSpan?: (span: Span) => void
+    } = {}
   ) {
     let span: Span | undefined
     startHook.tap(pluginName, () => {
-      span = stackPush(this.compiler, spanName, attrs)
+      span = trace(spanName, parentSpan?.()?.id, attrs ? attrs() : attrs)
       onSetSpan?.(span)
     })
     stopHook.tap(pluginName, () => {
@@ -41,56 +52,39 @@ export class ProfilingPlugin {
       if (!span) {
         return
       }
-      stackPop(this.compiler, span)
-    })
-  }
-
-  traceLoopedHook(spanName: string, startHook: any, stopHook: any) {
-    let span: Span | undefined
-    startHook.tap(pluginName, () => {
-      if (!span) {
-        span = stackPush(this.compiler, spanName)
-      }
-    })
-    stopHook.tap(pluginName, () => {
-      stackPop(this.compiler, span)
+      span.stop()
     })
   }
 
   traceTopLevelHooks(compiler: any) {
     this.traceHookPair(
-      'webpack-compile',
-      compiler.hooks.compile,
-      compiler.hooks.done,
-      () => {
-        return { name: compiler.name }
-      },
-      (span) => spans.set(compiler, span)
-    )
-    this.traceHookPair(
-      'webpack-prepare-env',
-      compiler.hooks.environment,
-      compiler.hooks.afterEnvironment
+      'webpack-compilation',
+      isWebpack5 ? compiler.hooks.beforeCompile : compiler.hooks.compile,
+      isWebpack5 ? compiler.hooks.afterCompile : compiler.hooks.done,
+      {
+        parentSpan: () => this.runWebpackSpan,
+        attrs: () => ({ name: compiler.name }),
+        onSetSpan: (span) => spans.set(compiler, span),
+      }
     )
+
     if (compiler.options.mode === 'development') {
       this.traceHookPair(
         'webpack-invalidated',
         compiler.hooks.invalid,
         compiler.hooks.done,
-        () => ({ name: compiler.name })
+        { attrs: () => ({ name: compiler.name }) }
       )
     }
   }
 
   traceCompilationHooks(compiler: any) {
-    if (isWebpack5) {
-      this.traceHookPair(
-        'webpack-compilation',
-        compiler.hooks.beforeCompile,
-        compiler.hooks.afterCompile,
-        () => ({ name: compiler.name })
-      )
-    }
+    this.traceHookPair(
+      'webpack-emit',
+      compiler.hooks.emit,
+      compiler.hooks.afterEmit,
+      { parentSpan: () => this.runWebpackSpan }
+    )
 
     compiler.hooks.compilation.tap(pluginName, (compilation: any) => {
       compilation.hooks.buildModule.tap(pluginName, (module: any) => {
@@ -118,8 +112,8 @@ export class ProfilingPlugin {
       getNormalModuleLoaderHook(compilation).tap(
         pluginName,
         (loaderContext: any, module: any) => {
-          const parentSpan = spans.get(module)
-          loaderContext.currentTraceSpan = parentSpan
+          const moduleSpan = spans.get(module)
+          loaderContext.currentTraceSpan = moduleSpan
         }
       )
 
@@ -130,32 +124,38 @@ export class ProfilingPlugin {
       this.traceHookPair(
         'webpack-compilation-chunk-graph',
         compilation.hooks.beforeChunks,
-        compilation.hooks.afterChunks
+        compilation.hooks.afterChunks,
+        { parentSpan: () => this.runWebpackSpan }
       )
       this.traceHookPair(
         'webpack-compilation-optimize',
         compilation.hooks.optimize,
-        compilation.hooks.reviveModules
+        compilation.hooks.reviveModules,
+        { parentSpan: () => this.runWebpackSpan }
       )
-      this.traceLoopedHook(
+      this.traceHookPair(
         'webpack-compilation-optimize-modules',
         compilation.hooks.optimizeModules,
-        compilation.hooks.afterOptimizeModules
+        compilation.hooks.afterOptimizeModules,
+        { parentSpan: () => this.runWebpackSpan }
       )
-      this.traceLoopedHook(
+      this.traceHookPair(
         'webpack-compilation-optimize-chunks',
         compilation.hooks.optimizeChunks,
-        compilation.hooks.afterOptimizeChunks
+        compilation.hooks.afterOptimizeChunks,
+        { parentSpan: () => this.runWebpackSpan }
       )
       this.traceHookPair(
         'webpack-compilation-optimize-tree',
         compilation.hooks.optimizeTree,
-        compilation.hooks.afterOptimizeTree
+        compilation.hooks.afterOptimizeTree,
+        { parentSpan: () => this.runWebpackSpan }
       )
       this.traceHookPair(
         'webpack-compilation-hash',
         compilation.hooks.beforeHash,
-        compilation.hooks.afterHash
+        compilation.hooks.afterHash,
+        { parentSpan: () => this.runWebpackSpan }
       )
     })
   }