Instrument the Promise for async modules instead of using a cache

We typically do this for Promises in general. Usually the ones we provide.

This ensures that if Webpack replaces the module (HMR) then we'll get the
new Promise when we require it again.

Author: sebmarkbage

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

@@ -111,7 +111,33 @@ export function resolveServerReference<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>> = new Map();
-const asyncModuleCache: Map<string, Thenable<any>> = new Map();
+
+function requireAsyncModule(id: string): null | Thenable<any> {
+  // We've already loaded all the chunks. We can require the module.
+  const promise = __webpack_require__(id);
+  if (typeof promise.then !== 'function') {
+    // This wasn't a promise after all.
+    return null;
+  } else if (promise.status === 'fulfilled') {
+    // This module was already resolved earlier.
+    return null;
+  } else {
+    // Instrument the Promise to stash the result.
+    promise.then(
+      value => {
+        const fulfilledThenable: FulfilledThenable<mixed> = (promise: any);
+        fulfilledThenable.status = 'fulfilled';
+        fulfilledThenable.value = value;
+      },
+      reason => {
+        const rejectedThenable: RejectedThenable<mixed> = (promise: any);
+        rejectedThenable.status = 'rejected';
+        rejectedThenable.reason = reason;
+      },
+    );
+    return promise;
+  }
+}
 
 function ignoreReject() {
   // We rely on rejected promises to be handled by another listener.
@@ -138,32 +164,12 @@ export function preloadModule<T>(
     }
   }
   if (metadata.async) {
-    const existingPromise = asyncModuleCache.get(metadata.id);
-    if (existingPromise) {
-      if (existingPromise.status === 'fulfilled') {
-        return null;
-      }
-      return existingPromise;
+    if (promises.length === 0) {
+      return requireAsyncModule(metadata.id);
     } else {
-      const modulePromise: Thenable<T> = Promise.all(promises).then(() => {
-        return __webpack_require__(metadata.id);
+      return Promise.all(promises).then(() => {
+        return requireAsyncModule(metadata.id);
       });
-      modulePromise.then(
-        value => {
-          const fulfilledThenable: FulfilledThenable<mixed> =
-            (modulePromise: any);
-          fulfilledThenable.status = 'fulfilled';
-          fulfilledThenable.value = value;
-        },
-        reason => {
-          const rejectedThenable: RejectedThenable<mixed> =
-            (modulePromise: any);
-          rejectedThenable.status = 'rejected';
-          rejectedThenable.reason = reason;
-        },
-      );
-      asyncModuleCache.set(metadata.id, modulePromise);
-      return modulePromise;
     }
   } else if (promises.length > 0) {
     return Promise.all(promises);
@@ -175,18 +181,16 @@ export function preloadModule<T>(
 // Actually require the module or suspend if it's not yet ready.
 // Increase priority if necessary.
 export function requireModule<T>(metadata: ClientReference<T>): T {
-  let moduleExports;
+  let moduleExports = __webpack_require__(metadata.id);
   if (metadata.async) {
-    // We assume that preloadModule has been called before, which
-    // should have added something to the module cache.
-    const promise: any = asyncModuleCache.get(metadata.id);
-    if (promise.status === 'fulfilled') {
-      moduleExports = promise.value;
+    if (typeof moduleExports.then !== 'function') {
+      // This wasn't a promise after all.
+    } else if (moduleExports.status === 'fulfilled') {
+      // This Promise should've been instrumented by preloadModule.
+      moduleExports = moduleExports.value;
     } else {
-      throw promise.reason;
+      throw moduleExports.reason;
     }
-  } else {
-    moduleExports = __webpack_require__(metadata.id);
   }
   if (metadata.name === '*') {
     // This is a placeholder value that represents that the caller imported this