[dynamicIO] cache tracking for import() (#74152)

This PR introduces some extra tracking which makes us treat
`import(...)` like a call to a cached function. Thanks to this, `await
import(...)` will no longer cause dynamicity errors in `dynamicIO`.
The motivation is the same as allowing `fs.readFileSync` -- if something
is available on the server at prerender time, we don't consider it IO.

Fixes #72589
Closes #75132

### Implementation notes

The tracking is implemented via an SWC transform
(`track_dynamic_imports.ts`) that turns `import(...)` into
`trackDynamicImport(import(...))`.

`trackDynamicImport(promise)` tracks the promise globally, without using
`workUnitStore.cacheSignal`. The prospective render subscribes to
pending modules using `trackPendingModules(cacheSignal)`, which causes
the prospective render to wait for all `import()`s to finish before
proceeding to the final render. The mechanism is analogous to `'use
cache'`, but the "result" of an `import()` is stored *in the module
cache* instead; when we invoke the `import()` again in the actual
prerender, it will resolve at microtask-speed, like we need it to.

The transform is enabled for all modules that run server-side, both RSC
and SSR, because the prerender runs both. This also includes route
handlers, because we also use a `cacheSignal` when prerendering those.
Notably, we also instrument `import()` in `node_modules` to account for
libraries that do lazy initialization.

I've also had to adjust the prospective client render in PPR mode to
wait for `cacheSignal.cacheReady()` - otherwise, it wouldn't wait for
`import()`s to resolve before doing the final client render, which'd
subsequently cause a dynamicity error.
(we didn't need to wait for `cacheSignal` before, because all cached
promises would've already been awaited in the server prerender)

Finally, we no longer need `warmFlightResponse`, because
`trackPendingModules` already makes the `cacheSignal` wait for all
loading chunks to finish, so we don't need to do it ahead of time.

Author: lubieowoce

.gitignore

@@ -28,7 +28,7 @@ pids
 coverage
 
 # test output
-test/**/out*
+test/**/out/*
 test/**/next-env.d.ts
 .DS_Store
 /e2e-tests
@@ -42,7 +42,7 @@ test/traces
 .nvmrc
 
 # examples
-examples/**/out
+examples/**/out/*
 examples/**/.env*.local
 
 pr-stats.md

crates/next-core/src/next_import_map.rs

@@ -965,6 +965,13 @@ async fn insert_next_shared_aliases(
             "next/dist/build/webpack/loaders/next-flight-loader/cache-wrapper",
         ),
     );
+    import_map.insert_exact_alias(
+        "private-next-rsc-track-dynamic-import",
+        request_to_import_mapping(
+            project_path,
+            "next/dist/build/webpack/loaders/next-flight-loader/track-dynamic-import",
+        ),
+    );
 
     insert_turbopack_dev_alias(import_map).await?;
     insert_package_alias(

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

@@ -12,8 +12,8 @@ use crate::{
     next_shared::transforms::{
         get_next_dynamic_transform_rule, get_next_font_transform_rule, get_next_image_rule,
         get_next_lint_transform_rule, get_next_modularize_imports_rule,
-        get_next_pages_transforms_rule, get_server_actions_transform_rule,
-        next_amp_attributes::get_next_amp_attr_rule,
+        get_next_pages_transforms_rule, get_next_track_dynamic_imports_transform_rule,
+        get_server_actions_transform_rule, next_amp_attributes::get_next_amp_attr_rule,
         next_cjs_optimizer::get_next_cjs_optimizer_rule,
         next_disallow_re_export_all_in_page::get_next_disallow_export_all_in_page_rule,
         next_edge_node_api_assert::next_edge_node_api_assert,
@@ -178,6 +178,10 @@ pub async fn get_next_server_transforms_rules(
         ServerContextType::Middleware { .. } | ServerContextType::Instrumentation { .. } => false,
     };
 
+    if is_app_dir && *next_config.enable_dynamic_io().await? {
+        rules.push(get_next_track_dynamic_imports_transform_rule(mdx_rs));
+    }
+
     if !foreign_code {
         rules.push(
             get_next_dynamic_transform_rule(true, is_server_components, is_app_dir, mode, mdx_rs)

crates/next-core/src/next_shared/transforms/mod.rs

@@ -16,6 +16,7 @@ pub(crate) mod next_pure;
 pub(crate) mod next_react_server_components;
 pub(crate) mod next_shake_exports;
 pub(crate) mod next_strip_page_exports;
+pub(crate) mod next_track_dynamic_imports;
 pub(crate) mod react_remove_properties;
 pub(crate) mod relay;
 pub(crate) mod remove_console;
@@ -30,6 +31,7 @@ pub use next_dynamic::get_next_dynamic_transform_rule;
 pub use next_font::get_next_font_transform_rule;
 pub use next_lint::get_next_lint_transform_rule;
 pub use next_strip_page_exports::get_next_pages_transforms_rule;
+pub use next_track_dynamic_imports::get_next_track_dynamic_imports_transform_rule;
 pub use server_actions::get_server_actions_transform_rule;
 use turbo_tasks::{ReadRef, ResolvedVc, Value};
 use turbo_tasks_fs::FileSystemPath;

crates/next-core/src/next_shared/transforms/next_track_dynamic_imports.rs

@@ -0,0 +1,24 @@
+use anyhow::Result;
+use async_trait::async_trait;
+use next_custom_transforms::transforms::track_dynamic_imports::*;
+use swc_core::ecma::ast::Program;
+use turbopack::module_options::ModuleRule;
+use turbopack_ecmascript::{CustomTransformer, TransformContext};
+
+use super::get_ecma_transform_rule;
+
+pub fn get_next_track_dynamic_imports_transform_rule(mdx_rs: bool) -> ModuleRule {
+    get_ecma_transform_rule(Box::new(NextTrackDynamicImports {}), mdx_rs, false)
+}
+
+#[derive(Debug)]
+struct NextTrackDynamicImports {}
+
+#[async_trait]
+impl CustomTransformer for NextTrackDynamicImports {
+    #[tracing::instrument(level = tracing::Level::TRACE, name = "next_track_dynamic_imports", skip_all)]
+    async fn transform(&self, program: &mut Program, ctx: &TransformContext<'_>) -> Result<()> {
+        program.mutate(track_dynamic_imports(ctx.unresolved_mark));
+        Ok(())
+    }
+}

crates/next-custom-transforms/src/chain_transforms.rs

@@ -121,6 +121,9 @@ pub struct TransformOptions {
 
     #[serde(default)]
     pub css_env: Option<swc_core::ecma::preset_env::Config>,
+
+    #[serde(default)]
+    pub track_dynamic_imports: bool,
 }
 
 pub fn custom_before_pass<'a, C>(
@@ -333,6 +336,14 @@ where
                 )),
                 None => Either::Right(noop_pass()),
             },
+            match &opts.track_dynamic_imports {
+                true => Either::Left(
+                    crate::transforms::track_dynamic_imports::track_dynamic_imports(
+                        unresolved_mark,
+                    ),
+                ),
+                false => Either::Right(noop_pass()),
+            },
             match &opts.cjs_require_optimizer {
                 Some(config) => Either::Left(visit_mut_pass(
                     crate::transforms::cjs_optimizer::cjs_optimizer(

crates/next-custom-transforms/src/transforms/mod.rs

@@ -18,6 +18,7 @@ pub mod react_server_components;
 pub mod server_actions;
 pub mod shake_exports;
 pub mod strip_page_exports;
+pub mod track_dynamic_imports;
 pub mod warn_for_edge_runtime;
 
 //[TODO] PACK-1564: need to decide reuse vs. turbopack specific

crates/next-custom-transforms/src/transforms/track_dynamic_imports.rs

@@ -0,0 +1,124 @@
+use swc_core::{
+    common::{source_map::PURE_SP, util::take::Take, Mark, SyntaxContext},
+    ecma::{
+        ast::*,
+        utils::{prepend_stmt, private_ident, quote_ident, quote_str},
+        visit::{noop_visit_mut_type, visit_mut_pass, VisitMut, VisitMutWith},
+    },
+    quote,
+};
+
+pub fn track_dynamic_imports(unresolved_mark: Mark) -> impl VisitMut + Pass {
+    visit_mut_pass(ImportReplacer::new(unresolved_mark))
+}
+
+struct ImportReplacer {
+    unresolved_ctxt: SyntaxContext,
+    has_dynamic_import: bool,
+    wrapper_function_local_ident: Ident,
+}
+
+impl ImportReplacer {
+    pub fn new(unresolved_mark: Mark) -> Self {
+        ImportReplacer {
+            unresolved_ctxt: SyntaxContext::empty().apply_mark(unresolved_mark),
+            has_dynamic_import: false,
+            wrapper_function_local_ident: private_ident!("$$trackDynamicImport__"),
+        }
+    }
+}
+
+impl VisitMut for ImportReplacer {
+    noop_visit_mut_type!();
+
+    fn visit_mut_program(&mut self, program: &mut Program) {
+        program.visit_mut_children_with(self);
+        // if we wrapped a dynamic import while visiting the children, we need to import the wrapper
+
+        if self.has_dynamic_import {
+            let import_args = MakeNamedImportArgs {
+                original_ident: quote_ident!("trackDynamicImport").into(),
+                local_ident: self.wrapper_function_local_ident.clone(),
+                source: "private-next-rsc-track-dynamic-import",
+                unresolved_ctxt: self.unresolved_ctxt,
+            };
+            match program {
+                Program::Module(module) => {
+                    prepend_stmt(&mut module.body, make_named_import_esm(import_args));
+                }
+                Program::Script(script) => {
+                    // CJS modules can still use `import()`. for CJS, we have to inject the helper
+                    // using `require` instead of `import` to avoid accidentally turning them
+                    // into ESM modules.
+                    prepend_stmt(&mut script.body, make_named_import_cjs(import_args));
+                }
+            }
+        }
+    }
+
+    fn visit_mut_expr(&mut self, expr: &mut Expr) {
+        expr.visit_mut_children_with(self);
+
+        // before: `import(...)`
+        // after:  `$$trackDynamicImport__(import(...))`
+
+        if let Expr::Call(CallExpr {
+            callee: Callee::Import(_),
+            ..
+        }) = expr
+        {
+            self.has_dynamic_import = true;
+            let replacement_expr = quote!(
+                "$wrapper_fn($expr)" as Expr,
+                wrapper_fn = self.wrapper_function_local_ident.clone(),
+                expr: Expr = expr.take()
+            )
+            .with_span(PURE_SP);
+            *expr = replacement_expr
+        }
+    }
+}
+
+struct MakeNamedImportArgs<'a> {
+    original_ident: Ident,
+    local_ident: Ident,
+    source: &'a str,
+    unresolved_ctxt: SyntaxContext,
+}
+
+fn make_named_import_esm(args: MakeNamedImportArgs) -> ModuleItem {
+    let MakeNamedImportArgs {
+        original_ident,
+        local_ident,
+        source,
+        ..
+    } = args;
+    let mut item = quote!(
+        "import { $original_ident as $local_ident } from 'dummy'" as ModuleItem,
+        original_ident = original_ident,
+        local_ident = local_ident,
+    );
+    // the import source cannot be parametrized in `quote!()`, so patch it manually
+    let decl = item.as_mut_module_decl().unwrap().as_mut_import().unwrap();
+    decl.src = Box::new(source.into());
+    item
+}
+
+fn make_named_import_cjs(args: MakeNamedImportArgs) -> Stmt {
+    let MakeNamedImportArgs {
+        original_ident,
+        local_ident,
+        source,
+        unresolved_ctxt,
+    } = args;
+    quote!(
+        "const { [$original_name]: $local_ident } = $require($source)" as Stmt,
+        original_name: Expr = quote_str!(original_ident.sym).into(),
+        local_ident = local_ident,
+        source: Expr = quote_str!(source).into(),
+        // the builtin `require` is considered an unresolved identifier.
+        // we have to match that, or it won't be recognized as
+        // a proper `require()` call.
+        require = quote_ident!(unresolved_ctxt, "require")
+    )
+}

crates/next-custom-transforms/tests/fixture.rs

@@ -21,6 +21,7 @@ use next_custom_transforms::transforms::{
     server_actions::{self, server_actions, ServerActionsMode},
     shake_exports::{shake_exports, Config as ShakeExportsConfig},
     strip_page_exports::{next_transform_strip_page_exports, ExportFilter},
+    track_dynamic_imports::track_dynamic_imports,
     warn_for_edge_runtime::warn_for_edge_runtime,
 };
 use rustc_hash::FxHashSet;
@@ -930,6 +931,29 @@ fn test_source_maps(input: PathBuf) {
     );
 }
 
+#[fixture("tests/fixture/track-dynamic-imports/**/input.js")]
+fn track_dynamic_imports_fixture(input: PathBuf) {
+    let output = input.parent().unwrap().join("output.js");
+    test_fixture(
+        syntax(),
+        &|_tr| {
+            let unresolved_mark = Mark::new();
+            let top_level_mark = Mark::new();
+            (
+                resolver(unresolved_mark, top_level_mark, false),
+                track_dynamic_imports(unresolved_mark),
+            )
+        },
+        &input,
+        &output,
+        FixtureTestConfig {
+            // auto detect script/module to test CJS handling
+            module: None,
+            ..Default::default()
+        },
+    );
+}
+
 fn lint_to_fold<R>(r: R) -> impl Pass
 where
     R: Visit,

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/1/input.js

@@ -0,0 +1,4 @@
+export default async function Page() {
+  const { foo } = await import('some-module')
+  return foo()
+}

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/1/output.js

@@ -0,0 +1,5 @@
+import { trackDynamicImport as $$trackDynamicImport__ } from "private-next-rsc-track-dynamic-import";
+export default async function Page() {
+    const { foo } = await /*#__PURE__*/ $$trackDynamicImport__(import('some-module'));
+    return foo();
+}

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/2/input.js

@@ -0,0 +1,4 @@
+export default async function Page() {
+  await import((await import('get-name')).default)
+  return null
+}

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/2/output.js

@@ -0,0 +1,5 @@
+import { trackDynamicImport as $$trackDynamicImport__ } from "private-next-rsc-track-dynamic-import";
+export default async function Page() {
+    await /*#__PURE__*/ $$trackDynamicImport__(import((await /*#__PURE__*/ $$trackDynamicImport__(import('get-name'))).default));
+    return null;
+}

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/3/input.js

@@ -0,0 +1,8 @@
+export default async function Page() {
+  const { foo } = await import('some-module')
+  // name conflict
+  $$trackDynamicImport__()
+  return foo()
+}
+
+export function $$trackDynamicImport__() {}

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/3/output.js

@@ -0,0 +1,9 @@
+import { trackDynamicImport as $$trackDynamicImport__ } from "private-next-rsc-track-dynamic-import";
+export default async function Page() {
+    const { foo } = await /*#__PURE__*/ $$trackDynamicImport__(import('some-module'));
+    // name conflict
+    $$trackDynamicImport__1();
+    return foo();
+}
+function $$trackDynamicImport__1() {}
+export { $$trackDynamicImport__1 as $$trackDynamicImport__ };

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/4/input.js

@@ -0,0 +1,6 @@
+const promise = import('some-module')
+
+export default async function Page() {
+  const { foo } = await promise
+  return foo()
+}

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/4/output.js

@@ -0,0 +1,6 @@
+import { trackDynamicImport as $$trackDynamicImport__ } from "private-next-rsc-track-dynamic-import";
+const promise = /*#__PURE__*/ $$trackDynamicImport__(import('some-module'));
+export default async function Page() {
+    const { foo } = await promise;
+    return foo();
+}

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/5/input.js

@@ -0,0 +1,6 @@
+async function foo() {
+  const { foo } = await import('some-module')
+  return foo()
+}
+
+exports.foo = foo

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/5/output.js

@@ -0,0 +1,6 @@
+const { ["trackDynamicImport"]: $$trackDynamicImport__ } = require("private-next-rsc-track-dynamic-import");
+async function foo() {
+    const { foo } = await /*#__PURE__*/ $$trackDynamicImport__(import('some-module'));
+    return foo();
+}
+exports.foo = foo;

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/index.ts

@@ -0,0 +1,2 @@
+/// <reference types="./next" />
+/// <reference types="./modules" />

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/modules.d.ts

@@ -0,0 +1,7 @@
+declare module 'some-module' {
+  export function foo(): null
+}
+declare module 'get-name' {
+  const name: string
+  export default name
+}

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/next.d.ts

@@ -0,0 +1,3 @@
+declare module 'private-next-rsc-track-dynamic-import' {
+  export function trackDynamicImport<T>(promise: Promise<T>): Promise<T>
+}

crates/next-custom-transforms/tests/fixture/track-dynamic-imports/tsconfig.json

@@ -0,0 +1,23 @@
+{
+  "compilerOptions": {
+    "noEmit": true,
+    "rootDir": ".",
+
+    "allowJs": true,
+    "checkJs": true,
+
+    "lib": ["ESNext", "DOM"],
+    "skipLibCheck": true,
+
+    "strict": true,
+    "jsx": "preserve",
+
+    "target": "ESNext",
+    "esModuleInterop": true,
+    "module": "Preserve",
+    "moduleResolution": "bundler",
+    "moduleDetection": "force"
+  },
+  "files": ["./index.ts"], // loads ambient declarations for modules used in tests
+  "include": ["./**/*/input.js", "./**/*/output.js"]
+}