[Flight] Add support for Webpack Async Modules

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;
-  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;
+  let moduleExports;
+  if (moduleData.async) {
+    // We assume that preloadModule has been called before, which
+    // should have added something to the module cache.
+    const promise: any = asyncModuleCache.get(moduleData.id);
+    if (promise.status === 'fulfilled') {
+      moduleExports = promise.value;
+    } else if (promise.status === 'rejected') {
+      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,20 +20,27 @@ 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;
 
 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,19 +43,49 @@ 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) {
         cachedReference = target[name] = {
           $$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]: {