[Flight] Add support for Webpack Async Modules (#25138)

This lets you await the result of require(...) which will then
mark the result as async which will then let the client unwrap the Promise
before handing it over in the same way.

packages/react-server-dom-webpack/src/ReactFlightClientWebpackBundlerConfig.js

@@ -7,6 +7,8 @@
  * @flow
  */
 
+import type {Thenable} from 'shared/ReactTypes';
+
 export type WebpackSSRMap = {
   [clientId: string]: {
     [clientExportName: string]: ModuleMetaData,
@@ -19,6 +21,7 @@ export opaque type ModuleMetaData = {
   id: string,
   chunks: Array<string>,
   name: string,
+  async: boolean,
 };
 
 // eslint-disable-next-line no-unused-vars
@@ -29,7 +32,17 @@ export function resolveModuleReference<T>(
   moduleData: ModuleMetaData,
 ): ModuleReference<T> {
   if (bundlerConfig) {
-    return bundlerConfig[moduleData.id][moduleData.name];
+    const resolvedModuleData = bundlerConfig[moduleData.id][moduleData.name];
+    if (moduleData.async) {
+      return {
+        id: resolvedModuleData.id,
+        chunks: resolvedModuleData.chunks,
+        name: resolvedModuleData.name,
+        async: true,
+      };
+    } else {
+      return resolvedModuleData;
+    }
   }
   return moduleData;
 }
@@ -39,39 +52,72 @@ export function resolveModuleReference<T>(
 // in Webpack but unfortunately it's not exposed so we have to
 // replicate it in user space. null means that it has already loaded.
 const chunkCache: Map<string, null | Promise<any> | Error> = new Map();
+const asyncModuleCache: Map<string, Thenable<any>> = new Map();
 
 // Start preloading the modules since we might need them soon.
 // This function doesn't suspend.
 export function preloadModule<T>(moduleData: ModuleReference<T>): void {
   const chunks = moduleData.chunks;
+  const promises = [];
   for (let i = 0; i < chunks.length; i++) {
     const chunkId = chunks[i];
     const entry = chunkCache.get(chunkId);
     if (entry === undefined) {
       const thenable = __webpack_chunk_load__(chunkId);
+      promises.push(thenable);
       const resolve = chunkCache.set.bind(chunkCache, chunkId, null);
       const reject = chunkCache.set.bind(chunkCache, chunkId);
       thenable.then(resolve, reject);
       chunkCache.set(chunkId, thenable);
     }
   }
+  if (moduleData.async) {
+    const modulePromise: any = Promise.all(promises).then(() => {
+      return __webpack_require__(moduleData.id);
+    });
+    modulePromise.then(
+      value => {
+        modulePromise.status = 'fulfilled';
+        modulePromise.value = value;
+      },
+      reason => {
+        modulePromise.status = 'rejected';
+        modulePromise.reason = reason;
+      },
+    );
+    asyncModuleCache.set(moduleData.id, modulePromise);
+  }
 }
 
 // Actually require the module or suspend if it's not yet ready.
 // Increase priority if necessary.
 export function requireModule<T>(moduleData: ModuleReference<T>): T {
-  const chunks = moduleData.chunks;
+  let moduleExports;
-  for (let i = 0; i < chunks.length; i++) {
+  if (moduleData.async) {
-    const chunkId = chunks[i];
+    // We assume that preloadModule has been called before, which
-    const entry = chunkCache.get(chunkId);
+    // should have added something to the module cache.
-    if (entry !== null) {
+    const promise: any = asyncModuleCache.get(moduleData.id);
-      // We assume that preloadModule has been called before.
+    if (promise.status === 'fulfilled') {
-      // So we don't expect to see entry being undefined here, that's an error.
+      moduleExports = promise.value;
-      // Let's throw either an error or the Promise.
+    } else if (promise.status === 'rejected') {
-      throw entry;
+      throw promise.reason;
+    } else {
+      throw promise;
     }
+  } else {
+    const chunks = moduleData.chunks;
+    for (let i = 0; i < chunks.length; i++) {
+      const chunkId = chunks[i];
+      const entry = chunkCache.get(chunkId);
+      if (entry !== null) {
+        // We assume that preloadModule has been called before.
+        // So we don't expect to see entry being undefined here, that's an error.
+        // Let's throw either an error or the Promise.
+        throw entry;
+      }
+    }
+    moduleExports = __webpack_require__(moduleData.id);
   }
-  const moduleExports = __webpack_require__(moduleData.id);
   if (moduleData.name === '*') {
     // This is a placeholder value that represents that the caller imported this
     // as a CommonJS module as is.

packages/react-server-dom-webpack/src/ReactFlightServerWebpackBundlerConfig.js

@@ -20,12 +20,14 @@ export type ModuleReference<T> = {
   $$typeof: Symbol,
   filepath: string,
   name: string,
+  async: boolean,
 };
 
 export type ModuleMetaData = {
   id: string,
   chunks: Array<string>,
   name: string,
+  async: boolean,
 };
 
 export type ModuleKey = string;
@@ -33,7 +35,12 @@ export type ModuleKey = string;
 const MODULE_TAG = Symbol.for('react.module.reference');
 
 export function getModuleKey(reference: ModuleReference<any>): ModuleKey {
-  return reference.filepath + '#' + reference.name;
+  return (
+    reference.filepath +
+    '#' +
+    reference.name +
+    (reference.async ? '#async' : '')
+  );
 }
 
 export function isModuleReference(reference: Object): boolean {
@@ -44,5 +51,16 @@ export function resolveModuleMetaData<T>(
   config: BundlerConfig,
   moduleReference: ModuleReference<T>,
 ): ModuleMetaData {
-  return config[moduleReference.filepath][moduleReference.name];
+  const resolvedModuleData =
+    config[moduleReference.filepath][moduleReference.name];
+  if (moduleReference.async) {
+    return {
+      id: resolvedModuleData.id,
+      chunks: resolvedModuleData.chunks,
+      name: resolvedModuleData.name,
+      async: true,
+    };
+  } else {
+    return resolvedModuleData;
+  }
 }

packages/react-server-dom-webpack/src/ReactFlightWebpackNodeRegister.js

@@ -14,6 +14,8 @@ const Module = require('module');
 
 module.exports = function register() {
   const MODULE_REFERENCE = Symbol.for('react.module.reference');
+  const PROMISE_PROTOTYPE = Promise.prototype;
+
   const proxyHandlers = {
     get: function(target, name, receiver) {
       switch (name) {
@@ -26,6 +28,8 @@ module.exports = function register() {
           return target.filepath;
         case 'name':
           return target.name;
+        case 'async':
+          return target.async;
         // We need to special case this because createElement reads it if we pass this
         // reference.
         case 'defaultProps':
@@ -39,8 +43,33 @@ module.exports = function register() {
             // This a placeholder value that tells the client to conditionally use the
             // whole object or just the default export.
             name: '',
+            async: target.async,
           };
           return true;
+        case 'then':
+          if (!target.async) {
+            // If this module is expected to return a Promise (such as an AsyncModule) then
+            // we should resolve that with a client reference that unwraps the Promise on
+            // the client.
+            const then = function then(resolve, reject) {
+              const moduleReference: {[string]: any} = {
+                $$typeof: MODULE_REFERENCE,
+                filepath: target.filepath,
+                name: '*', // Represents the whole object instead of a particular import.
+                async: true,
+              };
+              return Promise.resolve(
+                resolve(new Proxy(moduleReference, proxyHandlers)),
+              );
+            };
+            // If this is not used as a Promise but is treated as a reference to a `.then`
+            // export then we should treat it as a reference to that name.
+            then.$$typeof = MODULE_REFERENCE;
+            then.filepath = target.filepath;
+            // then.name is conveniently already "then" which is the export name we need.
+            // This will break if it's minified though.
+            return then;
+          }
       }
       let cachedReference = target[name];
       if (!cachedReference) {
@@ -48,10 +77,15 @@ module.exports = function register() {
           $$typeof: MODULE_REFERENCE,
           filepath: target.filepath,
           name: name,
+          async: target.async,
         };
       }
       return cachedReference;
     },
+    getPrototypeOf(target) {
+      // Pretend to be a Promise in case anyone asks.
+      return PROMISE_PROTOTYPE;
+    },
     set: function() {
       throw new Error('Cannot assign to a client module from a server module.');
     },
@@ -63,6 +97,7 @@ module.exports = function register() {
       $$typeof: MODULE_REFERENCE,
       filepath: moduleId,
       name: '*', // Represents the whole object instead of a particular import.
+      async: false,
     };
     module.exports = new Proxy(moduleReference, proxyHandlers);
   };

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js

@@ -237,6 +237,83 @@ describe('ReactFlightDOM', () => {
     expect(container.innerHTML).toBe('<p>@div</p>');
   });
 
+  it('should unwrap async module references', async () => {
+    const AsyncModule = Promise.resolve(function AsyncModule({text}) {
+      return 'Async: ' + text;
+    });
+
+    const AsyncModule2 = Promise.resolve({
+      exportName: 'Module',
+    });
+
+    function Print({response}) {
+      return <p>{response.readRoot()}</p>;
+    }
+
+    function App({response}) {
+      return (
+        <Suspense fallback={<h1>Loading...</h1>}>
+          <Print response={response} />
+        </Suspense>
+      );
+    }
+
+    const AsyncModuleRef = await clientExports(AsyncModule);
+    const AsyncModuleRef2 = await clientExports(AsyncModule2);
+
+    const {writable, readable} = getTestStream();
+    const {pipe} = ReactServerDOMWriter.renderToPipeableStream(
+      <AsyncModuleRef text={AsyncModuleRef2.exportName} />,
+      webpackMap,
+    );
+    pipe(writable);
+    const response = ReactServerDOMReader.createFromReadableStream(readable);
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(<App response={response} />);
+    });
+    expect(container.innerHTML).toBe('<p>Async: Module</p>');
+  });
+
+  it('should be able to import a name called "then"', async () => {
+    const thenExports = {
+      then: function then() {
+        return 'and then';
+      },
+    };
+
+    function Print({response}) {
+      return <p>{response.readRoot()}</p>;
+    }
+
+    function App({response}) {
+      return (
+        <Suspense fallback={<h1>Loading...</h1>}>
+          <Print response={response} />
+        </Suspense>
+      );
+    }
+
+    const ThenRef = clientExports(thenExports).then;
+
+    const {writable, readable} = getTestStream();
+    const {pipe} = ReactServerDOMWriter.renderToPipeableStream(
+      <ThenRef />,
+      webpackMap,
+    );
+    pipe(writable);
+    const response = ReactServerDOMReader.createFromReadableStream(readable);
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(<App response={response} />);
+    });
+    expect(container.innerHTML).toBe('<p>and then</p>');
+  });
+
   it('should progressively reveal server components', async () => {
     let reportedErrors = [];
 

packages/react-server-dom-webpack/src/__tests__/utils/WebpackMock.js

@@ -52,6 +52,19 @@ exports.clientExports = function clientExports(moduleExports) {
       name: '*',
     },
   };
+  if (typeof moduleExports.then === 'function') {
+    moduleExports.then(asyncModuleExports => {
+      for (const name in asyncModuleExports) {
+        webpackMap[path] = {
+          [name]: {
+            id: idx,
+            chunks: [],
+            name: name,
+          },
+        };
+      }
+    });
+  }
   for (const name in moduleExports) {
     webpackMap[path] = {
       [name]: {