From 80f50fb535b273dc0b677947504a7d1da0929e8c Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Fri, 14 Oct 2022 14:32:22 -0400
Subject: [PATCH] Allow Async Functions to be used in Server Components

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      | 22 ++++++-------
 .../react-server/src/ReactFlightServer.js     | 31 +++++++++++++++++--
 2 files changed, 40 insertions(+), 13 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..5e501ce8c49fe 100644
--- a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js
+++ b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOM-test.js
@@ -362,13 +362,8 @@ describe('ReactFlightDOM', () => {
           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 +464,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 +479,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 +497,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 +509,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..10b0faadaf274 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/packages/react-server/src/ReactFlightServer.js
@@ -22,6 +22,7 @@ import type {
   ServerContextJSONValue,
   Wakeable,
 } from 'shared/ReactTypes';
+import type {LazyComponent} from 'react/src/ReactLazy';
 
 import {
   scheduleWork,
@@ -192,6 +193,25 @@ function createRootContext(
 
 const POP = {};
 
+function readWakeable(wakeable: Wakeable) {
+  if (wakeable.status === 'fulfilled') {
+    return wakeable.value;
+  } else if (wakeable.status === 'rejected') {
+    throw wakeable.reason;
+  }
+  throw wakeable;
+}
+
+function createLazyWrapperAroundWakeable(wakeable) {
+  trackSuspendedWakeable(wakeable);
+  const lazyType: LazyComponent<T, Wakeable<T>> = {
+    $$typeof: REACT_LAZY_TYPE,
+    _payload: wakeable,
+    _init: readWakeable,
+  };
+  return lazyType;
+}
+
 function attemptResolveElement(
   type: any,
   key: null | React$Key,
@@ -214,7 +234,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 +664,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.