Fix flight loader for shared components (#34692)

* fix flight loader for shared components

* add tests

* fix condition

packages/next/build/webpack-config.ts

@@ -1,7 +1,6 @@
 import ReactRefreshWebpackPlugin from 'next/dist/compiled/@next/react-refresh-utils/ReactRefreshWebpackPlugin'
 import chalk from 'next/dist/compiled/chalk'
 import crypto from 'crypto'
-import { stringify } from 'querystring'
 import { webpack } from 'next/dist/compiled/webpack/webpack'
 import type { webpack5 } from 'next/dist/compiled/webpack/webpack'
 import path, { join as pathJoin, relative as relativePath } from 'path'
@@ -1180,10 +1179,11 @@ export default async function getBaseWebpackConfig(
                 ...codeCondition,
                 test: serverComponentsRegex,
                 use: {
-                  loader: `next-flight-server-loader?${stringify({
+                  loader: 'next-flight-server-loader',
+                  options: {
                     client: 1,
-                    pageExtensions: JSON.stringify(rawPageExtensions),
-                  })}`,
+                    pageExtensions: rawPageExtensions,
+                  },
                 },
               },
             ]
@@ -1195,22 +1195,16 @@ export default async function getBaseWebpackConfig(
           ? [
               {
                 ...codeCondition,
-                test: serverComponentsRegex,
                 use: {
-                  loader: `next-flight-server-loader?${stringify({
-                    pageExtensions: JSON.stringify(rawPageExtensions),
-                  })}`,
-                },
-              },
-              {
-                ...codeCondition,
-                test: clientComponentsRegex,
-                use: {
-                  loader: 'next-flight-client-loader',
+                  loader: 'next-flight-server-loader',
+                  options: {
+                    pageExtensions: rawPageExtensions,
+                  },
                 },
               },
               {
-                test: /next[\\/](dist[\\/]client[\\/])?(link|image)/,
+                test: codeCondition.test,
+                resourceQuery: /__sc_client__/,
                 use: {
                   loader: 'next-flight-client-loader',
                 },

packages/next/build/webpack/loaders/next-flight-client-loader.ts

@@ -94,11 +94,8 @@ export default async function transformSource(
   this: any,
   source: string
 ): Promise<string> {
-  const { resourcePath, resourceQuery } = this
+  const { resourcePath } = this
 
-  if (resourceQuery !== '?flight') return source
-
-  let url = resourcePath
   const transformedSource = source
   if (typeof transformedSource !== 'string') {
     throw new Error('Expected source to have been transformed to a string.')
@@ -108,7 +105,7 @@ export default async function transformSource(
   await parseExportNamesInto(resourcePath, transformedSource, names)
 
   // next.js/packages/next/<component>.js
-  if (/[\\/]next[\\/](link|image)\.js$/.test(url)) {
+  if (/[\\/]next[\\/](link|image)\.js$/.test(resourcePath)) {
     names.push('default')
   }
 
@@ -122,7 +119,7 @@ export default async function transformSource(
       newSrc += 'export const ' + name + ' = '
     }
     newSrc += '{ $$typeof: MODULE_REFERENCE, filepath: '
-    newSrc += JSON.stringify(url)
+    newSrc += JSON.stringify(resourcePath)
     newSrc += ', name: '
     newSrc += JSON.stringify(name)
     newSrc += '};\n'

packages/next/build/webpack/loaders/next-flight-server-loader.ts

@@ -5,17 +5,19 @@ import { parse } from '../../swc'
 import { getBaseSWCOptions } from '../../swc/options'
 import { getRawPageExtensions } from '../../utils'
 
-function isClientComponent(importSource: string, pageExtensions: string[]) {
-  return new RegExp(`\\.client(\\.(${pageExtensions.join('|')}))?`).test(
-    importSource
-  )
-}
+const getIsClientComponent =
+  (pageExtensions: string[]) => (importSource: string) => {
+    return new RegExp(`\\.client(\\.(${pageExtensions.join('|')}))?`).test(
+      importSource
+    )
+  }
 
-function isServerComponent(importSource: string, pageExtensions: string[]) {
-  return new RegExp(`\\.server(\\.(${pageExtensions.join('|')}))?`).test(
-    importSource
-  )
-}
+const getIsServerComponent =
+  (pageExtensions: string[]) => (importSource: string) => {
+    return new RegExp(`\\.server(\\.(${pageExtensions.join('|')}))?`).test(
+      importSource
+    )
+  }
 
 function isNextComponent(importSource: string) {
   return (
@@ -31,21 +33,28 @@ export function isImageImport(importSource: string) {
   )
 }
 
-async function parseImportsInfo(
-  resourcePath: string,
-  source: string,
-  imports: Array<string>,
-  isClientCompilation: boolean,
-  pageExtensions: string[]
-): Promise<{
+async function parseImportsInfo({
+  resourcePath,
+  source,
+  imports,
+  isClientCompilation,
+  isServerComponent,
+  isClientComponent,
+}: {
+  resourcePath: string
+  source: string
+  imports: Array<string>
+  isClientCompilation: boolean
+  isServerComponent: (name: string) => boolean
+  isClientComponent: (name: string) => boolean
+}): Promise<{
   source: string
   defaultExportName: string
 }> {
   const opts = getBaseSWCOptions({
     filename: resourcePath,
     globalWindow: isClientCompilation,
   })
-
   const ast = await parse(source, { ...opts.jsc.parser, isModule: true })
   const { body } = ast
   const beginPos = ast.span.start
@@ -58,29 +67,49 @@ async function parseImportsInfo(
       case 'ImportDeclaration': {
         const importSource = node.source.value
         if (!isClientCompilation) {
+          // Server compilation for .server.js.
+          if (isServerComponent(importSource)) {
+            continue
+          }
+
+          const importDeclarations = source.substring(
+            lastIndex,
+            node.source.span.start - beginPos
+          )
+
           if (
             !(
-              isClientComponent(importSource, pageExtensions) ||
+              isClientComponent(importSource) ||
               isNextComponent(importSource) ||
               isImageImport(importSource)
             )
           ) {
-            continue
+            if (
+              ['react/jsx-runtime', 'react/jsx-dev-runtime'].includes(
+                importSource
+              )
+            ) {
+              continue
+            }
+
+            // A shared component. It should be handled as a server
+            // component.
+            transformedSource += importDeclarations
+            transformedSource += JSON.stringify(`${importSource}?__sc_server__`)
+          } else {
+            // A client component. It should be loaded as module reference.
+            transformedSource += importDeclarations
+            transformedSource += JSON.stringify(`${importSource}?__sc_client__`)
+            imports.push(`require(${JSON.stringify(importSource)})`)
           }
-          const importDeclarations = source.substring(
-            lastIndex,
-            node.source.span.start - beginPos
-          )
-          transformedSource += importDeclarations
-          transformedSource += JSON.stringify(`${node.source.value}?flight`)
         } else {
           // For the client compilation, we skip all modules imports but
           // always keep client components in the bundle. All client components
           // have to be imported from either server or client components.
           if (
             !(
-              isClientComponent(importSource, pageExtensions) ||
-              isServerComponent(importSource, pageExtensions) ||
+              isClientComponent(importSource) ||
+              isServerComponent(importSource) ||
               // Special cases for Next.js APIs that are considered as client
               // components:
               isNextComponent(importSource) ||
@@ -89,11 +118,12 @@ async function parseImportsInfo(
           ) {
             continue
           }
+
+          imports.push(`require(${JSON.stringify(importSource)})`)
         }
 
         lastIndex = node.source.span.end - beginPos
-        imports.push(`require(${JSON.stringify(importSource)})`)
-        continue
+        break
       }
       case 'ExportDefaultDeclaration': {
         const def = node.decl
@@ -126,28 +156,44 @@ export default async function transformSource(
   this: any,
   source: string
 ): Promise<string> {
-  const { client: isClientCompilation, pageExtensions: pageExtensionsJson } =
-    this.getOptions()
-  const { resourcePath } = this
-  const pageExtensions = JSON.parse(pageExtensionsJson)
+  const { client: isClientCompilation, pageExtensions } = this.getOptions()
+  const { resourcePath, resourceQuery } = this
 
   if (typeof source !== 'string') {
     throw new Error('Expected source to have been transformed to a string.')
   }
 
+  // We currently assume that all components are shared components (unsuffixed)
+  // from node_modules.
   if (resourcePath.includes('/node_modules/')) {
     return source
   }
 
+  const rawRawPageExtensions = getRawPageExtensions(pageExtensions)
+  const isServerComponent = getIsServerComponent(rawRawPageExtensions)
+  const isClientComponent = getIsClientComponent(rawRawPageExtensions)
+
+  if (!isClientCompilation) {
+    // We only apply the loader to server components, or shared components that
+    // are imported by a server component.
+    if (
+      !isServerComponent(resourcePath) &&
+      resourceQuery !== '?__sc_server__'
+    ) {
+      return source
+    }
+  }
+
   const imports: string[] = []
   const { source: transformedSource, defaultExportName } =
-    await parseImportsInfo(
+    await parseImportsInfo({
       resourcePath,
       source,
       imports,
       isClientCompilation,
-      getRawPageExtensions(pageExtensions)
-    )
+      isServerComponent,
+      isClientComponent,
+    })
 
   /**
    * For .server.js files, we handle this loader differently.
@@ -177,6 +223,5 @@ export default async function transformSource(
   }
 
   const transformed = transformedSource + '\n' + noop + '\n' + defaultExportNoop
-
   return transformed
 }

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

@@ -69,7 +69,7 @@ export class FlightManifestPlugin {
     const { clientComponentsRegex } = this
     compilation.chunkGroups.forEach((chunkGroup: any) => {
       function recordModule(id: string, _chunk: any, mod: any) {
-        const resource = mod.resource?.replace(/\?flight$/, '')
+        const resource = mod.resource?.replace(/\?__sc_client__$/, '')
 
         // TODO: Hook into deps instead of the target module.
         // That way we know by the type of dep whether to include.

test/integration/react-streaming-and-server-components/app/components/client.client.js

@@ -0,0 +1,7 @@
+import { useState } from 'react'
+
+export default function Client() {
+  // To ensure that this component is rendered as a client component, we use a
+  // state here.
+  return useState('client_component')[0]
+}

test/integration/react-streaming-and-server-components/app/components/shared.client.js

@@ -0,0 +1,3 @@
+import Shared from './shared'
+
+export default Shared

test/integration/react-streaming-and-server-components/app/components/shared.js

@@ -0,0 +1,21 @@
+import { useState } from 'react'
+import Client from './client.client'
+
+const random = ~~(Math.random() * 10000)
+
+export default function Shared() {
+  let isServerComponent
+  try {
+    useState()
+    isServerComponent = false
+  } catch (e) {
+    isServerComponent = true
+  }
+
+  return (
+    <>
+      <Client />,{' '}
+      {(isServerComponent ? 'shared:server' : 'shared:client') + ':' + random}
+    </>
+  )
+}

test/integration/react-streaming-and-server-components/app/pages/shared.server.js

@@ -0,0 +1,22 @@
+import ClientFromDirect from '../components/client.client'
+import ClientFromShared from '../components/shared'
+import SharedFromClient from '../components/shared.client'
+
+export default function Page() {
+  // All three client components should be rendered correctly, but only
+  // shared component is a server component, and another is a client component.
+  // These two shared components should be created as two module instances.
+  return (
+    <div id="main">
+      <ClientFromDirect />
+      <br />
+      <ClientFromShared />
+      <br />
+      <ClientFromShared />
+      <br />
+      <SharedFromClient />
+      <br />
+      <SharedFromClient />
+    </div>
+  )
+}

test/integration/react-streaming-and-server-components/test/rsc.js

@@ -55,6 +55,26 @@ export default function (context, { runtime, env }) {
     expect(html).toContain('foo.client')
   })
 
+  it('should resolve different kinds of components correctly', async () => {
+    const html = await renderViaHTTP(context.appPort, '/shared')
+    const main = getNodeBySelector(html, '#main').html()
+
+    // Should have 5 occurrences of "client_component".
+    expect([...main.matchAll(/client_component/g)].length).toBe(5)
+
+    // Should have 2 occurrences of "shared:server", and 2 occurrences of
+    // "shared:client".
+    const sharedServerModule = [...main.matchAll(/shared:server:(\d+)/g)]
+    const sharedClientModule = [...main.matchAll(/shared:client:(\d+)/g)]
+    expect(sharedServerModule.length).toBe(2)
+    expect(sharedClientModule.length).toBe(2)
+
+    // Should have 2 modules created for the shared component.
+    expect(sharedServerModule[0][1]).toBe(sharedServerModule[1][1])
+    expect(sharedClientModule[0][1]).toBe(sharedClientModule[1][1])
+    expect(sharedServerModule[0][1]).not.toBe(sharedClientModule[0][1])
+  })
+
   it('should support next/link in server components', async () => {
     const linkHTML = await renderViaHTTP(context.appPort, '/next-api/link')
     const linkText = getNodeBySelector(

yarn.lock

@@ -4855,7 +4855,7 @@
   version "4.1.5"
   resolved "https://registry.yarnpkg.com/@types/debug/-/debug-4.1.5.tgz#b14efa8852b7768d898906613c23f688713e02cd"
 
-"@types/eslint-scope@^3.7.3":
+"@types/eslint-scope@^3.7.0", "@types/eslint-scope@^3.7.3":
   version "3.7.3"
   resolved "https://registry.yarnpkg.com/@types/eslint-scope/-/eslint-scope-3.7.3.tgz#125b88504b61e3c8bc6f870882003253005c3224"
   integrity sha512-PB3ldyrcnAicT35TWPs5IcwKD8S333HMaa2VVv4+wdvebJkjWuW/xESoB8IwRcog8HYVYamb1g/R31Qv5Bx03g==
@@ -4880,6 +4880,11 @@
   version "0.0.39"
   resolved "https://registry.yarnpkg.com/@types/estree/-/estree-0.0.39.tgz#e177e699ee1b8c22d23174caaa7422644389509f"
 
+"@types/estree@^0.0.50":
+  version "0.0.50"
+  resolved "https://registry.yarnpkg.com/@types/estree/-/estree-0.0.50.tgz#1e0caa9364d3fccd2931c3ed96fdbeaa5d4cca83"
+  integrity sha512-C6N5s2ZFtuZRj54k2/zyRhNDjJwwcViAM3Nbm8zjBpbqAdZ00mr0CFxvSKeO8Y/e03WVFLpQMdHYVfUd6SB+Hw==
+
 "@types/estree@^0.0.51":
   version "0.0.51"
   resolved "https://registry.yarnpkg.com/@types/estree/-/estree-0.0.51.tgz#cfd70924a25a3fd32b218e5e420e6897e1ac4f40"