[Flight] Assume __turbopack_load_by_url__ returns a cached Promise (facebook#33792)

Author: eps1lon

packages/react-server-dom-turbopack/src/client/ReactFlightClientConfigBundlerTurbopack.js

@@ -132,12 +132,6 @@ export function resolveServerReference<T>(
   return [resolvedModuleData.id, resolvedModuleData.chunks, name];
 }
 
-// The chunk cache contains all the chunks we've preloaded so far.
-// If they're still pending they're a thenable. This map also exists
-// in Turbopack 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();
-
 function requireAsyncModule(id: string): null | Thenable<any> {
   // We've already loaded all the chunks. We can require the module.
   const promise = __turbopack_require__(id);
@@ -165,6 +159,13 @@ function requireAsyncModule(id: string): null | Thenable<any> {
   }
 }
 
+// Turbopack will return cached promises for the same chunk.
+// We still want to keep track of which chunks we have already instrumented
+// and which chunks have already been loaded until Turbopack returns instrumented
+// thenables directly.
+const instrumentedChunks: WeakSet<Thenable<any>> = new WeakSet();
+const loadedChunks: WeakSet<Thenable<any>> = new WeakSet();
+
 function ignoreReject() {
   // We rely on rejected promises to be handled by another listener.
 }
@@ -174,19 +175,19 @@ export function preloadModule<T>(
   metadata: ClientReference<T>,
 ): null | Thenable<any> {
   const chunks = metadata[CHUNKS];
-  const promises = [];
+  const promises: Promise<any>[] = [];
   for (let i = 0; i < chunks.length; i++) {
     const chunkFilename = chunks[i];
-    const entry = chunkCache.get(chunkFilename);
-    if (entry === undefined) {
-      const thenable = loadChunk(chunkFilename);
+    const thenable = loadChunk(chunkFilename);
+    if (!loadedChunks.has(thenable)) {
       promises.push(thenable);
+    }
+
+    if (!instrumentedChunks.has(thenable)) {
       // $FlowFixMe[method-unbinding]
-      const resolve = chunkCache.set.bind(chunkCache, chunkFilename, null);
+      const resolve = loadedChunks.add.bind(loadedChunks, thenable);
       thenable.then(resolve, ignoreReject);
-      chunkCache.set(chunkFilename, thenable);
-    } else if (entry !== null) {
-      promises.push(entry);
+      instrumentedChunks.add(thenable);
     }
   }
   if (isAsyncImport(metadata)) {