[Flight] add support for Lazy components in Flight server (#24068)

* [Flight] add support for Lazy components in Flight server

Lazy components suspend until resolved just like in Fizz. Add tests to confirm Lazy works with Shared Components and Client Component references.

* Support Lazy elements

React.Lazy can now return an element instead of a Component. This commit implements support for Lazy elements when server rendering.

* add lazy initialization to resolveModelToJson

adding lazying initialization toResolveModelToJson means we use attemptResolveElement's full logic on whatever the resolved type ends up being. This better aligns handling of misued Lazy types like a lazy element being used as a Component or a lazy Component being used as an element.

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

@@ -138,6 +138,213 @@ describe('ReactFlight', () => {
     expect(ReactNoop).toMatchRenderedOutput(<span>Hello, Seb Smith</span>);
   });
 
+  it('can render a lazy component as a shared component on the server', async () => {
+    function SharedComponent({text}) {
+      return (
+        <div>
+          shared<span>{text}</span>
+        </div>
+      );
+    }
+
+    let load = null;
+    const loadSharedComponent = () => {
+      return new Promise(res => {
+        load = () => res({default: SharedComponent});
+      });
+    };
+
+    const LazySharedComponent = React.lazy(loadSharedComponent);
+
+    function ServerComponent() {
+      return (
+        <React.Suspense fallback={'Loading...'}>
+          <LazySharedComponent text={'a'} />
+        </React.Suspense>
+      );
+    }
+
+    const transport = ReactNoopFlightServer.render(<ServerComponent />);
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      ReactNoop.render(rootModel);
+    });
+    expect(ReactNoop).toMatchRenderedOutput('Loading...');
+    await load();
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      ReactNoop.render(rootModel);
+    });
+    expect(ReactNoop).toMatchRenderedOutput(
+      <div>
+        shared<span>a</span>
+      </div>,
+    );
+  });
+
+  it('errors on a Lazy element being used in Component position', async () => {
+    function SharedComponent({text}) {
+      return (
+        <div>
+          shared<span>{text}</span>
+        </div>
+      );
+    }
+
+    let load = null;
+
+    const LazyElementDisguisedAsComponent = React.lazy(() => {
+      return new Promise(res => {
+        load = () => res({default: <SharedComponent text={'a'} />});
+      });
+    });
+
+    function ServerComponent() {
+      return (
+        <React.Suspense fallback={'Loading...'}>
+          <LazyElementDisguisedAsComponent text={'b'} />
+        </React.Suspense>
+      );
+    }
+
+    const transport = ReactNoopFlightServer.render(<ServerComponent />);
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      ReactNoop.render(rootModel);
+    });
+    expect(ReactNoop).toMatchRenderedOutput('Loading...');
+    spyOnDevAndProd(console, 'error');
+    await load();
+    expect(console.error).toHaveBeenCalledTimes(1);
+  });
+
+  it('can render a lazy element', async () => {
+    function SharedComponent({text}) {
+      return (
+        <div>
+          shared<span>{text}</span>
+        </div>
+      );
+    }
+
+    let load = null;
+
+    const lazySharedElement = React.lazy(() => {
+      return new Promise(res => {
+        load = () => res({default: <SharedComponent text={'a'} />});
+      });
+    });
+
+    function ServerComponent() {
+      return (
+        <React.Suspense fallback={'Loading...'}>
+          {lazySharedElement}
+        </React.Suspense>
+      );
+    }
+
+    const transport = ReactNoopFlightServer.render(<ServerComponent />);
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      ReactNoop.render(rootModel);
+    });
+    expect(ReactNoop).toMatchRenderedOutput('Loading...');
+    await load();
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      ReactNoop.render(rootModel);
+    });
+    expect(ReactNoop).toMatchRenderedOutput(
+      <div>
+        shared<span>a</span>
+      </div>,
+    );
+  });
+
+  it('errors with lazy value in element position that resolves to Component', async () => {
+    function SharedComponent({text}) {
+      return (
+        <div>
+          shared<span>{text}</span>
+        </div>
+      );
+    }
+
+    let load = null;
+
+    const componentDisguisedAsElement = React.lazy(() => {
+      return new Promise(res => {
+        load = () => res({default: SharedComponent});
+      });
+    });
+
+    function ServerComponent() {
+      return (
+        <React.Suspense fallback={'Loading...'}>
+          {componentDisguisedAsElement}
+        </React.Suspense>
+      );
+    }
+
+    const transport = ReactNoopFlightServer.render(<ServerComponent />);
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      ReactNoop.render(rootModel);
+    });
+    expect(ReactNoop).toMatchRenderedOutput('Loading...');
+    spyOnDevAndProd(console, 'error');
+    await load();
+    expect(console.error).toHaveBeenCalledTimes(1);
+  });
+
+  it('can render a lazy module reference', async () => {
+    function ClientComponent() {
+      return <div>I am client</div>;
+    }
+
+    const ClientComponentReference = moduleReference(ClientComponent);
+
+    let load = null;
+    const loadClientComponentReference = () => {
+      return new Promise(res => {
+        load = () => res({default: ClientComponentReference});
+      });
+    };
+
+    const LazyClientComponentReference = React.lazy(
+      loadClientComponentReference,
+    );
+
+    function ServerComponent() {
+      return (
+        <React.Suspense fallback={'Loading...'}>
+          <LazyClientComponentReference />
+        </React.Suspense>
+      );
+    }
+
+    const transport = ReactNoopFlightServer.render(<ServerComponent />);
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      ReactNoop.render(rootModel);
+    });
+    expect(ReactNoop).toMatchRenderedOutput('Loading...');
+    await load();
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      ReactNoop.render(rootModel);
+    });
+    expect(ReactNoop).toMatchRenderedOutput(<div>I am client</div>);
+  });
+
   it('should error if a non-serializable value is passed to a host component', () => {
     function EventHandlerProp() {
       return (

packages/react-server/src/ReactFlightServer.js

@@ -200,6 +200,12 @@ function attemptResolveElement(
       return [REACT_ELEMENT_TYPE, type, key, props];
     }
     switch (type.$$typeof) {
+      case REACT_LAZY_TYPE: {
+        const payload = type._payload;
+        const init = type._init;
+        const wrappedType = init(payload);
+        return attemptResolveElement(wrappedType, key, ref, props);
+      }
       case REACT_FORWARD_REF_TYPE: {
         const render = type.render;
         return render(props, undefined);
@@ -452,10 +458,6 @@ export function resolveModelToJSON(
   switch (value) {
     case REACT_ELEMENT_TYPE:
       return '$';
-    case REACT_LAZY_TYPE:
-      throw new Error(
-        'React Lazy Components are not yet supported on the server.',
-      );
   }
 
   if (__DEV__) {
@@ -477,23 +479,36 @@ export function resolveModelToJSON(
   while (
     typeof value === 'object' &&
     value !== null &&
-    (value: any).$$typeof === REACT_ELEMENT_TYPE
+    ((value: any).$$typeof === REACT_ELEMENT_TYPE ||
+      (value: any).$$typeof === REACT_LAZY_TYPE)
   ) {
     if (__DEV__) {
       if (isInsideContextValue) {
         console.error('React elements are not allowed in ServerContext');
       }
     }
-    // TODO: Concatenate keys of parents onto children.
-    const element: React$Element<any> = (value: any);
+
     try {
-      // Attempt to render the server component.
-      value = attemptResolveElement(
-        element.type,
-        element.key,
-        element.ref,
-        element.props,
-      );
+      switch ((value: any).$$typeof) {
+        case REACT_ELEMENT_TYPE: {
+          // TODO: Concatenate keys of parents onto children.
+          const element: React$Element<any> = (value: any);
+          // Attempt to render the server component.
+          value = attemptResolveElement(
+            element.type,
+            element.key,
+            element.ref,
+            element.props,
+          );
+          break;
+        }
+        case REACT_LAZY_TYPE: {
+          const payload = (value: any)._payload;
+          const init = (value: any)._init;
+          value = init(payload);
+          break;
+        }
+      }
     } catch (x) {
       if (typeof x === 'object' && x !== null && typeof x.then === 'function') {
         // Something suspended, we'll need to create a new segment and resolve it later.