[Flight] Encode ReactIOInfo as its own row type (facebook#33390)

Stacked on facebook#33388.

This encodes the I/O entries as their own row type (`"J"`). This makes
it possible to parse them directly without first parsing the debug info
for each component. E.g. if you're just interested in logging the I/O
without all the places it was awaited.

This is not strictly necessary since the debug info is also readily
available without parsing the actual trees. (That's how the Server
Components Performance Track works.) However, we might want to exclude
this information in profiling builds while retaining some limited form
of I/O tracking.

It also allows for logging side-effects that are not awaited if we
wanted to.

Author: sebmarkbage

packages/react-client/src/ReactFlightClient.js

@@ -13,6 +13,7 @@ import type {
   ReactComponentInfo,
   ReactEnvironmentInfo,
   ReactAsyncInfo,
+  ReactIOInfo,
   ReactTimeInfo,
   ReactStackTrace,
   ReactFunctionLocation,
@@ -47,6 +48,7 @@ import {
   enablePostpone,
   enableProfilerTimer,
   enableComponentPerformanceTrack,
+  enableAsyncDebugInfo,
 } from 'shared/ReactFeatureFlags';
 
 import {
@@ -672,6 +674,14 @@ function nullRefGetter() {
   }
 }
 
+function getIOInfoTaskName(ioInfo: ReactIOInfo): string {
+  return ''; // TODO
+}
+
+function getAsyncInfoTaskName(asyncInfo: ReactAsyncInfo): string {
+  return 'await'; // We could be smarter about this and give it a name like `then` or `Promise.all`.
+}
+
 function getServerComponentTaskName(componentInfo: ReactComponentInfo): string {
   return '<' + (componentInfo.name || '...') + '>';
 }
@@ -2447,30 +2457,27 @@ function getRootTask(
 
 function initializeFakeTask(
   response: Response,
-  debugInfo: ReactComponentInfo | ReactAsyncInfo,
+  debugInfo: ReactComponentInfo | ReactAsyncInfo | ReactIOInfo,
   childEnvironmentName: string,
 ): null | ConsoleTask {
   if (!supportsCreateTask) {
     return null;
   }
-  const componentInfo: ReactComponentInfo = (debugInfo: any); // Refined
   if (debugInfo.stack == null) {
     // If this is an error, we should've really already initialized the task.
     // If it's null, we can't initialize a task.
     return null;
   }
   const stack = debugInfo.stack;
   const env: string =
-    componentInfo.env == null
-      ? response._rootEnvironmentName
-      : componentInfo.env;
+    debugInfo.env == null ? response._rootEnvironmentName : debugInfo.env;
   if (env !== childEnvironmentName) {
     // This is the boundary between two environments so we'll annotate the task name.
     // That is unusual so we don't cache it.
     const ownerTask =
-      componentInfo.owner == null
+      debugInfo.owner == null
         ? null
-        : initializeFakeTask(response, componentInfo.owner, env);
+        : initializeFakeTask(response, debugInfo.owner, env);
     return buildFakeTask(
       response,
       ownerTask,
@@ -2479,20 +2486,27 @@ function initializeFakeTask(
       env,
     );
   } else {
-    const cachedEntry = componentInfo.debugTask;
+    const cachedEntry = debugInfo.debugTask;
     if (cachedEntry !== undefined) {
       return cachedEntry;
     }
     const ownerTask =
-      componentInfo.owner == null
+      debugInfo.owner == null
         ? null
-        : initializeFakeTask(response, componentInfo.owner, env);
+        : initializeFakeTask(response, debugInfo.owner, env);
+    // Some unfortunate pattern matching to refine the type.
+    const taskName =
+      debugInfo.key !== undefined
+        ? getServerComponentTaskName(((debugInfo: any): ReactComponentInfo))
+        : debugInfo.name !== undefined
+          ? getIOInfoTaskName(((debugInfo: any): ReactIOInfo))
+          : getAsyncInfoTaskName(((debugInfo: any): ReactAsyncInfo));
     // $FlowFixMe[cannot-write]: We consider this part of initialization.
-    return (componentInfo.debugTask = buildFakeTask(
+    return (debugInfo.debugTask = buildFakeTask(
       response,
       ownerTask,
       stack,
-      getServerComponentTaskName(componentInfo),
+      taskName,
       env,
     ));
   }
@@ -2555,7 +2569,7 @@ function fakeJSXCallSite() {
 
 function initializeFakeStack(
   response: Response,
-  debugInfo: ReactComponentInfo | ReactAsyncInfo,
+  debugInfo: ReactComponentInfo | ReactAsyncInfo | ReactIOInfo,
 ): void {
   const cachedEntry = debugInfo.debugStack;
   if (cachedEntry !== undefined) {
@@ -2740,6 +2754,54 @@ function resolveConsoleEntry(
   );
 }
 
+function initializeIOInfo(response: Response, ioInfo: ReactIOInfo): void {
+  const env =
+    // TODO: Pass env through I/O info.
+    // ioInfo.env !== undefined ? ioInfo.env :
+    response._rootEnvironmentName;
+  if (ioInfo.stack !== undefined) {
+    initializeFakeTask(response, ioInfo, env);
+    initializeFakeStack(response, ioInfo);
+  }
+  // TODO: Initialize owner.
+  // Adjust the time to the current environment's time space.
+  // $FlowFixMe[cannot-write]
+  ioInfo.start += response._timeOrigin;
+  // $FlowFixMe[cannot-write]
+  ioInfo.end += response._timeOrigin;
+}
+
+function resolveIOInfo(
+  response: Response,
+  id: number,
+  model: UninitializedModel,
+): void {
+  const chunks = response._chunks;
+  let chunk = chunks.get(id);
+  if (!chunk) {
+    chunk = createResolvedModelChunk(response, model);
+    chunks.set(id, chunk);
+    initializeModelChunk(chunk);
+  } else {
+    resolveModelChunk(chunk, model);
+    if (chunk.status === RESOLVED_MODEL) {
+      initializeModelChunk(chunk);
+    }
+  }
+  if (chunk.status === INITIALIZED) {
+    initializeIOInfo(response, chunk.value);
+  } else {
+    chunk.then(
+      v => {
+        initializeIOInfo(response, v);
+      },
+      e => {
+        // Ignore debug info errors for now. Unnecessary noise.
+      },
+    );
+  }
+}
+
 function mergeBuffer(
   buffer: Array<Uint8Array>,
   lastChunk: Uint8Array,
@@ -2844,7 +2906,7 @@ function flushComponentPerformance(
 
   // First find the start time of the first component to know if it was running
   // in parallel with the previous.
-  const debugInfo = root._debugInfo;
+  const debugInfo = __DEV__ && root._debugInfo;
   if (debugInfo) {
     for (let i = 1; i < debugInfo.length; i++) {
       const info = debugInfo[i];
@@ -3101,6 +3163,17 @@ function processFullStringRow(
       }
       // Fallthrough to share the error with Console entries.
     }
+    case 74 /* "J" */: {
+      if (
+        enableProfilerTimer &&
+        enableComponentPerformanceTrack &&
+        enableAsyncDebugInfo
+      ) {
+        resolveIOInfo(response, id, row);
+        return;
+      }
+      // Fallthrough to share the error with Console entries.
+    }
     case 87 /* "W" */: {
       if (__DEV__) {
         resolveConsoleEntry(response, row);

packages/react-server/src/ReactFlightServer.js

@@ -1905,7 +1905,7 @@ function visitAsyncNode(
             return ioNode;
           }
           // Outline the IO node.
-          emitIOChunk(request, ioNode);
+          serializeIONode(request, ioNode);
           // Then emit a reference to us awaiting it in the current task.
           request.pendingChunks++;
           emitDebugChunk(request, task.id, {
@@ -1942,7 +1942,7 @@ function emitAsyncSequence(
       // each occurrence. Right now we'll only track the first time it is invoked.
       awaitedNode.end = performance.now();
     }
-    emitIOChunk(request, awaitedNode);
+    serializeIONode(request, awaitedNode);
     request.pendingChunks++;
     emitDebugChunk(request, task.id, {
       awaited: ((awaitedNode: any): ReactIOInfo), // This is deduped by this reference.
@@ -3493,80 +3493,88 @@ function outlineComponentInfo(
   request.writtenObjects.set(componentInfo, serializeByValueID(id));
 }
 
-function outlineIOInfo(request: Request, ioInfo: ReactIOInfo): void {
+function emitIOInfoChunk(
+  request: Request,
+  id: number,
+  start: number,
+  end: number,
+  stack: ?ReactStackTrace,
+): void {
   if (!__DEV__) {
     // These errors should never make it into a build so we don't need to encode them in codes.json
     // eslint-disable-next-line react-internal/prod-error-codes
     throw new Error(
-      'outlineIOInfo should never be called in production mode. This is a bug in React.',
+      'emitIOInfoChunk should never be called in production mode. This is a bug in React.',
     );
   }
 
-  if (request.writtenObjects.has(ioInfo)) {
-    // Already written
-    return;
-  }
-
-  // Limit the number of objects we write to prevent emitting giant props objects.
   let objectLimit = 10;
-  if (ioInfo.stack != null) {
-    // Ensure we have enough object limit to encode the stack trace.
-    objectLimit += ioInfo.stack.length;
+  if (stack) {
+    objectLimit += stack.length;
   }
-
-  // We use the console encoding so that we can dedupe objects but don't necessarily
-  // use the full serialization that requires a task.
   const counter = {objectLimit};
+  function replacer(
+    this:
+      | {+[key: string | number]: ReactClientValue}
+      | $ReadOnlyArray<ReactClientValue>,
+    parentPropertyName: string,
+    value: ReactClientValue,
+  ): ReactJSONValue {
+    return renderConsoleValue(
+      request,
+      counter,
+      this,
+      parentPropertyName,
+      value,
+    );
+  }
 
-  // We can't serialize the ConsoleTask/Error objects so we need to omit them before serializing.
-  const relativeStartTimestamp = ioInfo.start - request.timeOrigin;
-  const relativeEndTimestamp = ioInfo.end - request.timeOrigin;
+  const relativeStartTimestamp = start - request.timeOrigin;
+  const relativeEndTimestamp = end - request.timeOrigin;
   const debugIOInfo: Omit<ReactIOInfo, 'debugTask' | 'debugStack'> = {
     start: relativeStartTimestamp,
     end: relativeEndTimestamp,
-    stack: ioInfo.stack,
+    stack: stack,
   };
-  const id = outlineConsoleValue(request, counter, debugIOInfo);
-  request.writtenObjects.set(ioInfo, serializeByValueID(id));
+  // $FlowFixMe[incompatible-type] stringify can return null
+  const json: string = stringify(debugIOInfo, replacer);
+  const row = id.toString(16) + ':J' + json + '\n';
+  const processedChunk = stringToChunk(row);
+  request.completedRegularChunks.push(processedChunk);
 }
 
-function emitIOChunk(request: Request, ioNode: IONode | PromiseNode): void {
-  if (!__DEV__) {
-    // These errors should never make it into a build so we don't need to encode them in codes.json
-    // eslint-disable-next-line react-internal/prod-error-codes
-    throw new Error(
-      'outlineIOInfo should never be called in production mode. This is a bug in React.',
-    );
+function outlineIOInfo(request: Request, ioInfo: ReactIOInfo): void {
+  if (request.writtenObjects.has(ioInfo)) {
+    // Already written
+    return;
   }
+  // We can't serialize the ConsoleTask/Error objects so we need to omit them before serializing.
+  request.pendingChunks++;
+  const id = request.nextChunkId++;
+  emitIOInfoChunk(request, id, ioInfo.start, ioInfo.end, ioInfo.stack);
+  request.writtenObjects.set(ioInfo, serializeByValueID(id));
+}
 
-  if (request.writtenObjects.has(ioNode)) {
+function serializeIONode(
+  request: Request,
+  ioNode: IONode | PromiseNode,
+): string {
+  const existingRef = request.writtenObjects.get(ioNode);
+  if (existingRef !== undefined) {
     // Already written
-    return;
+    return existingRef;
   }
 
-  // Limit the number of objects we write to prevent emitting giant props objects.
-  let objectLimit = 10;
   let stack = null;
   if (ioNode.stack !== null) {
     stack = filterStackTrace(request, ioNode.stack, 1);
-    // Ensure we have enough object limit to encode the stack trace.
-    objectLimit += stack.length;
   }
-
-  // We use the console encoding so that we can dedupe objects but don't necessarily
-  // use the full serialization that requires a task.
-  const counter = {objectLimit};
-
-  // We can't serialize the ConsoleTask/Error objects so we need to omit them before serializing.
-  const relativeStartTimestamp = ioNode.start - request.timeOrigin;
-  const relativeEndTimestamp = ioNode.end - request.timeOrigin;
-  const debugIOInfo: Omit<ReactIOInfo, 'debugTask' | 'debugStack'> = {
-    start: relativeStartTimestamp,
-    end: relativeEndTimestamp,
-    stack: stack,
-  };
-  const id = outlineConsoleValue(request, counter, debugIOInfo);
-  request.writtenObjects.set(ioNode, serializeByValueID(id));
+  request.pendingChunks++;
+  const id = request.nextChunkId++;
+  emitIOInfoChunk(request, id, ioNode.start, ioNode.end, stack);
+  const ref = serializeByValueID(id);
+  request.writtenObjects.set(ioNode, ref);
+  return ref;
 }
 
 function emitTypedArrayChunk(