[Flight] Add support for Module References in transport protocol (#20121)

* Refactor Flight to require a module reference to be brand checked

This exposes a host environment (bundler) specific hook to check if an
object is a module reference. This will be used so that they can be passed
directly into Flight without needing additional wrapper objects.

* Emit module references as a special type of value

We already have JSON and errors as special types of "rows". This encodes
module references as a special type of row value. This was always the
intention because it allows those values to be emitted first in the stream
so that as a large models stream down, we can start preloading as early
as possible.

We preload the module when they resolve but we lazily require them as they
are referenced.

* Emit module references where ever they occur

This emits module references where ever they occur. In blocks or even
directly in elements.

* Don't special case the root row

I originally did this so that a simple stream is also just plain JSON.

However, since we might want to emit things like modules before the root
module in the stream, this gets unnecessarily complicated. We could add
this back as a special case if it's the first byte written but meh.

* Update the protocol

* Add test for using a module reference as a client component

* Relax element type check

Since Flight now accepts a module reference as returned by any bundler
system, depending on the renderer running. We need to drastically relax
the check to include all of them. We can add more as we discover them.

* Move flow annotation

Seems like our compiler is not happy with stripping this.

* Some bookkeeping bug

* Can't use the private field to check

Author: sebmarkbage

packages/react-client/src/ReactFlightClient.js

@@ -41,8 +41,9 @@ export type JSONValue =
 
 const PENDING = 0;
 const RESOLVED_MODEL = 1;
-const INITIALIZED = 2;
-const ERRORED = 3;
+const RESOLVED_MODULE = 2;
+const INITIALIZED = 3;
+const ERRORED = 4;
 
 type PendingChunk = {
   _status: 0,
@@ -56,21 +57,28 @@ type ResolvedModelChunk = {
   _response: Response,
   then(resolve: () => mixed): void,
 };
-type InitializedChunk<T> = {
+type ResolvedModuleChunk<T> = {
   _status: 2,
+  _value: ModuleReference<T>,
+  _response: Response,
+  then(resolve: () => mixed): void,
+};
+type InitializedChunk<T> = {
+  _status: 3,
   _value: T,
   _response: Response,
   then(resolve: () => mixed): void,
 };
 type ErroredChunk = {
-  _status: 3,
+  _status: 4,
   _value: Error,
   _response: Response,
   then(resolve: () => mixed): void,
 };
 type SomeChunk<T> =
   | PendingChunk
   | ResolvedModelChunk
+  | ResolvedModuleChunk<T>
   | InitializedChunk<T>
   | ErroredChunk;
 
@@ -105,6 +113,8 @@ function readChunk<T>(chunk: SomeChunk<T>): T {
       return chunk._value;
     case RESOLVED_MODEL:
       return initializeModelChunk(chunk);
+    case RESOLVED_MODULE:
+      return initializeModuleChunk(chunk);
     case PENDING:
       // eslint-disable-next-line no-throw-literal
       throw (chunk: Wakeable);
@@ -155,6 +165,13 @@ function createResolvedModelChunk(
   return new Chunk(RESOLVED_MODEL, value, response);
 }
 
+function createResolvedModuleChunk<T>(
+  response: Response,
+  value: ModuleReference<T>,
+): ResolvedModuleChunk<T> {
+  return new Chunk(RESOLVED_MODULE, value, response);
+}
+
 function resolveModelChunk<T>(
   chunk: SomeChunk<T>,
   value: UninitializedModel,
@@ -170,6 +187,21 @@ function resolveModelChunk<T>(
   wakeChunk(listeners);
 }
 
+function resolveModuleChunk<T>(
+  chunk: SomeChunk<T>,
+  value: ModuleReference<T>,
+): void {
+  if (chunk._status !== PENDING) {
+    // We already resolved. We didn't expect to see this.
+    return;
+  }
+  const listeners = chunk._value;
+  const resolvedChunk: ResolvedModuleChunk<T> = (chunk: any);
+  resolvedChunk._status = RESOLVED_MODULE;
+  resolvedChunk._value = value;
+  wakeChunk(listeners);
+}
+
 function initializeModelChunk<T>(chunk: ResolvedModelChunk): T {
   const value: T = parseModel(chunk._response, chunk._value);
   const initializedChunk: InitializedChunk<T> = (chunk: any);
@@ -178,6 +210,14 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk): T {
   return value;
 }
 
+function initializeModuleChunk<T>(chunk: ResolvedModuleChunk<T>): T {
+  const value: T = requireModule(chunk._value);
+  const initializedChunk: InitializedChunk<T> = (chunk: any);
+  initializedChunk._status = INITIALIZED;
+  initializedChunk._value = value;
+  return value;
+}
+
 // Report that any missing chunks in the model is now going to throw this
 // error upon read. Also notify any pending promises.
 export function reportGlobalError(response: Response, error: Error): void {
@@ -241,7 +281,7 @@ function createElement(type, key, props): React$Element<any> {
 
 type UninitializedBlockPayload<Data> = [
   mixed,
-  ModuleMetaData | SomeChunk<ModuleMetaData>,
+  BlockRenderFunction<any, Data> | SomeChunk<BlockRenderFunction<any, Data>>,
   Data | SomeChunk<Data>,
   Response,
 ];
@@ -250,14 +290,7 @@ function initializeBlock<Props, Data>(
   tuple: UninitializedBlockPayload<Data>,
 ): BlockComponent<Props, Data> {
   // Require module first and then data. The ordering matters.
-  const moduleMetaData: ModuleMetaData = readMaybeChunk(tuple[1]);
-  const moduleReference: ModuleReference<
-    BlockRenderFunction<Props, Data>,
-  > = resolveModuleReference(moduleMetaData);
-  // TODO: Do this earlier, as the chunk is resolved.
-  preloadModule(moduleReference);
-
-  const moduleExport = requireModule(moduleReference);
+  const moduleExport = readMaybeChunk(tuple[1]);
 
   // The ordering here is important because this call might suspend.
   // We don't want that to prevent the module graph for being initialized.
@@ -363,6 +396,28 @@ export function resolveModel(
   }
 }
 
+export function resolveModule(
+  response: Response,
+  id: number,
+  model: UninitializedModel,
+): void {
+  const chunks = response._chunks;
+  const chunk = chunks.get(id);
+  const moduleMetaData: ModuleMetaData = parseModel(response, model);
+  const moduleReference = resolveModuleReference(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
+  // that we'll need them.
+  preloadModule(moduleReference);
+
+  if (!chunk) {
+    chunks.set(id, createResolvedModuleChunk(response, moduleReference));
+  } else {
+    resolveModuleChunk(chunk, moduleReference);
+  }
+}
+
 export function resolveError(
   response: Response,
   id: number,

packages/react-client/src/ReactFlightClientStream.js

@@ -10,6 +10,7 @@
 import type {Response} from './ReactFlightClientHostConfigStream';
 
 import {
+  resolveModule,
   resolveModel,
   resolveError,
   createResponse as createResponseBase,
@@ -39,6 +40,13 @@ function processFullRow(response: Response, row: string): void {
       resolveModel(response, id, json);
       return;
     }
+    case 'M': {
+      const colon = row.indexOf(':', 1);
+      const id = parseInt(row.substring(1, colon), 16);
+      const json = row.substring(colon + 1);
+      resolveModule(response, id, json);
+      return;
+    }
     case 'E': {
       const colon = row.indexOf(':', 1);
       const id = parseInt(row.substring(1, colon), 16);
@@ -48,9 +56,9 @@ function processFullRow(response: Response, row: string): void {
       return;
     }
     default: {
-      // Assume this is the root model.
-      resolveModel(response, 0, row);
-      return;
+      throw new Error(
+        "Error parsing the data. It's probably an error code or network corruption.",
+      );
     }
   }
 }

packages/react-client/src/__tests__/ReactFlight-test.js

@@ -53,17 +53,29 @@ describe('ReactFlight', () => {
     };
   });
 
+  function moduleReference(value) {
+    return {
+      $$typeof: Symbol.for('react.module.reference'),
+      value: value,
+    };
+  }
+
   function block(render, load) {
     if (load === undefined) {
       return () => {
-        return ReactNoopFlightServerRuntime.serverBlockNoData(render);
+        return ReactNoopFlightServerRuntime.serverBlockNoData(
+          moduleReference(render),
+        );
       };
     }
     return function(...args) {
       const curriedLoad = () => {
         return load(...args);
       };
-      return ReactNoopFlightServerRuntime.serverBlock(render, curriedLoad);
+      return ReactNoopFlightServerRuntime.serverBlock(
+        moduleReference(render),
+        curriedLoad,
+      );
     };
   }
 
@@ -97,6 +109,35 @@ describe('ReactFlight', () => {
     });
   });
 
+  it('can render a client component using a module reference and render there', () => {
+    function UserClient(props) {
+      return (
+        <span>
+          {props.greeting}, {props.name}
+        </span>
+      );
+    }
+    const User = moduleReference(UserClient);
+
+    function Greeting({firstName, lastName}) {
+      return <User greeting="Hello" name={firstName + ' ' + lastName} />;
+    }
+
+    const model = {
+      greeting: <Greeting firstName="Seb" lastName="Smith" />,
+    };
+
+    const transport = ReactNoopFlightServer.render(model);
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      const greeting = rootModel.greeting;
+      ReactNoop.render(greeting);
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(<span>Hello, Seb Smith</span>);
+  });
+
   if (ReactFeatureFlags.enableBlocksAPI) {
     it('can transfer a Block to the client and render there, without data', () => {
       function User(props, data) {

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

@@ -42,8 +42,14 @@ const ReactNoopFlightServer = ReactFlightServer({
   formatChunk(type: string, props: Object): Uint8Array {
     return Buffer.from(JSON.stringify({type, props}), 'utf8');
   },
-  resolveModuleMetaData(config: void, renderFn: Function) {
-    return saveModule(renderFn);
+  isModuleReference(reference: Object): boolean {
+    return reference.$$typeof === Symbol.for('react.module.reference');
+  },
+  resolveModuleMetaData(
+    config: void,
+    reference: {$$typeof: Symbol, value: any},
+  ) {
+    return saveModule(reference.value);
   },
 });