Resolve outlined models async in Reply just like in Flight Client

Author: sebmarkbage

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

@@ -1130,6 +1130,47 @@ describe('ReactFlightDOMBrowser', () => {
     expect(result).toBe('Hello world');
   });
 
+  it('can pass an async server exports that resolves later to an outline object like a Map', async () => {
+    let resolve;
+    const chunkPromise = new Promise(r => (resolve = r));
+
+    function action() {}
+    const serverModule = serverExports(
+      {
+        action: action,
+      },
+      chunkPromise,
+    );
+
+    // Send the action to the client
+    const stream = ReactServerDOMServer.renderToReadableStream(
+      {action: serverModule.action},
+      webpackMap,
+    );
+    const response =
+      await ReactServerDOMClient.createFromReadableStream(stream);
+
+    // Pass the action back to the server inside a Map
+
+    const map = new Map();
+    map.set('action', response.action);
+
+    const body = await ReactServerDOMClient.encodeReply(map);
+    const resultPromise = ReactServerDOMServer.decodeReply(
+      body,
+      webpackServerMap,
+    );
+
+    // We couldn't yet resolve the server reference because we haven't loaded
+    // its chunk yet in the new server instance. We now resolve it which loads
+    // it asynchronously.
+    await resolve();
+
+    const result = await resultPromise;
+    expect(result instanceof Map).toBe(true);
+    expect(result.get('action')).toBe(action);
+  });
+
   it('supports Float hints before the first await in server components in Fiber', async () => {
     function Component() {
       return <p>hello world</p>;

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

@@ -91,6 +91,23 @@ describe('ReactFlightDOMReplyEdge', () => {
     expect(result).toEqual(buffers);
   });
 
+  // @gate enableBinaryFlight
+  it('should be able to serialize a typed array inside a Map', async () => {
+    const array = new Uint8Array([
+      123, 4, 10, 5, 100, 255, 244, 45, 56, 67, 43, 124, 67, 89, 100, 20,
+    ]);
+    const map = new Map();
+    map.set('array', array);
+
+    const body = await ReactServerDOMClient.encodeReply(map);
+    const result = await ReactServerDOMServer.decodeReply(
+      body,
+      webpackServerMap,
+    );
+
+    expect(result.get('array')).toEqual(array);
+  });
+
   // @gate enableBinaryFlight
   it('should be able to serialize a blob', async () => {
     const bytes = new Uint8Array([

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

@@ -11,11 +11,16 @@ const url = require('url');
 const Module = require('module');
 
 let webpackModuleIdx = 0;
+let webpackChunkIdx = 0;
 const webpackServerModules = {};
 const webpackClientModules = {};
 const webpackErroredModules = {};
 const webpackServerMap = {};
 const webpackClientMap = {};
+const webpackChunkMap = {};
+global.__webpack_chunk_load__ = function (id) {
+  return webpackChunkMap[id];
+};
 global.__webpack_require__ = function (id) {
   if (webpackErroredModules[id]) {
     throw webpackErroredModules[id];
@@ -117,13 +122,20 @@ exports.clientExports = function clientExports(
 };
 
 // This tests server to server references. There's another case of client to server references.
-exports.serverExports = function serverExports(moduleExports) {
+exports.serverExports = function serverExports(moduleExports, blockOnChunk) {
   const idx = '' + webpackModuleIdx++;
   webpackServerModules[idx] = moduleExports;
   const path = url.pathToFileURL(idx).href;
+
+  const chunks = [];
+  if (blockOnChunk) {
+    const chunkId = webpackChunkIdx++;
+    webpackChunkMap[chunkId] = blockOnChunk;
+    chunks.push(chunkId);
+  }
   webpackServerMap[path] = {
     id: idx,
-    chunks: [],
+    chunks: chunks,
     name: '*',
   };
   // We only add this if this test is testing ESM compat.

packages/react-server/src/ReactFlightReplyServer.js

@@ -255,7 +255,14 @@ function loadServerReference<T>(
     }
   }
   promise.then(
-    createModelResolver(parentChunk, parentObject, key),
+    createModelResolver(
+      parentChunk,
+      parentObject,
+      key,
+      false,
+      response,
+      createModel,
+    ),
     createModelReject(parentChunk),
   );
   // We need a placeholder value that will be replaced later.
@@ -334,19 +341,31 @@ function createModelResolver<T>(
   chunk: SomeChunk<T>,
   parentObject: Object,
   key: string,
+  cyclic: boolean,
+  response: Response,
+  map: (response: Response, model: any) => T,
 ): (value: any) => void {
   let blocked;
   if (initializingChunkBlockedModel) {
     blocked = initializingChunkBlockedModel;
-    blocked.deps++;
+    if (!cyclic) {
+      blocked.deps++;
+    }
   } else {
     blocked = initializingChunkBlockedModel = {
-      deps: 1,
-      value: null,
+      deps: cyclic ? 0 : 1,
+      value: (null: any),
     };
   }
   return value => {
-    parentObject[key] = value;
+    parentObject[key] = map(response, value);
+
+    // If this is the root object for a model reference, where `blocked.value`
+    // is a stale `null`, the resolved value can be used directly.
+    if (key === '' && blocked.value === null) {
+      blocked.value = parentObject[key];
+    }
+
     blocked.deps--;
     if (blocked.deps === 0) {
       if (chunk.status !== BLOCKED) {
@@ -367,16 +386,61 @@ function createModelReject<T>(chunk: SomeChunk<T>): (error: mixed) => void {
   return (error: mixed) => triggerErrorOnChunk(chunk, error);
 }
 
-function getOutlinedModel(response: Response, id: number): any {
+function getOutlinedModel<T>(
+  response: Response,
+  id: number,
+  parentObject: Object,
+  key: string,
+  map: (response: Response, model: any) => T,
+): T {
   const chunk = getChunk(response, id);
-  if (chunk.status === RESOLVED_MODEL) {
-    initializeModelChunk(chunk);
+  switch (chunk.status) {
+    case RESOLVED_MODEL:
+      initializeModelChunk(chunk);
+      break;
   }
-  if (chunk.status !== INITIALIZED) {
-    // We know that this is emitted earlier so otherwise it's an error.
-    throw chunk.reason;
+  // The status might have changed after initialization.
+  switch (chunk.status) {
+    case INITIALIZED:
+      return map(response, chunk.value);
+    case PENDING:
+    case BLOCKED:
+      const parentChunk = initializingChunk;
+      chunk.then(
+        createModelResolver(
+          parentChunk,
+          parentObject,
+          key,
+          false,
+          response,
+          map,
+        ),
+        createModelReject(parentChunk),
+      );
+      return (null: any);
+    default:
+      throw chunk.reason;
   }
-  return chunk.value;
+}
+
+function createMap(
+  response: Response,
+  model: Array<[any, any]>,
+): Map<any, any> {
+  return new Map(model);
+}
+
+function createSet(response: Response, model: Array<any>): Set<any> {
+  return new Set(model);
+}
+
+function extractIterator(response: Response, model: Array<any>): Iterator<any> {
+  // $FlowFixMe[incompatible-use]: This uses raw Symbols because we're extracting from a native array.
+  return model[Symbol.iterator]();
+}
+
+function createModel(response: Response, model: any): any {
+  return model;
 }
 
 function parseTypedArray(
@@ -402,10 +466,17 @@ function parseTypedArray(
         });
 
   // Since loading the buffer is an async operation we'll be blocking the parent
-  // chunk. TODO: This is not safe if the parent chunk needs a mapper like Map.
+  // chunk.
   const parentChunk = initializingChunk;
   promise.then(
-    createModelResolver(parentChunk, parentObject, parentKey),
+    createModelResolver(
+      parentChunk,
+      parentObject,
+      parentKey,
+      false,
+      response,
+      createModel,
+    ),
     createModelReject(parentChunk),
   );
   return null;
@@ -434,7 +505,7 @@ function parseModelString(
         const id = parseInt(value.slice(2), 16);
         // TODO: Just encode this in the reference inline instead of as a model.
         const metaData: {id: ServerReferenceId, bound: Thenable<Array<any>>} =
-          getOutlinedModel(response, id);
+          getOutlinedModel(response, id, obj, key, createModel);
         return loadServerReference(
           response,
           metaData.id,
@@ -451,14 +522,12 @@ function parseModelString(
       case 'Q': {
         // Map
         const id = parseInt(value.slice(2), 16);
-        const data = getOutlinedModel(response, id);
-        return new Map(data);
+        return getOutlinedModel(response, id, obj, key, createMap);
       }
       case 'W': {
         // Set
         const id = parseInt(value.slice(2), 16);
-        const data = getOutlinedModel(response, id);
-        return new Set(data);
+        return getOutlinedModel(response, id, obj, key, createSet);
       }
       case 'K': {
         // FormData
@@ -480,8 +549,7 @@ function parseModelString(
       case 'i': {
         // Iterator
         const id = parseInt(value.slice(2), 16);
-        const data = getOutlinedModel(response, id);
-        return data[Symbol.iterator]();
+        return getOutlinedModel(response, id, obj, key, extractIterator);
       }
       case 'I': {
         // $Infinity
@@ -563,27 +631,7 @@ function parseModelString(
 
     // We assume that anything else is a reference ID.
     const id = parseInt(value.slice(1), 16);
-    const chunk = getChunk(response, id);
-    switch (chunk.status) {
-      case RESOLVED_MODEL:
-        initializeModelChunk(chunk);
-        break;
-    }
-    // The status might have changed after initialization.
-    switch (chunk.status) {
-      case INITIALIZED:
-        return chunk.value;
-      case PENDING:
-      case BLOCKED:
-        const parentChunk = initializingChunk;
-        chunk.then(
-          createModelResolver(parentChunk, obj, key),
-          createModelReject(parentChunk),
-        );
-        return null;
-      default:
-        throw chunk.reason;
-    }
+    return getOutlinedModel(response, id, obj, key, createModel);
   }
   return value;
 }