From 3b814327e2390e6f3dfd4c855d849662ddd52b2a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Markb=C3=A5ge?= <sebastian@calyptus.eu>
Date: Fri, 14 Oct 2022 15:09:33 -0400
Subject: [PATCH] Allow Async Functions to be used in Server Components
 (#25479)

This is a temporary step until we allow Promises everywhere.

Currently this serializes to a Lazy which can then be consumed in this same
slot by the client.
---
 .../src/__tests__/ReactFlightDOM-test.js      | 25 +++++++-------
 .../react-server/src/ReactFlightServer.js     | 33 +++++++++++++++++--
 2 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js
index 4a835834152f3..1d1835ee31db8 100644
--- a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js
+++ b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js
@@ -350,25 +350,19 @@ describe('ReactFlightDOM', () => {
     }
 
     function makeDelayedText() {
-      let error, _resolve, _reject;
+      let _resolve, _reject;
       let promise = new Promise((resolve, reject) => {
         _resolve = () => {
           promise = null;
           resolve();
         };
         _reject = e => {
-          error = e;
           promise = null;
           reject(e);
         };
       });
-      function DelayedText({children}, data) {
-        if (promise) {
-          throw promise;
-        }
-        if (error) {
-          throw error;
-        }
+      async function DelayedText({children}) {
+        await promise;
         return <Text>{children}</Text>;
       }
       return [DelayedText, _resolve, _reject];
@@ -469,7 +463,9 @@ describe('ReactFlightDOM', () => {
       resolveName();
     });
     // Advance time enough to trigger a nested fallback.
-    jest.advanceTimersByTime(500);
+    await act(async () => {
+      jest.advanceTimersByTime(500);
+    });
     expect(container.innerHTML).toBe(
       '<div>:name::avatar:</div>' +
         '<p>(loading sidebar)</p>' +
@@ -482,7 +478,8 @@ describe('ReactFlightDOM', () => {
     const theError = new Error('Game over');
     // Let's *fail* loading games.
     await act(async () => {
-      rejectGames(theError);
+      await rejectGames(theError);
+      await 'the inner async function';
     });
     const expectedGamesValue = __DEV__
       ? '<p>Game over + a dev digest</p>'
@@ -499,7 +496,8 @@ describe('ReactFlightDOM', () => {
 
     // We can now show the sidebar.
     await act(async () => {
-      resolvePhotos();
+      await resolvePhotos();
+      await 'the inner async function';
     });
     expect(container.innerHTML).toBe(
       '<div>:name::avatar:</div>' +
@@ -510,7 +508,8 @@ describe('ReactFlightDOM', () => {
 
     // Show everything.
     await act(async () => {
-      resolvePosts();
+      await resolvePosts();
+      await 'the inner async function';
     });
     expect(container.innerHTML).toBe(
       '<div>:name::avatar:</div>' +
diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
index a4b1d8d7eae8f..b0d0875b41747 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/packages/react-server/src/ReactFlightServer.js
@@ -21,7 +21,9 @@ import type {
   ReactProviderType,
   ServerContextJSONValue,
   Wakeable,
+  Thenable,
 } from 'shared/ReactTypes';
+import type {LazyComponent} from 'react/src/ReactLazy';
 
 import {
   scheduleWork,
@@ -87,6 +89,7 @@ type ReactJSONValue =
 
 export type ReactModel =
   | React$Element<any>
+  | LazyComponent<any, any>
   | string
   | boolean
   | number
@@ -192,6 +195,25 @@ function createRootContext(
 
 const POP = {};
 
+function readThenable<T>(thenable: Thenable<T>): T {
+  if (thenable.status === 'fulfilled') {
+    return thenable.value;
+  } else if (thenable.status === 'rejected') {
+    throw thenable.reason;
+  }
+  throw thenable;
+}
+
+function createLazyWrapperAroundWakeable(wakeable: Wakeable) {
+  trackSuspendedWakeable(wakeable);
+  const lazyType: LazyComponent<any, Thenable<any>> = {
+    $$typeof: REACT_LAZY_TYPE,
+    _payload: (wakeable: any),
+    _init: readThenable,
+  };
+  return lazyType;
+}
+
 function attemptResolveElement(
   type: any,
   key: null | React$Key,
@@ -214,7 +236,15 @@ function attemptResolveElement(
     }
     // This is a server-side component.
     prepareToUseHooksForComponent(prevThenableState);
-    return type(props);
+    const result = type(props);
+    if (
+      typeof result === 'object' &&
+      result !== null &&
+      typeof result.then === 'function'
+    ) {
+      return createLazyWrapperAroundWakeable(result);
+    }
+    return result;
   } else if (typeof type === 'string') {
     // This is a host element. E.g. HTML.
     return [REACT_ELEMENT_TYPE, type, key, props];
@@ -636,7 +666,6 @@ export function resolveModelToJSON(
 
         return serializeByRefID(newTask.id);
       } else {
-        logRecoverableError(request, x);
         // 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
         // once it gets rendered.