Improvements to webpack tracing, including hot-reload

packages/next/build/tracer.ts

@@ -2,6 +2,49 @@ import api, { Span } from '@opentelemetry/api'
 
 export const tracer = api.trace.getTracer('next', process.env.__NEXT_VERSION)
 
+const compilerStacks = new WeakMap()
+
+export function stackPush(compiler: any, spanName: string, attrs?: any): any {
+  let stack = compilerStacks.get(compiler)
+  let span
+
+  if (!stack) {
+    compilerStacks.set(compiler, (stack = []))
+    span = tracer.startSpan(spanName, attrs ? attrs() : undefined)
+  } else {
+    const parent = stack[stack.length - 1]
+    tracer.withSpan(parent, () => {
+      span = tracer.startSpan(spanName, attrs ? attrs() : undefined)
+    })
+  }
+
+  stack.push(span)
+  return span
+}
+
+export function stackPop(compiler: any, span: any) {
+  span.end()
+
+  let stack = compilerStacks.get(compiler)
+  if (!stack) {
+    console.warn(
+      'Attempted to pop from non-existent stack. Compiler reference must be bad.'
+    )
+    return
+  }
+  const poppedSpan = stack.pop()
+  if (poppedSpan !== span) {
+    stack.push(poppedSpan)
+    const spanIdx = stack.indexOf(span)
+    console.warn('Attempted to pop span that was not at top of stack.')
+    if (spanIdx !== -1) {
+      console.info(
+        `Span was found at index ${spanIdx} with stack size ${stack.length}`
+      )
+    }
+  }
+}
+
 export function traceFn<T extends (...args: unknown[]) => ReturnType<T>>(
   span: Span,
   fn: T

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

@@ -1,8 +1,13 @@
-import { tracer } from '../../tracer'
+import { tracer, stackPush, stackPop } from '../../tracer'
 import { webpack, isWebpack5 } from 'next/dist/compiled/webpack/webpack'
+import {
+  Span,
+  trace,
+  ProxyTracerProvider,
+  NoopTracerProvider,
+} from '@opentelemetry/api'
 
 const pluginName = 'ProfilingPlugin'
-
 export const spans = new WeakMap()
 
 function getNormalModuleLoaderHook(compilation: any) {
@@ -14,26 +19,87 @@ function getNormalModuleLoaderHook(compilation: any) {
   return compilation.hooks.normalModuleLoader
 }
 
+function tracingIsEnabled() {
+  const tracerProvider: any = trace.getTracerProvider()
+  if (tracerProvider instanceof ProxyTracerProvider) {
+    const proxyDelegate: any = tracerProvider.getDelegate()
+    return !(proxyDelegate instanceof NoopTracerProvider)
+  }
+  return false
+}
+
 export class ProfilingPlugin {
+  compiler: any
+
   apply(compiler: any) {
-    // Only enabled when instrumentation is loaded
-    const currentSpan = tracer.getCurrentSpan()
-    if (!currentSpan || !currentSpan.isRecording()) {
+    // Only enable plugin when instrumentation is loaded
+    if (!tracingIsEnabled()) {
       return
     }
+    this.traceTopLevelHooks(compiler)
+    this.traceCompilationHooks(compiler)
+    this.compiler = compiler
+  }
 
-    compiler.hooks.compile.tap(pluginName, () => {
-      const span = tracer.startSpan('webpack-compile', {
-        attributes: { name: compiler.name },
-      })
-      spans.set(compiler, span)
+  traceHookPair(
+    spanName: string,
+    startHook: any,
+    stopHook: any,
+    attrs?: any,
+    onSetSpan?: (span: Span | undefined) => void
+  ) {
+    let span: Span | undefined
+    startHook.tap(pluginName, () => {
+      span = stackPush(this.compiler, spanName, attrs)
+      onSetSpan?.(span)
+    })
+    stopHook.tap(pluginName, () => {
+      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)
+      }
     })
-    compiler.hooks.done.tap(pluginName, () => {
-      spans.get(compiler).end()
+    stopHook.tap(pluginName, () => {
+      stackPop(this.compiler, span)
     })
+  }
+
+  traceTopLevelHooks(compiler: any) {
+    this.traceHookPair(
+      'webpack-compile',
+      compiler.hooks.compile,
+      compiler.hooks.done,
+      () => {
+        return { attributes: { name: compiler.name } }
+      },
+      (span) => spans.set(compiler, span)
+    )
+    this.traceHookPair(
+      'webpack-prepare-env',
+      compiler.hooks.environment,
+      compiler.hooks.afterEnvironment
+    )
+    this.traceHookPair(
+      'webpack-invalidated',
+      compiler.hooks.invalid,
+      compiler.hooks.done
+    )
+  }
+
+  traceCompilationHooks(compiler: any) {
     compiler.hooks.compilation.tap(pluginName, (compilation: any) => {
       compilation.hooks.buildModule.tap(pluginName, (module: any) => {
-        tracer.withSpan(spans.get(compiler), () => {
+        const compilerSpan = spans.get(compiler)
+        if (!compilerSpan) {
+          return
+        }
+        tracer.withSpan(compilerSpan, () => {
           const span = tracer.startSpan('build-module')
           span.setAttribute('name', module.userRequest)
           spans.set(module, span)
@@ -51,6 +117,45 @@ export class ProfilingPlugin {
       compilation.hooks.succeedModule.tap(pluginName, (module: any) => {
         spans.get(module).end()
       })
+
+      if (isWebpack5) {
+        this.traceHookPair(
+          'webpack-compilation',
+          compilation.hooks.beforeCompile,
+          compilation.hooks.afterCompile
+        )
+      }
+
+      this.traceHookPair(
+        'webpack-compilation-chunk-graph',
+        compilation.hooks.beforeChunks,
+        compilation.hooks.afterChunks
+      )
+      this.traceHookPair(
+        'webpack-compilation-optimize',
+        compilation.hooks.optimize,
+        compilation.hooks.reviveModules
+      )
+      this.traceLoopedHook(
+        'webpack-compilation-optimize-modules',
+        compilation.hooks.optimizeModules,
+        compilation.hooks.afterOptimizeModules
+      )
+      this.traceLoopedHook(
+        'webpack-compilation-optimize-chunks',
+        compilation.hooks.optimizeChunks,
+        compilation.hooks.afterOptimizeChunks
+      )
+      this.traceHookPair(
+        'webpack-compilation-optimize-tree',
+        compilation.hooks.optimizeTree,
+        compilation.hooks.afterOptimizeTree
+      )
+      this.traceHookPair(
+        'webpack-compilation-hash',
+        compilation.hooks.beforeHash,
+        compilation.hooks.afterHash
+      )
     })
   }
 }