Add a module map option to the Webpack Flight Client (#24629)

On the server we have a similar translation map from the file path that the
loader uses to the refer to the original module and to the bundled module ID.

The Flight server is optimized to emit the smallest format for the client.
However during SSR, the same client component might go by a different
module ID since it's a different bundle than the client bundle.

This provides an option to add a translation map from client ID to SSR ID
when reading the Flight stream.

Ideally we should have a special SSR Flight Client that takes this option
but for now we only have one Client for both.

Author: sebmarkbage

packages/react-client/src/ReactFlightClient.js

@@ -15,6 +15,7 @@ import type {
   ModuleMetaData,
   UninitializedModel,
   Response,
+  BundlerConfig,
 } from './ReactFlightClientHostConfig';
 
 import {
@@ -97,6 +98,7 @@ Chunk.prototype.then = function<T>(resolve: () => mixed) {
 };
 
 export type ResponseBase = {
+  _bundlerConfig: BundlerConfig,
   _chunks: Map<number, SomeChunk<any>>,
   readRoot<T>(): T,
   ...
@@ -338,9 +340,10 @@ export function parseModelTuple(
   return value;
 }
 
-export function createResponse(): ResponseBase {
+export function createResponse(bundlerConfig: BundlerConfig): ResponseBase {
   const chunks: Map<number, SomeChunk<any>> = new Map();
   const response = {
+    _bundlerConfig: bundlerConfig,
     _chunks: chunks,
     readRoot: readRoot,
   };
@@ -384,7 +387,10 @@ export function resolveModule(
   const chunks = response._chunks;
   const chunk = chunks.get(id);
   const moduleMetaData: ModuleMetaData = parseModel(response, model);
-  const moduleReference = resolveModuleReference(moduleMetaData);
+  const moduleReference = resolveModuleReference(
+    response._bundlerConfig,
+    moduleMetaData,
+  );
 
   // TODO: Add an option to encode modules that are lazy loaded.
   // For now we preload all modules as early as possible since it's likely

packages/react-client/src/ReactFlightClientStream.js

@@ -9,6 +9,8 @@
 
 import type {Response} from './ReactFlightClientHostConfigStream';
 
+import type {BundlerConfig} from './ReactFlightClientHostConfig';
+
 import {
   resolveModule,
   resolveModel,
@@ -121,11 +123,11 @@ function createFromJSONCallback(response: Response) {
   };
 }
 
-export function createResponse(): Response {
+export function createResponse(bundlerConfig: BundlerConfig): Response {
   // NOTE: CHECK THE COMPILER OUTPUT EACH TIME YOU CHANGE THIS.
   // It should be inlined to one object literal but minor changes can break it.
   const stringDecoder = supportsBinaryStreams ? createStringDecoder() : null;
-  const response: any = createResponseBase();
+  const response: any = createResponseBase(bundlerConfig);
   response._partialRow = '';
   if (supportsBinaryStreams) {
     response._stringDecoder = stringDecoder;

packages/react-client/src/forks/ReactFlightClientHostConfig.custom.js

@@ -26,6 +26,7 @@
 declare var $$$hostConfig: any;
 
 export type Response = any;
+export opaque type BundlerConfig = mixed; // eslint-disable-line no-undef
 export opaque type ModuleMetaData = mixed; // eslint-disable-line no-undef
 export opaque type ModuleReference<T> = mixed; // eslint-disable-line no-undef
 export const resolveModuleReference = $$$hostConfig.resolveModuleReference;

packages/react-noop-renderer/src/ReactNoopFlightClient.js

@@ -22,7 +22,7 @@ type Source = Array<string>;
 
 const {createResponse, processStringChunk, close} = ReactFlightClient({
   supportsBinaryStreams: false,
-  resolveModuleReference(idx: string) {
+  resolveModuleReference(bundlerConfig: null, idx: string) {
     return idx;
   },
   preloadModule(idx: string) {},
@@ -35,7 +35,7 @@ const {createResponse, processStringChunk, close} = ReactFlightClient({
 });
 
 function read<T>(source: Source): T {
-  const response = createResponse(source);
+  const response = createResponse(source, null);
   for (let i = 0; i < source.length; i++) {
     processStringChunk(response, source[i], 0);
   }

packages/react-server-dom-relay/src/ReactFlightDOMRelayClientHostConfig.js

@@ -11,6 +11,8 @@ import type {JSONValue, ResponseBase} from 'react-client/src/ReactFlightClient';
 
 import type {JSResourceReference} from 'JSResourceReference';
 
+import type {ModuleMetaData} from 'ReactFlightDOMRelayClientIntegration';
+
 export type ModuleReference<T> = JSResourceReference<T>;
 
 import {
@@ -19,19 +21,29 @@ import {
 } from 'react-client/src/ReactFlightClient';
 
 export {
-  resolveModuleReference,
   preloadModule,
   requireModule,
 } from 'ReactFlightDOMRelayClientIntegration';
 
+import {resolveModuleReference as resolveModuleReferenceImpl} from 'ReactFlightDOMRelayClientIntegration';
+
 import isArray from 'shared/isArray';
 
 export type {ModuleMetaData} from 'ReactFlightDOMRelayClientIntegration';
 
+export type BundlerConfig = null;
+
 export type UninitializedModel = JSONValue;
 
 export type Response = ResponseBase;
 
+export function resolveModuleReference<T>(
+  bundlerConfig: BundlerConfig,
+  moduleData: ModuleMetaData,
+): ModuleReference<T> {
+  return resolveModuleReferenceImpl(moduleData);
+}
+
 function parseModelRecursively(response: Response, parentObj, value) {
   if (typeof value === 'string') {
     return parseModelString(response, parentObj, value);

packages/react-server-dom-relay/src/__tests__/ReactFlightDOMRelay-test.internal.js

@@ -31,7 +31,7 @@ describe('ReactFlightDOMRelay', () => {
   });
 
   function readThrough(data) {
-    const response = ReactDOMFlightRelayClient.createResponse();
+    const response = ReactDOMFlightRelayClient.createResponse(null);
     for (let i = 0; i < data.length; i++) {
       const chunk = data[i];
       ReactDOMFlightRelayClient.resolveRow(response, chunk);

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

@@ -7,6 +7,14 @@
  * @flow
  */
 
+export type WebpackSSRMap = {
+  [clientId: string]: {
+    [clientExportName: string]: ModuleMetaData,
+  },
+};
+
+export type BundlerConfig = null | WebpackSSRMap;
+
 export opaque type ModuleMetaData = {
   id: string,
   chunks: Array<string>,
@@ -17,8 +25,12 @@ export opaque type ModuleMetaData = {
 export opaque type ModuleReference<T> = ModuleMetaData;
 
 export function resolveModuleReference<T>(
+  bundlerConfig: BundlerConfig,
   moduleData: ModuleMetaData,
 ): ModuleReference<T> {
+  if (bundlerConfig) {
+    return bundlerConfig[moduleData.id][moduleData.name];
+  }
   return moduleData;
 }
 

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

@@ -9,6 +9,8 @@
 
 import type {Response as FlightResponse} from 'react-client/src/ReactFlightClientStream';
 
+import type {BundlerConfig} from './ReactFlightClientWebpackBundlerConfig';
+
 import {
   createResponse,
   reportGlobalError,
@@ -17,6 +19,10 @@ import {
   close,
 } from 'react-client/src/ReactFlightClientStream';
 
+export type Options = {
+  moduleMap?: BundlerConfig,
+};
+
 function startReadingFromStream(
   response: FlightResponse,
   stream: ReadableStream,
@@ -37,16 +43,24 @@ function startReadingFromStream(
   reader.read().then(progress, error);
 }
 
-function createFromReadableStream(stream: ReadableStream): FlightResponse {
-  const response: FlightResponse = createResponse();
+function createFromReadableStream(
+  stream: ReadableStream,
+  options?: Options,
+): FlightResponse {
+  const response: FlightResponse = createResponse(
+    options && options.moduleMap ? options.moduleMap : null,
+  );
   startReadingFromStream(response, stream);
   return response;
 }
 
 function createFromFetch(
   promiseForResponse: Promise<Response>,
+  options?: Options,
 ): FlightResponse {
-  const response: FlightResponse = createResponse();
+  const response: FlightResponse = createResponse(
+    options && options.moduleMap ? options.moduleMap : null,
+  );
   promiseForResponse.then(
     function(r) {
       startReadingFromStream(response, (r.body: any));
@@ -58,8 +72,13 @@ function createFromFetch(
   return response;
 }
 
-function createFromXHR(request: XMLHttpRequest): FlightResponse {
-  const response: FlightResponse = createResponse();
+function createFromXHR(
+  request: XMLHttpRequest,
+  options?: Options,
+): FlightResponse {
+  const response: FlightResponse = createResponse(
+    options && options.moduleMap ? options.moduleMap : null,
+  );
   let processedLength = 0;
   function progress(e: ProgressEvent): void {
     const chunk = request.responseText;

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

@@ -24,6 +24,7 @@ global.__webpack_require__ = function(id) {
 let act;
 let React;
 let ReactDOMClient;
+let ReactDOMServer;
 let ReactServerDOMWriter;
 let ReactServerDOMReader;
 
@@ -35,6 +36,7 @@ describe('ReactFlightDOMBrowser', () => {
     act = require('jest-react').act;
     React = require('react');
     ReactDOMClient = require('react-dom/client');
+    ReactDOMServer = require('react-dom/server.browser');
     ReactServerDOMWriter = require('react-server-dom-webpack/writer.browser.server');
     ReactServerDOMReader = require('react-server-dom-webpack');
   });
@@ -69,6 +71,18 @@ describe('ReactFlightDOMBrowser', () => {
     }
   }
 
+  async function readResult(stream) {
+    const reader = stream.getReader();
+    let result = '';
+    while (true) {
+      const {done, value} = await reader.read();
+      if (done) {
+        return result;
+      }
+      result += Buffer.from(value).toString('utf8');
+    }
+  }
+
   function makeDelayedText(Model) {
     let error, _resolve, _reject;
     let promise = new Promise((resolve, reject) => {
@@ -453,4 +467,49 @@ describe('ReactFlightDOMBrowser', () => {
     // Final pending chunk is written; stream should be closed.
     expect(isDone).toBeTruthy();
   });
+
+  it('should allow an alternative module mapping to be used for SSR', async () => {
+    function ClientComponent() {
+      return <span>Client Component</span>;
+    }
+    // The Client build may not have the same IDs as the Server bundles for the same
+    // component.
+    const ClientComponentOnTheClient = moduleReference(ClientComponent);
+    const ClientComponentOnTheServer = moduleReference(ClientComponent);
+
+    // In the SSR bundle this module won't exist. We simulate this by deleting it.
+    const clientId = webpackMap[ClientComponentOnTheClient.filepath].default.id;
+    delete webpackModules[clientId];
+
+    // Instead, we have to provide a translation from the client meta data to the SSR
+    // meta data.
+    const ssrMetaData = webpackMap[ClientComponentOnTheServer.filepath].default;
+    const translationMap = {
+      [clientId]: {
+        d: ssrMetaData,
+      },
+    };
+
+    function App() {
+      return <ClientComponentOnTheClient />;
+    }
+
+    const stream = ReactServerDOMWriter.renderToReadableStream(
+      <App />,
+      webpackMap,
+    );
+    const response = ReactServerDOMReader.createFromReadableStream(stream, {
+      moduleMap: translationMap,
+    });
+
+    function ClientRoot() {
+      return response.readRoot();
+    }
+
+    const ssrStream = await ReactDOMServer.renderToReadableStream(
+      <ClientRoot />,
+    );
+    const result = await readResult(ssrStream);
+    expect(result).toEqual('<span>Client Component</span>');
+  });
 });

packages/react-server-native-relay/src/ReactFlightNativeRelayClientHostConfig.js

@@ -11,6 +11,8 @@ import type {JSONValue, ResponseBase} from 'react-client/src/ReactFlightClient';
 
 import type {JSResourceReference} from 'JSResourceReference';
 
+import type {ModuleMetaData} from 'ReactFlightNativeRelayClientIntegration';
+
 export type ModuleReference<T> = JSResourceReference<T>;
 
 import {
@@ -19,19 +21,29 @@ import {
 } from 'react-client/src/ReactFlightClient';
 
 export {
-  resolveModuleReference,
   preloadModule,
   requireModule,
 } from 'ReactFlightNativeRelayClientIntegration';
 
+import {resolveModuleReference as resolveModuleReferenceImpl} from 'ReactFlightNativeRelayClientIntegration';
+
 import isArray from 'shared/isArray';
 
 export type {ModuleMetaData} from 'ReactFlightNativeRelayClientIntegration';
 
+export type BundlerConfig = null;
+
 export type UninitializedModel = JSONValue;
 
 export type Response = ResponseBase;
 
+export function resolveModuleReference<T>(
+  bundlerConfig: BundlerConfig,
+  moduleData: ModuleMetaData,
+): ModuleReference<T> {
+  return resolveModuleReferenceImpl(moduleData);
+}
+
 function parseModelRecursively(response: Response, parentObj, value) {
   if (typeof value === 'string') {
     return parseModelString(response, parentObj, value);