[dynamicIO] fix: do not apply import tracking transform in edge (vercel#79284)

Dynamic import tracking should never be applied to edge runtime code,
because:
1. `dynamicIO` prerendering is not supported for edge
2. the implementation of `dynamicIO` prerendering relies on
node-specific APIs.

I missed this in vercel#74152. This PR fixes a couple things:
- We no longer apply the dynamic import tracking to any edge modules
- Just to be safe, `trackDynamicImport` and `new CacheSignal` will now
throw if it's called from the edge runtime
- However, if `track-module-loading` does end up in an edge module, we
don't want to error at the top-level when creating
`moduleLoadingSignal`, so it is now initialized lazily. This is a bit
ugly, but `AppRouteRouteModule` currently pulls all prerendering-related
code in at the top level (instead of e.g. a conditional `require()` when
prerendering), so it *is* included in edge route handlers, even if
unused.
- this also required changing `__next_require__`/`__next_chunk_load__`
to only track module loading in DIO prerenders, otherwise we'd attempt
to create `moduleLoadingSignal` in edge RSC too

I also fixed a potential unhandled rejection in `CacheSignal.trackRead`,
because `.finally` rejects even if the original error was caught
somewhere else. (thanks for finding this one @gnoff!)

Author: lubieowoce

crates/next-core/src/next_server/transforms.rs

@@ -178,7 +178,12 @@ pub async fn get_next_server_transforms_rules(
         ServerContextType::Middleware { .. } | ServerContextType::Instrumentation { .. } => false,
     };
 
-    if is_app_dir && *next_config.enable_dynamic_io().await? {
+    if is_app_dir &&
+        // `dynamicIO` is not supported in the edge runtime.
+        // (also, the code generated by the dynamic imports transform relies on `CacheSignal`, which uses nodejs-specific APIs)
+        next_runtime != NextRuntime::Edge &&
+        *next_config.enable_dynamic_io().await?
+    {
         rules.push(get_next_track_dynamic_imports_transform_rule(mdx_rs));
     }
 

packages/next/errors.json

@@ -677,5 +677,7 @@
   "676": "Invariant: missing internal router-server-methods this is an internal bug",
   "677": "`trackDynamicImport` should always receive a promise. Something went wrong in the dynamic imports transform.",
   "678": "CacheSignal got more endRead() calls than beginRead() calls",
-  "679": "A CacheSignal cannot subscribe to itself"
+  "679": "A CacheSignal cannot subscribe to itself",
+  "680": "CacheSignal cannot be used in the edge runtime, because `dynamicIO` does not support it.",
+  "681": "Dynamic imports should not be instrumented in the edge runtime, because `dynamicIO` doesn't support it"
 }

packages/next/src/build/webpack-config.ts

@@ -534,6 +534,7 @@ export default async function getBaseWebpackConfig(
       loader: 'next-swc-loader',
       options: {
         isServer: isNodeOrEdgeCompilation,
+        compilerType,
         rootDir: dir,
         pagesDir,
         appDir,

packages/next/src/build/webpack/loaders/next-swc-loader.ts

@@ -39,6 +39,10 @@ import {
   type SwcTransformTelemetryOutput,
 } from '../plugins/telemetry-plugin/update-telemetry-loader-context-from-swc'
 import type { LoaderContext } from 'webpack'
+import {
+  COMPILER_NAMES,
+  type CompilerNameValues,
+} from '../../../shared/lib/constants'
 
 const maybeExclude = (
   excludePath: string,
@@ -57,6 +61,7 @@ const maybeExclude = (
 export interface SWCLoaderOptions {
   rootDir: string
   isServer: boolean
+  compilerType: CompilerNameValues
   pagesDir?: string
   appDir?: string
   hasReactRefresh: boolean
@@ -214,10 +219,12 @@ async function loaderTransform(
 function shouldTrackDynamicImports(loaderOptions: SWCLoaderOptions): boolean {
   // we only need to track `import()` 1. in dynamicIO, 2. on the server (RSC and SSR)
   // (Note: logic duplicated in crates/next-core/src/next_server/transforms.rs)
-  const { nextConfig, isServer, bundleLayer } = loaderOptions
+  const { nextConfig, bundleLayer, compilerType } = loaderOptions
   return (
     !!nextConfig.experimental?.dynamicIO &&
-    isServer &&
+    // NOTE: `server` means nodejs. `dynamicIO` is not supported in the edge runtime, so we want to exclude it.
+    // (also, the code generated by the dynamic imports transform relies on `CacheSignal`, which uses nodejs-specific APIs)
+    compilerType === COMPILER_NAMES.server &&
     (bundleLayer === WEBPACK_LAYERS.reactServerComponents ||
       bundleLayer === WEBPACK_LAYERS.serverSideRendering)
   )

packages/next/src/server/app-render/app-render.tsx

@@ -1203,18 +1203,33 @@ async function renderToHTMLOrFlightImpl(
     // to treat chunk/module loading with similar semantics as cache reads to avoid
     // module loading from causing a prerender to abort too early.
 
+    const shouldTrackModuleLoading = () => {
+      if (!renderOpts.experimental.dynamicIO) {
+        return false
+      }
+      const workUnitStore = workUnitAsyncStorage.getStore()
+      return !!(
+        workUnitStore &&
+        (workUnitStore.type === 'prerender' || workUnitStore.type === 'cache')
+      )
+    }
+
     const __next_require__: typeof instrumented.require = (...args) => {
       const exportsOrPromise = instrumented.require(...args)
-      // requiring an async module returns a promise.
-      trackPendingImport(exportsOrPromise)
+      if (shouldTrackModuleLoading()) {
+        // requiring an async module returns a promise.
+        trackPendingImport(exportsOrPromise)
+      }
       return exportsOrPromise
     }
     // @ts-expect-error
     globalThis.__next_require__ = __next_require__
 
     const __next_chunk_load__: typeof instrumented.loadChunk = (...args) => {
       const loadingChunk = instrumented.loadChunk(...args)
-      trackPendingChunkLoad(loadingChunk)
+      if (shouldTrackModuleLoading()) {
+        trackPendingChunkLoad(loadingChunk)
+      }
       return loadingChunk
     }
     // @ts-expect-error

packages/next/src/server/app-render/cache-signal.ts

@@ -16,6 +16,15 @@ export class CacheSignal {
 
   private subscribedSignals: Set<CacheSignal> | null = null
 
+  constructor() {
+    if (process.env.NEXT_RUNTIME === 'edge') {
+      // we rely on `process.nextTick`, which is not supported in edge
+      throw new InvariantError(
+        'CacheSignal cannot be used in the edge runtime, because `dynamicIO` does not support it.'
+      )
+    }
+  }
+
   private noMorePendingCaches() {
     if (!this.tickPending) {
       this.tickPending = true
@@ -107,7 +116,9 @@ export class CacheSignal {
 
   trackRead<T>(promise: Promise<T>) {
     this.beginRead()
-    promise.finally(this.endRead.bind(this))
+    // `promise.finally()` still rejects, so don't use it here to avoid unhandled rejections
+    const onFinally = this.endRead.bind(this)
+    promise.then(onFinally, onFinally)
     return promise
   }
 

packages/next/src/server/app-render/module-loading/track-dynamic-import.ts

@@ -12,6 +12,12 @@ import { trackPendingImport } from './track-module-loading.external'
 export function trackDynamicImport<TExports extends Record<string, any>>(
   modulePromise: Promise<TExports>
 ): Promise<TExports> {
+  if (process.env.NEXT_RUNTIME === 'edge') {
+    throw new InvariantError(
+      "Dynamic imports should not be instrumented in the edge runtime, because `dynamicIO` doesn't support it"
+    )
+  }
+
   if (!isThenable(modulePromise)) {
     // We're expecting `import()` to always return a promise. If it's not, something's very wrong.
     throw new InvariantError(

packages/next/src/server/app-render/module-loading/track-module-loading.instance.ts

@@ -1,14 +1,26 @@
 import { CacheSignal } from '../cache-signal'
 import { isThenable } from '../../../shared/lib/is-thenable'
 
-/** Tracks all in-flight async imports and chunk loads. */
-const moduleLoadingSignal = new CacheSignal()
+/**
+ * Tracks all in-flight async imports and chunk loads.
+ * Initialized lazily, because we don't want this to error in case it gets pulled into an edge runtime module.
+ */
+let _moduleLoadingSignal: CacheSignal | null
+function getModuleLoadingSignal() {
+  if (!_moduleLoadingSignal) {
+    _moduleLoadingSignal = new CacheSignal()
+  }
+  return _moduleLoadingSignal
+}
 
 export function trackPendingChunkLoad(promise: Promise<unknown>) {
+  const moduleLoadingSignal = getModuleLoadingSignal()
   moduleLoadingSignal.trackRead(promise)
 }
 
 export function trackPendingImport(exportsOrPromise: unknown) {
+  const moduleLoadingSignal = getModuleLoadingSignal()
+
   // requiring an async module returns a promise.
   // if it's sync, there's nothing to track.
   if (isThenable(exportsOrPromise)) {
@@ -28,6 +40,8 @@ export function trackPendingImport(exportsOrPromise: unknown) {
  * So if we see one, we want to extend the duration of `cacheSignal` at least until the import/chunk-load is done.
  */
 export function trackPendingModules(cacheSignal: CacheSignal): void {
+  const moduleLoadingSignal = getModuleLoadingSignal()
+
   // We can't just use `cacheSignal.trackRead(moduleLoadingSignal.cacheReady())`,
   // because we might start and finish multiple batches of module loads while waiting for caches,
   // and `moduleLoadingSignal.cacheReady()` would resolve after the first batch.

test/e2e/app-dir/dynamic-io-dynamic-imports/bundled/app/not-instrumented/edge-route-handler/messages.ts

@@ -0,0 +1 @@
+export default { title: 'hello' }

test/e2e/app-dir/dynamic-io-dynamic-imports/bundled/app/not-instrumented/edge-route-handler/route.ts

@@ -0,0 +1,9 @@
+export const runtime = 'edge'
+
+export async function GET(_request: Request) {
+  // This import should not be instrumented, because edge routes are never prerendered.
+  // `trackDynamicImport` will throw if it's used in the edge runtime,
+  // so it's enough to just do an import() here and see if it succeeds.
+  const messages = (await import('./messages')).default
+  return new Response(messages.title)
+}

test/e2e/app-dir/dynamic-io-dynamic-imports/bundled/app/not-instrumented/middleware/page.tsx

@@ -0,0 +1,5 @@
+// This page is only used as a matcher target for the middleware.
+
+export default function Page() {
+  return 'hello'
+}

test/e2e/app-dir/dynamic-io-dynamic-imports/bundled/messages.ts

@@ -0,0 +1 @@
+export default { title: 'hello' }

test/e2e/app-dir/dynamic-io-dynamic-imports/bundled/middleware.ts

@@ -0,0 +1,12 @@
+import type { MiddlewareConfig } from 'next/server'
+
+export default async function middleware() {
+  // This import should not be instrumented.
+  // `trackDynamicImport` will throw if it's used in the edge runtime,
+  // so it's enough to just do an import() here and see if it succeeds.
+  await import('./messages')
+}
+
+export const config: MiddlewareConfig = {
+  matcher: ['/not-instrumented/middleware'],
+}

test/e2e/app-dir/dynamic-io-dynamic-imports/dynamic-io-dynamic-imports.test.ts

@@ -50,6 +50,7 @@ describe('async imports in dynamicIO', () => {
          "/inside-render/server/from-node-modules/esm/async-module",
          "/inside-render/server/from-node-modules/esm/sync-module",
          "/inside-render/server/sync-module",
+         "/not-instrumented/middleware",
          "/outside-of-render/client/async-module",
          "/outside-of-render/client/sync-module",
          "/outside-of-render/server/async-module",
@@ -154,6 +155,20 @@ describe('async imports in dynamicIO', () => {
       })
     })
   })
+
+  describe('are not instrumented in edge', () => {
+    it('middleware', async () => {
+      // indirectly tests the behavior of middleware by rendering a page which the middleware matches
+      await testPage('/not-instrumented/middleware')
+    })
+
+    it('edge route handler', async () => {
+      const result = await next
+        .fetch('/not-instrumented/edge-route-handler')
+        .then((res) => res.text())
+      expect(result).toBe('hello')
+    })
+  })
 })
 
 describe('async imports in dynamicIO - external packages', () => {