[Flight] Implement FlightClient in terms of Thenable/Promises instead of throwing Promises (#25260)

* [Flight] Align Chunks with Thenable used with experimental_use

Use the field names used by the Thenable data structure passed to use().
These are considered public in this model.

This adds another field since we use a separate field name for "reason".

* Implement Thenable Protocol on Chunks

This doesn't just ping but resolves/rejects with the value.

* Subclass Promises

* Pass key through JSON parsing

* Wait for preloadModules before resolving module chunks

* Initialize lazy resolved values before reading the result

* Block a model from initializing if its direct dependencies are pending

If a module is blocked, then we can't complete initializing a model.
However, we can still let it parse, and then fill in the missing pieces
later.

We need to block it from resolving until all dependencies have filled in
which we can do with a ref count.

* Treat blocked modules or models as a special status

We currently loop over all chunks at the end to error them if they're
still pending. We shouldn't do this if they're pending because they're
blocked on an external resource like a module because the module might not
resolve before the Flight connection closes and that's not an error.

In an alternative solution I had a set that tracked pending chunks and
removed one at a time. While the loop at the end is faster it's more
work as we go.

I figured the extra status might also help debugging.

For modules we can probably assume no forward references, and the first
async module we can just use the promise as the chunk.

So we could probably get away with this only on models that are blocked by
modules.

Author: sebmarkbage

packages/react-client/src/ReactFlightClient.js

@@ -114,7 +114,7 @@ function createFromJSONCallback(response: Response) {
   return function(key: string, value: JSONValue) {
     if (typeof value === 'string') {
       // We can't use .bind here because we need the "this" value.
-      return parseModelString(response, this, value);
+      return parseModelString(response, this, key, value);
     }
     if (typeof value === 'object' && value !== null) {
       return parseModelTuple(response, value);

packages/react-client/src/ReactFlightClientStream.js

@@ -61,10 +61,6 @@ export function trackSuspendedWakeable(wakeable: Wakeable) {
   // If the thenable doesn't have a status, set it to "pending" and attach
   // a listener that will update its status and result when it resolves.
   switch (thenable.status) {
-    case 'pending':
-      // Since the status is already "pending", we can assume it will be updated
-      // when it resolves, either by React or something in userspace.
-      break;
     case 'fulfilled':
     case 'rejected':
       // A thenable that already resolved shouldn't have been thrown, so this is
@@ -75,9 +71,12 @@ export function trackSuspendedWakeable(wakeable: Wakeable) {
       suspendedThenable = null;
       break;
     default: {
-      // TODO: 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.
+      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(

packages/react-reconciler/src/ReactFiberWakeable.new.js

@@ -61,10 +61,6 @@ export function trackSuspendedWakeable(wakeable: Wakeable) {
   // If the thenable doesn't have a status, set it to "pending" and attach
   // a listener that will update its status and result when it resolves.
   switch (thenable.status) {
-    case 'pending':
-      // Since the status is already "pending", we can assume it will be updated
-      // when it resolves, either by React or something in userspace.
-      break;
     case 'fulfilled':
     case 'rejected':
       // A thenable that already resolved shouldn't have been thrown, so this is
@@ -75,9 +71,12 @@ export function trackSuspendedWakeable(wakeable: Wakeable) {
       suspendedThenable = null;
       break;
     default: {
-      // TODO: 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.
+      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(

packages/react-reconciler/src/ReactFiberWakeable.old.js

@@ -44,9 +44,9 @@ export function resolveModuleReference<T>(
   return resolveModuleReferenceImpl(moduleData);
 }
 
-function parseModelRecursively(response: Response, parentObj, value) {
+function parseModelRecursively(response: Response, parentObj, key, value) {
   if (typeof value === 'string') {
-    return parseModelString(response, parentObj, value);
+    return parseModelString(response, parentObj, key, value);
   }
   if (typeof value === 'object' && value !== null) {
     if (isArray(value)) {
@@ -55,6 +55,7 @@ function parseModelRecursively(response: Response, parentObj, value) {
         (parsedValue: any)[i] = parseModelRecursively(
           response,
           value,
+          '' + i,
           value[i],
         );
       }
@@ -65,6 +66,7 @@ function parseModelRecursively(response: Response, parentObj, value) {
         (parsedValue: any)[innerKey] = parseModelRecursively(
           response,
           value,
+          innerKey,
           value[innerKey],
         );
       }
@@ -77,5 +79,5 @@ function parseModelRecursively(response: Response, parentObj, value) {
 const dummy = {};
 
 export function parseModel<T>(response: Response, json: UninitializedModel): T {
-  return (parseModelRecursively(response, dummy, json): any);
+  return (parseModelRecursively(response, dummy, '', json): any);
 }

packages/react-server-dom-relay/src/ReactFlightDOMRelayClientHostConfig.js

@@ -7,7 +7,11 @@
  * @flow
  */
 
-import type {Thenable} from 'shared/ReactTypes';
+import type {
+  Thenable,
+  FulfilledThenable,
+  RejectedThenable,
+} from 'shared/ReactTypes';
 
 export type WebpackSSRMap = {
   [clientId: string]: {
@@ -56,7 +60,9 @@ const asyncModuleCache: Map<string, Thenable<any>> = new Map();
 
 // Start preloading the modules since we might need them soon.
 // This function doesn't suspend.
-export function preloadModule<T>(moduleData: ModuleReference<T>): void {
+export function preloadModule<T>(
+  moduleData: ModuleReference<T>,
+): null | Thenable<any> {
   const chunks = moduleData.chunks;
   const promises = [];
   for (let i = 0; i < chunks.length; i++) {
@@ -72,20 +78,35 @@ export function preloadModule<T>(moduleData: ModuleReference<T>): void {
     }
   }
   if (moduleData.async) {
-    const modulePromise: any = Promise.all(promises).then(() => {
-      return __webpack_require__(moduleData.id);
-    });
-    modulePromise.then(
-      value => {
-        modulePromise.status = 'fulfilled';
-        modulePromise.value = value;
-      },
-      reason => {
-        modulePromise.status = 'rejected';
-        modulePromise.reason = reason;
-      },
-    );
-    asyncModuleCache.set(moduleData.id, modulePromise);
+    const existingPromise = asyncModuleCache.get(moduleData.id);
+    if (existingPromise) {
+      if (existingPromise.status === 'fulfilled') {
+        return null;
+      }
+      return existingPromise;
+    } else {
+      const modulePromise: Thenable<T> = Promise.all(promises).then(() => {
+        return __webpack_require__(moduleData.id);
+      });
+      modulePromise.then(
+        value => {
+          const fulfilledThenable: FulfilledThenable<mixed> = (modulePromise: any);
+          fulfilledThenable.status = 'fulfilled';
+          fulfilledThenable.value = value;
+        },
+        reason => {
+          const rejectedThenable: RejectedThenable<mixed> = (modulePromise: any);
+          rejectedThenable.status = 'rejected';
+          rejectedThenable.reason = reason;
+        },
+      );
+      asyncModuleCache.set(moduleData.id, modulePromise);
+      return modulePromise;
+    }
+  } else if (promises.length > 0) {
+    return Promise.all(promises);
+  } else {
+    return null;
   }
 }
 
@@ -99,23 +120,10 @@ export function requireModule<T>(moduleData: ModuleReference<T>): T {
     const promise: any = asyncModuleCache.get(moduleData.id);
     if (promise.status === 'fulfilled') {
       moduleExports = promise.value;
-    } else if (promise.status === 'rejected') {
-      throw promise.reason;
     } else {
-      throw promise;
+      throw promise.reason;
     }
   } else {
-    const chunks = moduleData.chunks;
-    for (let i = 0; i < chunks.length; i++) {
-      const chunkId = chunks[i];
-      const entry = chunkCache.get(chunkId);
-      if (entry !== null) {
-        // We assume that preloadModule has been called before.
-        // So we don't expect to see entry being undefined here, that's an error.
-        // Let's throw either an error or the Promise.
-        throw entry;
-      }
-    }
     moduleExports = __webpack_require__(moduleData.id);
   }
   if (moduleData.name === '*') {

packages/react-server-dom-webpack/src/ReactFlightClientWebpackBundlerConfig.js

@@ -44,9 +44,9 @@ export function resolveModuleReference<T>(
   return resolveModuleReferenceImpl(moduleData);
 }
 
-function parseModelRecursively(response: Response, parentObj, value) {
+function parseModelRecursively(response: Response, parentObj, key, value) {
   if (typeof value === 'string') {
-    return parseModelString(response, parentObj, value);
+    return parseModelString(response, parentObj, key, value);
   }
   if (typeof value === 'object' && value !== null) {
     if (isArray(value)) {
@@ -55,6 +55,7 @@ function parseModelRecursively(response: Response, parentObj, value) {
         (parsedValue: any)[i] = parseModelRecursively(
           response,
           value,
+          '' + i,
           value[i],
         );
       }
@@ -65,6 +66,7 @@ function parseModelRecursively(response: Response, parentObj, value) {
         (parsedValue: any)[innerKey] = parseModelRecursively(
           response,
           value,
+          innerKey,
           value[innerKey],
         );
       }
@@ -77,5 +79,5 @@ function parseModelRecursively(response: Response, parentObj, value) {
 const dummy = {};
 
 export function parseModel<T>(response: Response, json: UninitializedModel): T {
-  return (parseModelRecursively(response, dummy, json): any);
+  return (parseModelRecursively(response, dummy, '', json): any);
 }

packages/react-server-native-relay/src/ReactFlightNativeRelayClientHostConfig.js

@@ -44,10 +44,6 @@ export function trackSuspendedWakeable(wakeable: Wakeable) {
   // If the thenable doesn't have a status, set it to "pending" and attach
   // a listener that will update its status and result when it resolves.
   switch (thenable.status) {
-    case 'pending':
-      // Since the status is already "pending", we can assume it will be updated
-      // when it resolves, either by React or something in userspace.
-      break;
     case 'fulfilled':
     case 'rejected':
       // A thenable that already resolved shouldn't have been thrown, so this is
@@ -57,9 +53,12 @@ export function trackSuspendedWakeable(wakeable: Wakeable) {
       // TODO: Log a warning?
       break;
     default: {
-      // TODO: 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.
+      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(

packages/react-server/src/ReactFizzWakeable.js

@@ -44,10 +44,6 @@ export function trackSuspendedWakeable(wakeable: Wakeable) {
   // If the thenable doesn't have a status, set it to "pending" and attach
   // a listener that will update its status and result when it resolves.
   switch (thenable.status) {
-    case 'pending':
-      // Since the status is already "pending", we can assume it will be updated
-      // when it resolves, either by React or something in userspace.
-      break;
     case 'fulfilled':
     case 'rejected':
       // A thenable that already resolved shouldn't have been thrown, so this is
@@ -57,9 +53,12 @@ export function trackSuspendedWakeable(wakeable: Wakeable) {
       // TODO: Log a warning?
       break;
     default: {
-      // TODO: 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.
+      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(

packages/react-server/src/ReactFlightWakeable.js

@@ -62,7 +62,7 @@ declare module 'ReactFlightDOMRelayClientIntegration' {
   ): JSResourceReference<T>;
   declare export function preloadModule<T>(
     moduleReference: JSResourceReference<T>,
-  ): void;
+  ): null | Promise<void>;
   declare export function requireModule<T>(
     moduleReference: JSResourceReference<T>,
   ): T;
@@ -95,7 +95,7 @@ declare module 'ReactFlightNativeRelayClientIntegration' {
   ): JSResourceReference<T>;
   declare export function preloadModule<T>(
     moduleReference: JSResourceReference<T>,
-  ): void;
+  ): null | Promise<void>;
   declare export function requireModule<T>(
     moduleReference: JSResourceReference<T>,
   ): T;