[Fizz] Push halted await to the owner stack for late-arriving I/O info (facebook#35019)

Author: unstubbable

packages/react-client/src/ReactFlightClient.js

@@ -367,6 +367,7 @@ type Response = {
   _debugRootStack?: null | Error, // DEV-only
   _debugRootTask?: null | ConsoleTask, // DEV-only
   _debugStartTime: number, // DEV-only
+  _debugEndTime?: number, // DEV-only
   _debugIOStarted: boolean, // DEV-only
   _debugFindSourceMapURL?: void | FindSourceMapURLCallback, // DEV-only
   _debugChannel?: void | DebugChannel, // DEV-only
@@ -500,6 +501,34 @@ function createErrorChunk<T>(
   return new ReactPromise(ERRORED, null, error);
 }
 
+function filterDebugInfo(
+  response: Response,
+  value: {_debugInfo: ReactDebugInfo, ...},
+) {
+  if (response._debugEndTime === null) {
+    // No end time was defined, so we keep all debug info entries.
+    return;
+  }
+
+  // Remove any debug info entries that arrived after the defined end time.
+  const relativeEndTime =
+    response._debugEndTime -
+    // $FlowFixMe[prop-missing]
+    performance.timeOrigin;
+  const debugInfo = [];
+  for (let i = 0; i < value._debugInfo.length; i++) {
+    const info = value._debugInfo[i];
+    if (typeof info.time === 'number' && info.time > relativeEndTime) {
+      break;
+    }
+    if (info.awaited != null && info.awaited.end > relativeEndTime) {
+      break;
+    }
+    debugInfo.push(info);
+  }
+  value._debugInfo = debugInfo;
+}
+
 function moveDebugInfoFromChunkToInnerValue<T>(
   chunk: InitializedChunk<T> | InitializedStreamChunk<any>,
   value: T,
@@ -534,7 +563,17 @@ function moveDebugInfoFromChunkToInnerValue<T>(
   }
 }
 
+function processChunkDebugInfo<T>(
+  response: Response,
+  chunk: InitializedChunk<T> | InitializedStreamChunk<any>,
+  value: T,
+): void {
+  filterDebugInfo(response, chunk);
+  moveDebugInfoFromChunkToInnerValue(chunk, value);
+}
+
 function wakeChunk<T>(
+  response: Response,
   listeners: Array<InitializationReference | (T => mixed)>,
   value: T,
   chunk: InitializedChunk<T>,
@@ -544,16 +583,17 @@ function wakeChunk<T>(
     if (typeof listener === 'function') {
       listener(value);
     } else {
-      fulfillReference(listener, value, chunk);
+      fulfillReference(response, listener, value, chunk);
     }
   }
 
   if (__DEV__) {
-    moveDebugInfoFromChunkToInnerValue(chunk, value);
+    processChunkDebugInfo(response, chunk, value);
   }
 }
 
 function rejectChunk(
+  response: Response,
   listeners: Array<InitializationReference | (mixed => mixed)>,
   error: mixed,
 ): void {
@@ -562,7 +602,7 @@ function rejectChunk(
     if (typeof listener === 'function') {
       listener(error);
     } else {
-      rejectReference(listener, error);
+      rejectReference(response, listener.handler, error);
     }
   }
 }
@@ -595,13 +635,14 @@ function resolveBlockedCycle<T>(
 }
 
 function wakeChunkIfInitialized<T>(
+  response: Response,
   chunk: SomeChunk<T>,
   resolveListeners: Array<InitializationReference | (T => mixed)>,
   rejectListeners: null | Array<InitializationReference | (mixed => mixed)>,
 ): void {
   switch (chunk.status) {
     case INITIALIZED:
-      wakeChunk(resolveListeners, chunk.value, chunk);
+      wakeChunk(response, resolveListeners, chunk.value, chunk);
       break;
     case BLOCKED:
       // It is possible that we're blocked on our own chunk if it's a cycle.
@@ -615,7 +656,7 @@ function wakeChunkIfInitialized<T>(
           if (cyclicHandler !== null) {
             // This reference points back to this chunk. We can resolve the cycle by
             // using the value from that handler.
-            fulfillReference(reference, cyclicHandler.value, chunk);
+            fulfillReference(response, reference, cyclicHandler.value, chunk);
             resolveListeners.splice(i, 1);
             i--;
             if (rejectListeners !== null) {
@@ -629,14 +670,15 @@ function wakeChunkIfInitialized<T>(
               case INITIALIZED:
                 const initializedChunk: InitializedChunk<T> = (chunk: any);
                 wakeChunk(
+                  response,
                   resolveListeners,
                   initializedChunk.value,
                   initializedChunk,
                 );
                 return;
               case ERRORED:
                 if (rejectListeners !== null) {
-                  rejectChunk(rejectListeners, chunk.reason);
+                  rejectChunk(response, rejectListeners, chunk.reason);
                 }
                 return;
             }
@@ -666,7 +708,7 @@ function wakeChunkIfInitialized<T>(
       break;
     case ERRORED:
       if (rejectListeners) {
-        rejectChunk(rejectListeners, chunk.reason);
+        rejectChunk(response, rejectListeners, chunk.reason);
       }
       break;
   }
@@ -724,7 +766,7 @@ function triggerErrorOnChunk<T>(
   erroredChunk.status = ERRORED;
   erroredChunk.reason = error;
   if (listeners !== null) {
-    rejectChunk(listeners, error);
+    rejectChunk(response, listeners, error);
   }
 }
 
@@ -832,7 +874,7 @@ function resolveModelChunk<T>(
     // longer be rendered or might not be the highest pri.
     initializeModelChunk(resolvedChunk);
     // The status might have changed after initialization.
-    wakeChunkIfInitialized(chunk, resolveListeners, rejectListeners);
+    wakeChunkIfInitialized(response, chunk, resolveListeners, rejectListeners);
   }
 }
 
@@ -861,12 +903,11 @@ function resolveModuleChunk<T>(
   }
   if (resolveListeners !== null) {
     initializeModuleChunk(resolvedChunk);
-    wakeChunkIfInitialized(chunk, resolveListeners, rejectListeners);
+    wakeChunkIfInitialized(response, chunk, resolveListeners, rejectListeners);
   }
 }
 
 type InitializationReference = {
-  response: Response, // TODO: Remove Response from here and pass it through instead.
   handler: InitializationHandler,
   parentObject: Object,
   key: string,
@@ -1005,7 +1046,7 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
         if (typeof listener === 'function') {
           listener(value);
         } else {
-          fulfillReference(listener, value, cyclicChunk);
+          fulfillReference(response, listener, value, cyclicChunk);
         }
       }
     }
@@ -1026,7 +1067,7 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
     initializedChunk.value = value;
 
     if (__DEV__) {
-      moveDebugInfoFromChunkToInnerValue(initializedChunk, value);
+      processChunkDebugInfo(response, initializedChunk, value);
     }
   } catch (error) {
     const erroredChunk: ErroredChunk<T> = (chunk: any);
@@ -1413,11 +1454,12 @@ function getChunk(response: Response, id: number): SomeChunk<any> {
 }
 
 function fulfillReference(
+  response: Response,
   reference: InitializationReference,
   value: any,
   fulfilledChunk: SomeChunk<any>,
 ): void {
-  const {response, handler, parentObject, key, map, path} = reference;
+  const {handler, parentObject, key, map, path} = reference;
 
   for (let i = 1; i < path.length; i++) {
     while (
@@ -1487,7 +1529,11 @@ function fulfillReference(
             return;
           }
           default: {
-            rejectReference(reference, referencedChunk.reason);
+            rejectReference(
+              response,
+              reference.handler,
+              referencedChunk.reason,
+            );
             return;
           }
         }
@@ -1585,21 +1631,20 @@ function fulfillReference(
     initializedChunk.value = handler.value;
     initializedChunk.reason = handler.reason; // Used by streaming chunks
     if (resolveListeners !== null) {
-      wakeChunk(resolveListeners, handler.value, initializedChunk);
+      wakeChunk(response, resolveListeners, handler.value, initializedChunk);
     } else {
       if (__DEV__) {
-        moveDebugInfoFromChunkToInnerValue(initializedChunk, handler.value);
+        processChunkDebugInfo(response, initializedChunk, handler.value);
       }
     }
   }
 }
 
 function rejectReference(
-  reference: InitializationReference,
+  response: Response,
+  handler: InitializationHandler,
   error: mixed,
 ): void {
-  const {handler, response} = reference;
-
   if (handler.errored) {
     // We've already errored. We could instead build up an AggregateError
     // but if there are multiple errors we just take the first one like
@@ -1690,7 +1735,6 @@ function waitForReference<T>(
   }
 
   const reference: InitializationReference = {
-    response,
     handler,
     parentObject,
     key,
@@ -1838,10 +1882,10 @@ function loadServerReference<A: Iterable<any>, T>(
       initializedChunk.status = INITIALIZED;
       initializedChunk.value = handler.value;
       if (resolveListeners !== null) {
-        wakeChunk(resolveListeners, handler.value, initializedChunk);
+        wakeChunk(response, resolveListeners, handler.value, initializedChunk);
       } else {
         if (__DEV__) {
-          moveDebugInfoFromChunkToInnerValue(initializedChunk, handler.value);
+          processChunkDebugInfo(response, initializedChunk, handler.value);
         }
       }
     }
@@ -2578,6 +2622,7 @@ function ResponseInstance(
   replayConsole: boolean, // DEV-only
   environmentName: void | string, // DEV-only
   debugStartTime: void | number, // DEV-only
+  debugEndTime: void | number, // DEV-only
   debugChannel: void | DebugChannel, // DEV-only
 ) {
   const chunks: Map<number, SomeChunk<any>> = new Map();
@@ -2645,6 +2690,7 @@ function ResponseInstance(
       // and is not considered I/O required to load the stream.
       setTimeout(markIOStarted.bind(this), 0);
     }
+    this._debugEndTime = debugEndTime == null ? null : debugEndTime;
     this._debugFindSourceMapURL = findSourceMapURL;
     this._debugChannel = debugChannel;
     this._blockedConsole = null;
@@ -2688,6 +2734,7 @@ export function createResponse(
   replayConsole: boolean, // DEV-only
   environmentName: void | string, // DEV-only
   debugStartTime: void | number, // DEV-only
+  debugEndTime: void | number, // DEV-only
   debugChannel: void | DebugChannel, // DEV-only
 ): WeakResponse {
   return getWeakResponse(
@@ -2704,6 +2751,7 @@ export function createResponse(
       replayConsole,
       environmentName,
       debugStartTime,
+      debugEndTime,
       debugChannel,
     ),
   );
@@ -3075,10 +3123,10 @@ function resolveStream<T: ReadableStream | $AsyncIterable<any, any, void>>(
   resolvedChunk.value = stream;
   resolvedChunk.reason = controller;
   if (resolveListeners !== null) {
-    wakeChunk(resolveListeners, chunk.value, (chunk: any));
+    wakeChunk(response, resolveListeners, chunk.value, (chunk: any));
   } else {
     if (__DEV__) {
-      moveDebugInfoFromChunkToInnerValue(resolvedChunk, stream);
+      processChunkDebugInfo(response, resolvedChunk, stream);
     }
   }
 }
@@ -3218,7 +3266,12 @@ function startAsyncIterable<T>(
         initializedChunk.status = INITIALIZED;
         initializedChunk.value = {done: false, value: value};
         if (resolveListeners !== null) {
-          wakeChunkIfInitialized(chunk, resolveListeners, rejectListeners);
+          wakeChunkIfInitialized(
+            response,
+            chunk,
+            resolveListeners,
+            rejectListeners,
+          );
         }
       }
       nextWriteIndex++;

packages/react-server-dom-esm/src/client/ReactFlightDOMClientBrowser.js

@@ -53,6 +53,7 @@ export type Options = {
   replayConsoleLogs?: boolean,
   environmentName?: string,
   startTime?: number,
+  endTime?: number,
 };
 
 function createDebugCallbackFromWritableStream(
@@ -107,6 +108,7 @@ function createResponseFromOptions(options: void | Options) {
     __DEV__ && options && options.startTime != null
       ? options.startTime
       : undefined,
+    __DEV__ && options && options.endTime != null ? options.endTime : undefined,
     debugChannel,
   );
 }

packages/react-server-dom-esm/src/client/ReactFlightDOMClientNode.js

@@ -58,6 +58,7 @@ export type Options = {
   replayConsoleLogs?: boolean,
   environmentName?: string,
   startTime?: number,
+  endTime?: number,
   // For the Node.js client we only support a single-direction debug channel.
   debugChannel?: Readable,
 };
@@ -116,6 +117,7 @@ function createFromNodeStream<T>(
     __DEV__ && options && options.startTime != null
       ? options.startTime
       : undefined,
+    __DEV__ && options && options.endTime != null ? options.endTime : undefined,
     debugChannel,
   );
 

packages/react-server-dom-parcel/src/client/ReactFlightDOMClientBrowser.js

@@ -132,6 +132,7 @@ function createResponseFromOptions(options: void | Options) {
     __DEV__ && options && options.startTime != null
       ? options.startTime
       : undefined,
+    __DEV__ && options && options.endTime != null ? options.endTime : undefined,
     debugChannel,
   );
 }
@@ -209,6 +210,7 @@ export type Options = {
   replayConsoleLogs?: boolean,
   environmentName?: string,
   startTime?: number,
+  endTime?: number,
 };
 
 export function createFromReadableStream<T>(

packages/react-server-dom-parcel/src/client/ReactFlightDOMClientEdge.js

@@ -80,6 +80,7 @@ export type Options = {
   replayConsoleLogs?: boolean,
   environmentName?: string,
   startTime?: number,
+  endTime?: number,
   // For the Edge client we only support a single-direction debug channel.
   debugChannel?: {readable?: ReadableStream, ...},
 };
@@ -111,6 +112,7 @@ function createResponseFromOptions(options?: Options) {
     __DEV__ && options && options.startTime != null
       ? options.startTime
       : undefined,
+    __DEV__ && options && options.endTime != null ? options.endTime : undefined,
     debugChannel,
   );
 }