Serialize Promises through Flight (#26086)

This lets you pass Promises from server components to client components
and `use()` them there.

We still don't support Promises as children on the client, so we need to
support both. This will be a lot simpler when we remove the need to
encode children as lazy since we don't need the lazy encoding anymore
then.

I noticed that this test failed because we don't synchronously resolve
instrumented Promises if they're lazy. The second fix calls `.then()`
early to ensure that this lazy initialization can happen eagerly. ~It
felt silly to do this with an empty function or something, so I just did
the attachment of ping listeners early here. It's also a little silly
since they will ping the currently running render for no reason if it's
synchronously available.~ EDIT: That didn't work because a ping might
interrupt the current render. Probably need a bigger refactor.

We could add another extension but we've already taken a lot of
liberties with the Promise protocol. At least this is one that doesn't
need extension of the protocol as much. Any sub-class of promises could
do this.

Author: sebmarkbage

packages/react-client/src/ReactFlightClient.js

@@ -493,6 +493,12 @@ export function parseModelString(
         // When passed into React, we'll know how to suspend on this.
         return createLazyChunkWrapper(chunk);
       }
+      case '@': {
+        // Promise
+        const id = parseInt(value.substring(2), 16);
+        const chunk = getChunk(response, id);
+        return chunk;
+      }
       case 'S': {
         return Symbol.for(value.substring(2));
       }

packages/react-reconciler/src/ReactFiberThenable.js

@@ -88,6 +88,9 @@ export function trackUsedThenable<T>(
         // Only instrument the thenable if the status if not defined. If
         // it's defined, but an unknown value, assume it's been instrumented by
         // some custom userspace implementation. We treat it as "pending".
+        // Attach a dummy listener, to ensure that any lazy initialization can
+        // happen. Flight lazily parses JSON when the value is actually awaited.
+        thenable.then(noop, noop);
       } else {
         const pendingThenable: PendingThenable<T> = (thenable: any);
         pendingThenable.status = 'pending';
@@ -107,17 +110,17 @@ export function trackUsedThenable<T>(
             }
           },
         );
+      }
 
-        // Check one more time in case the thenable resolved synchronously
-        switch (thenable.status) {
-          case 'fulfilled': {
-            const fulfilledThenable: FulfilledThenable<T> = (thenable: any);
-            return fulfilledThenable.value;
-          }
-          case 'rejected': {
-            const rejectedThenable: RejectedThenable<T> = (thenable: any);
-            throw rejectedThenable.reason;
-          }
+      // Check one more time in case the thenable resolved synchronously.
+      switch (thenable.status) {
+        case 'fulfilled': {
+          const fulfilledThenable: FulfilledThenable<T> = (thenable: any);
+          return fulfilledThenable.value;
+        }
+        case 'rejected': {
+          const rejectedThenable: RejectedThenable<T> = (thenable: any);
+          throw rejectedThenable.reason;
         }
       }
 

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

@@ -905,4 +905,50 @@ describe('ReactFlightDOM', () => {
 
     expect(reportedErrors).toEqual(['bug in the bundler']);
   });
+
+  // @gate enableUseHook
+  it('should pass a Promise through props and be able use() it on the client', async () => {
+    async function getData() {
+      return 'async hello';
+    }
+
+    function Component({data}) {
+      const text = use(data);
+      return <p>{text}</p>;
+    }
+
+    const ClientComponent = clientExports(Component);
+
+    function ServerComponent() {
+      const data = getData(); // no await here
+      return <ClientComponent data={data} />;
+    }
+
+    function Print({response}) {
+      return use(response);
+    }
+
+    function App({response}) {
+      return (
+        <Suspense fallback={<h1>Loading...</h1>}>
+          <Print response={response} />
+        </Suspense>
+      );
+    }
+
+    const {writable, readable} = getTestStream();
+    const {pipe} = ReactServerDOMWriter.renderToPipeableStream(
+      <ServerComponent />,
+      webpackMap,
+    );
+    pipe(writable);
+    const response = ReactServerDOMReader.createFromReadableStream(readable);
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(<App response={response} />);
+    });
+    expect(container.innerHTML).toBe('<p>async hello</p>');
+  });
 });

packages/react-server/src/ReactFlightServer.js

@@ -216,6 +216,82 @@ const POP = {};
 const jsxPropsParents: WeakMap<any, any> = new WeakMap();
 const jsxChildrenParents: WeakMap<any, any> = new WeakMap();
 
+function serializeThenable(request: Request, thenable: Thenable<any>): number {
+  request.pendingChunks++;
+  const newTask = createTask(
+    request,
+    null,
+    getActiveContext(),
+    request.abortableTasks,
+  );
+
+  switch (thenable.status) {
+    case 'fulfilled': {
+      // We have the resolved value, we can go ahead and schedule it for serialization.
+      newTask.model = thenable.value;
+      pingTask(request, newTask);
+      return newTask.id;
+    }
+    case 'rejected': {
+      const x = thenable.reason;
+      const digest = logRecoverableError(request, x);
+      if (__DEV__) {
+        const {message, stack} = getErrorMessageAndStackDev(x);
+        emitErrorChunkDev(request, newTask.id, digest, message, stack);
+      } else {
+        emitErrorChunkProd(request, newTask.id, digest);
+      }
+      return newTask.id;
+    }
+    default: {
+      if (typeof thenable.status === 'string') {
+        // Only instrument the thenable if the status if not defined. If
+        // it's defined, but an unknown value, assume it's been instrumented by
+        // some custom userspace implementation. We treat it as "pending".
+        break;
+      }
+      const pendingThenable: PendingThenable<mixed> = (thenable: any);
+      pendingThenable.status = 'pending';
+      pendingThenable.then(
+        fulfilledValue => {
+          if (thenable.status === 'pending') {
+            const fulfilledThenable: FulfilledThenable<mixed> = (thenable: any);
+            fulfilledThenable.status = 'fulfilled';
+            fulfilledThenable.value = fulfilledValue;
+          }
+        },
+        (error: mixed) => {
+          if (thenable.status === 'pending') {
+            const rejectedThenable: RejectedThenable<mixed> = (thenable: any);
+            rejectedThenable.status = 'rejected';
+            rejectedThenable.reason = error;
+          }
+        },
+      );
+      break;
+    }
+  }
+
+  thenable.then(
+    value => {
+      newTask.model = value;
+      pingTask(request, newTask);
+    },
+    reason => {
+      // TODO: Is it safe to directly emit these without being inside a retry?
+      const digest = logRecoverableError(request, reason);
+      if (__DEV__) {
+        const {message, stack} = getErrorMessageAndStackDev(reason);
+        emitErrorChunkDev(request, newTask.id, digest, message, stack);
+      } else {
+        emitErrorChunkProd(request, newTask.id, digest);
+      }
+    },
+  );
+
+  return newTask.id;
+}
+
 function readThenable<T>(thenable: Thenable<T>): T {
   if (thenable.status === 'fulfilled') {
     return thenable.value;
@@ -270,6 +346,7 @@ function createLazyWrapperAroundWakeable(wakeable: Wakeable) {
 }
 
 function attemptResolveElement(
+  request: Request,
   type: any,
   key: null | React$Key,
   ref: mixed,
@@ -303,6 +380,14 @@ function attemptResolveElement(
       result !== null &&
       typeof result.then === 'function'
     ) {
+      // When the return value is in children position we can resolve it immediately,
+      // to its value without a wrapper if it's synchronously available.
+      const thenable: Thenable<any> = result;
+      if (thenable.status === 'fulfilled') {
+        return thenable.value;
+      }
+      // TODO: Once we accept Promises as children on the client, we can just return
+      // the thenable here.
       return createLazyWrapperAroundWakeable(result);
     }
     return result;
@@ -331,6 +416,7 @@ function attemptResolveElement(
         const init = type._init;
         const wrappedType = init(payload);
         return attemptResolveElement(
+          request,
           wrappedType,
           key,
           ref,
@@ -345,6 +431,7 @@ function attemptResolveElement(
       }
       case REACT_MEMO_TYPE: {
         return attemptResolveElement(
+          request,
           type.type,
           key,
           ref,
@@ -414,10 +501,14 @@ function serializeByValueID(id: number): string {
   return '$' + id.toString(16);
 }
 
-function serializeByRefID(id: number): string {
+function serializeLazyID(id: number): string {
   return '$L' + id.toString(16);
 }
 
+function serializePromiseID(id: number): string {
+  return '$@' + id.toString(16);
+}
+
 function serializeSymbolReference(name: string): string {
   return '$S' + name;
 }
@@ -442,7 +533,7 @@ function serializeClientReference(
       // knows how to deal with lazy values. This lets us suspend
       // on this component rather than its parent until the code has
       // loaded.
-      return serializeByRefID(existingId);
+      return serializeLazyID(existingId);
     }
     return serializeByValueID(existingId);
   }
@@ -461,7 +552,7 @@ function serializeClientReference(
       // knows how to deal with lazy values. This lets us suspend
       // on this component rather than its parent until the code has
       // loaded.
-      return serializeByRefID(moduleId);
+      return serializeLazyID(moduleId);
     }
     return serializeByValueID(moduleId);
   } catch (x) {
@@ -835,6 +926,7 @@ export function resolveModelToJSON(
           const element: React$Element<any> = (value: any);
           // Attempt to render the Server Component.
           value = attemptResolveElement(
+            request,
             element.type,
             element.key,
             element.ref,
@@ -873,7 +965,7 @@ export function resolveModelToJSON(
         const ping = newTask.ping;
         x.then(ping, ping);
         newTask.thenableState = getThenableStateAfterSuspending();
-        return serializeByRefID(newTask.id);
+        return serializeLazyID(newTask.id);
       } else {
         // Something errored. We'll still send everything we have up until this point.
         // We'll replace this element with a lazy reference that throws on the client
@@ -887,7 +979,7 @@ export function resolveModelToJSON(
         } else {
           emitErrorChunkProd(request, errorId, digest);
         }
-        return serializeByRefID(errorId);
+        return serializeLazyID(errorId);
       }
     }
   }
@@ -899,6 +991,11 @@ export function resolveModelToJSON(
   if (typeof value === 'object') {
     if (isClientReference(value)) {
       return serializeClientReference(request, parent, key, (value: any));
+    } else if (typeof value.then === 'function') {
+      // We assume that any object with a .then property is a "Thenable" type,
+      // or a Promise type. Either of which can be represented by a Promise.
+      const promiseId = serializeThenable(request, (value: any));
+      return serializePromiseID(promiseId);
     } else if ((value: any).$$typeof === REACT_PROVIDER_TYPE) {
       const providerKey = ((value: any): ReactProviderType<any>)._context
         ._globalName;
@@ -1157,6 +1254,7 @@ function retryTask(request: Request, task: Task): void {
       // also suspends.
       task.model = value;
       value = attemptResolveElement(
+        request,
         element.type,
         element.key,
         element.ref,
@@ -1180,6 +1278,7 @@ function retryTask(request: Request, task: Task): void {
         const nextElement: React$Element<any> = (value: any);
         task.model = value;
         value = attemptResolveElement(
+          request,
           nextElement.type,
           nextElement.key,
           nextElement.ref,