From b1a92960301d286917e95766c1016fb2467c215e Mon Sep 17 00:00:00 2001
From: Andrew Clark <git@andrewclark.io>
Date: Sat, 22 Oct 2022 22:52:39 -0400
Subject: [PATCH] Revert yieldy behavior for non-use Suspense (in Flight, too)

Same as #25537 but for Flight.

I was going to wait to do this later because the temporary
implementation of async components uses some of the same code that
non-used wakables do, but it's not so bad. I just had to inline one bit
of code, which we'll remove when we unify the implementation with `use`.
---
 .../react-server/src/ReactFlightServer.js     | 49 +++++++++++++++----
 .../react-server/src/ReactFlightThenable.js   | 29 +++--------
 2 files changed, 45 insertions(+), 33 deletions(-)

diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
index ff19ea3abc839..a50f2113b1230 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/packages/react-server/src/ReactFlightServer.js
@@ -22,6 +22,9 @@ import type {
   ServerContextJSONValue,
   Wakeable,
   Thenable,
+  PendingThenable,
+  FulfilledThenable,
+  RejectedThenable,
 } from 'shared/ReactTypes';
 import type {LazyComponent} from 'react/src/ReactLazy';
 
@@ -64,7 +67,6 @@ import {
   getActiveContext,
   rootContextSnapshot,
 } from './ReactFlightNewContext';
-import {trackSuspendedWakeable} from './ReactFlightThenable';
 
 import {
   REACT_ELEMENT_TYPE,
@@ -212,10 +214,44 @@ function readThenable<T>(thenable: Thenable<T>): T {
 }
 
 function createLazyWrapperAroundWakeable(wakeable: Wakeable) {
-  trackSuspendedWakeable(wakeable);
+  // This is a temporary fork of the `use` implementation until we accept
+  // promises everywhere.
+  const thenable: Thenable<mixed> = (wakeable: any);
+  switch (thenable.status) {
+    case 'fulfilled':
+    case 'rejected':
+      break;
+    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;
+    }
+  }
   const lazyType: LazyComponent<any, Thenable<any>> = {
     $$typeof: REACT_LAZY_TYPE,
-    _payload: (wakeable: any),
+    _payload: thenable,
     _init: readThenable,
   };
   return lazyType;
@@ -806,11 +842,7 @@ export function resolveModelToJSON(
         );
         const ping = newTask.ping;
         x.then(ping, ping);
-
-        const wakeable: Wakeable = x;
-        trackSuspendedWakeable(wakeable);
         newTask.thenableState = getThenableStateAfterSuspending();
-
         return serializeByRefID(newTask.id);
       } else {
         // Something errored. We'll still send everything we have up until this point.
@@ -1134,9 +1166,6 @@ function retryTask(request: Request, task: Task): void {
       // Something suspended again, let's pick it back up later.
       const ping = task.ping;
       x.then(ping, ping);
-
-      const wakeable: Wakeable = x;
-      trackSuspendedWakeable(wakeable);
       task.thenableState = getThenableStateAfterSuspending();
       return;
     } else {
diff --git a/packages/react-server/src/ReactFlightThenable.js b/packages/react-server/src/ReactFlightThenable.js
index ac47b76bc1acb..eaeb635217d68 100644
--- a/packages/react-server/src/ReactFlightThenable.js
+++ b/packages/react-server/src/ReactFlightThenable.js
@@ -14,7 +14,6 @@
 // instead of "Wakeable". Or some other more appropriate name.
 
 import type {
-  Wakeable,
   Thenable,
   PendingThenable,
   FulfilledThenable,
@@ -30,14 +29,12 @@ export function createThenableState(): ThenableState {
   return [];
 }
 
-// TODO: Unify this with trackSuspendedThenable. It needs to support not only
-// `use`, but async components, too.
-export function trackSuspendedWakeable(wakeable: Wakeable) {
-  // If this wakeable isn't already a thenable, turn it into one now. Then,
-  // when we resume the work loop, we can check if its status is
-  // still pending.
-  // TODO: Get rid of the Wakeable type? It's superseded by UntrackedThenable.
-  const thenable: Thenable<mixed> = (wakeable: any);
+export function trackUsedThenable<T>(
+  thenableState: ThenableState,
+  thenable: Thenable<T>,
+  index: number,
+) {
+  thenableState[index] = thenable;
 
   // We use an expando to track the status and result of a thenable so that we
   // can synchronously unwrap the value. Think of this as an extension of the
@@ -84,20 +81,6 @@ export function trackSuspendedWakeable(wakeable: Wakeable) {
   }
 }
 
-export function trackUsedThenable<T>(
-  thenableState: ThenableState,
-  thenable: Thenable<T>,
-  index: number,
-) {
-  // This is only a separate function from trackSuspendedWakeable for symmetry
-  // with Fiber.
-  // TODO: Disallow throwing a thenable directly. It must go through `use` (or
-  // some equivalent for internal Suspense implementations). We can't do this in
-  // Fiber yet because it's a breaking change but we can do it in Server
-  // Components because Server Components aren't released yet.
-  thenableState[index] = thenable;
-}
-
 export function getPreviouslyUsedThenableAtIndex<T>(
   thenableState: ThenableState | null,
   index: number,