Webpack: abstract away getting compilation spans

Currently many of our plugins import a singleton mapping of webpack compilation objects to their corresponding tracing span.

This change abstracts that away into a function called `getCompilationSpan`, and allows us to use the corresponding mapping in the rspack case.

Previously, any plugin that used these would fail with rspack when the compilation object wouldn’t be found. This was the case when running `next build` with rspack.

Test Plan: Install and run `next build` with the rspack plugin in a test app.

Author: wbinnssmith

packages/next/src/build/webpack/plugins/build-manifest-plugin.ts

@@ -17,11 +17,8 @@ import type { BuildManifest } from '../../../server/get-page-files'
 import getRouteFromEntrypoint from '../../../server/get-route-from-entrypoint'
 import { ampFirstEntryNamesMap } from './next-drop-client-page-plugin'
 import { getSortedRoutes } from '../../../shared/lib/router/utils'
-import { spans as webpackSpans } from './profiling-plugin'
-import { compilationSpans as rspackSpans } from './rspack-profiling-plugin'
 import { Span } from '../../../trace'
-
-const compilationSpans = process.env.NEXT_RSPACK ? rspackSpans : webpackSpans
+import { getCompilationSpan } from '../utils'
 
 type DeepMutable<T> = { -readonly [P in keyof T]: DeepMutable<T[P]> }
 
@@ -109,9 +106,9 @@ export function generateClientManifest(
   compilation?: any
 ): string | undefined {
   const compilationSpan = compilation
-    ? compilationSpans.get(compilation)
+    ? getCompilationSpan(compilation)
     : compiler
-      ? compilationSpans.get(compiler)
+      ? getCompilationSpan(compiler)
       : new Span({ name: 'client-manifest' })
 
   const genClientManifestSpan = compilationSpan?.traceChild(
@@ -205,7 +202,7 @@ export default class BuildManifestPlugin {
 
   createAssets(compiler: any, compilation: any) {
     const compilationSpan =
-      compilationSpans.get(compilation) ?? compilationSpans.get(compiler)
+      getCompilationSpan(compilation) ?? getCompilationSpan(compiler)
     if (!compilationSpan) {
       throw new Error('No span found for compilation')
     }

packages/next/src/build/webpack/plugins/css-minimizer-plugin.ts

@@ -3,7 +3,7 @@ import postcssScss from 'next/dist/compiled/postcss-scss'
 import postcss from 'postcss'
 import type { Parser } from 'postcss'
 import { webpack, sources } from 'next/dist/compiled/webpack/webpack'
-import { spans } from './profiling-plugin'
+import { getCompilationSpan } from '../utils'
 
 // https://github.com/NMFR/optimize-css-assets-webpack-plugin/blob/0a410a9bf28c7b0e81a3470a13748e68ca2f50aa/src/index.js#L20
 const CSS_REGEX = /\.css(\?.*)?$/i
@@ -64,7 +64,8 @@ export class CssMinimizerPlugin {
           stage: webpack.Compilation.PROCESS_ASSETS_STAGE_OPTIMIZE_SIZE,
         },
         async (assets: any) => {
-          const compilationSpan = spans.get(compilation) || spans.get(compiler)
+          const compilationSpan =
+            getCompilationSpan(compilation) || getCompilationSpan(compiler)
           const cssMinimizerSpan = compilationSpan!.traceChild(
             'css-minimizer-plugin'
           )

packages/next/src/build/webpack/plugins/minify-webpack-plugin/src/index.ts

@@ -7,7 +7,7 @@ import {
   type Compilation,
 } from 'next/dist/compiled/webpack/webpack'
 import pLimit from 'next/dist/compiled/p-limit'
-import { spans } from '../../profiling-plugin'
+import { getCompilationSpan } from '../../../utils'
 
 function buildError(error: any, file: string) {
   if (error.line) {
@@ -48,7 +48,8 @@ export class MinifyPlugin {
     }
   ) {
     const mangle = !this.options.noMangling
-    const compilationSpan = spans.get(compilation)! || spans.get(compiler)
+    const compilationSpan =
+      getCompilationSpan(compilation)! || getCompilationSpan(compiler)
 
     const MinifierSpan = compilationSpan.traceChild(
       'minify-webpack-plugin-optimize'

packages/next/src/build/webpack/plugins/next-trace-entrypoints-plugin.ts

@@ -1,6 +1,5 @@
 import nodePath from 'path'
 import type { Span } from '../../../trace'
-import { spans } from './profiling-plugin'
 import isError from '../../../lib/is-error'
 import { nodeFileTrace } from 'next/dist/compiled/@vercel/nft'
 import type { NodeFileTraceReasons } from 'next/dist/compiled/@vercel/nft'
@@ -20,6 +19,7 @@ import { getModuleBuildInfo } from '../loaders/get-module-build-info'
 import { getPageFilePath } from '../../entries'
 import { resolveExternal } from '../../handle-externals'
 import { isStaticMetadataRoute } from '../../../lib/metadata/is-metadata-route'
+import { getCompilationSpan } from '../utils'
 
 const PLUGIN_NAME = 'TraceEntryPointsPlugin'
 export const TRACE_IGNORES = [
@@ -576,6 +576,12 @@ export class TraceEntryPointsPlugin implements webpack.WebpackPluginInstance {
 
   apply(compiler: webpack.Compiler) {
     compiler.hooks.compilation.tap(PLUGIN_NAME, (compilation) => {
+      const compilationSpan =
+        getCompilationSpan(compilation) || getCompilationSpan(compiler)!
+      const traceEntrypointsPluginSpan = compilationSpan.traceChild(
+        'next-trace-entrypoint-plugin'
+      )
+
       compilation.hooks.processAssets.tapAsync(
         {
           name: PLUGIN_NAME,
@@ -633,10 +639,6 @@ export class TraceEntryPointsPlugin implements webpack.WebpackPluginInstance {
         }
       }
 
-      const compilationSpan = spans.get(compilation) || spans.get(compiler)!
-      const traceEntrypointsPluginSpan = compilationSpan.traceChild(
-        'next-trace-entrypoint-plugin'
-      )
       traceEntrypointsPluginSpan.traceFn(() => {
         compilation.hooks.processAssets.tapAsync(
           {

packages/next/src/build/webpack/utils.ts

@@ -5,9 +5,13 @@ import type {
   NormalModule,
   Module,
   ModuleGraph,
+  Compiler,
 } from 'webpack'
 import type { ModuleGraphConnection } from 'webpack'
 import { getAppLoader } from '../entries'
+import { spans as webpackCompilationSpans } from './plugins/profiling-plugin'
+import { compilationSpans as rspackCompilationSpans } from './plugins/rspack-profiling-plugin'
+import type { Span } from '../../trace'
 
 export function traverseModules(
   compilation: Compilation,
@@ -114,3 +118,13 @@ export function getModuleReferencesInOrder(
   connections.sort((a, b) => a.index - b.index)
   return connections.map((c) => c.connection)
 }
+
+export function getCompilationSpan(
+  compilation: Compiler | Compilation
+): Span | undefined {
+  const compilationSpans = process.env.NEXT_RSPACK
+    ? rspackCompilationSpans
+    : webpackCompilationSpans
+
+  return compilationSpans.get(compilation)
+}