[Flight] Eval Fake Server Component Functions to Recreate Native Stacks (#29632)

We have three kinds of stacks that we send in the RSC protocol:

- The stack trace where a replayed `console.log` was called on the
server.
- The JSX callsite that created a Server Component which then later
called another component.
- The JSX callsite that created a Host or Client Component.

These stack frames disappear in native stacks on the client since
they're executed on the server. This evals a fake file which only has
one call in it on the same line/column as the server. Then we call
through these fake modules to "replay" the callstack. We then replay the
`console.log` within this stack, or call `console.createTask` in this
stack to recreate the stack.

The main concern with this approach is the performance. It adds
significant cost to create all these eval:ed functions but it should
eventually balance out.

This doesn't yet apply source maps to these. With source maps it'll be
able to show the server source code when clicking the links.

I don't love how these appear.

- Because we haven't yet initialized the client module we don't have the
name of the client component we're about to render yet which leads to
the `<...>` task name.
- The `(async)` suffix Chrome adds is still a problem.
- The VMxxxx prefix is used to disambiguate which is noisy. Might be
helped by source maps.
- The continuation of the async stacks end up rooted somewhere in the
bootstrapping of the app. This might be ok when the bootstrapping ends
up ignore listed but it's kind of a problem that you can't clear the
async stack.

<img width="927" alt="Screenshot 2024-05-28 at 11 58 56 PM"
src="https://github.com/facebook/react/assets/63648/1c9d32ce-e671-47c8-9d18-9fab3bffabd0">

<img width="431" alt="Screenshot 2024-05-28 at 11 58 07 PM"
src="https://github.com/facebook/react/assets/63648/52f57518-bbed-400e-952d-6650835ac6b6">
<img width="327" alt="Screenshot 2024-05-28 at 11 58 31 PM"
src="https://github.com/facebook/react/assets/63648/d311a639-79a1-457f-9a46-4f3298d07e65">

<img width="817" alt="Screenshot 2024-05-28 at 11 59 12 PM"
src="https://github.com/facebook/react/assets/63648/3aefd356-acf4-4daa-bdbf-b8c8345f6d4b">

Author: sebmarkbage

packages/react-client/src/ReactFlightClient.js

@@ -67,8 +67,11 @@ import {
   REACT_ELEMENT_TYPE,
   REACT_POSTPONE_TYPE,
   ASYNC_ITERATOR,
+  REACT_FRAGMENT_TYPE,
 } from 'shared/ReactSymbols';
 
+import getComponentNameFromType from 'shared/getComponentNameFromType';
+
 export type {CallServerCallback, EncodeFormActionCallback};
 
 interface FlightStreamController {
@@ -573,6 +576,43 @@ function nullRefGetter() {
   }
 }
 
+function getServerComponentTaskName(componentInfo: ReactComponentInfo): string {
+  return '<' + (componentInfo.name || '...') + '>';
+}
+
+function getTaskName(type: mixed): string {
+  if (type === REACT_FRAGMENT_TYPE) {
+    return '<>';
+  }
+  if (typeof type === 'function') {
+    // This is a function so it must have been a Client Reference that resolved to
+    // a function. We use "use client" to indicate that this is the boundary into
+    // the client. There should only be one for any given owner chain.
+    return '"use client"';
+  }
+  if (
+    typeof type === 'object' &&
+    type !== null &&
+    type.$$typeof === REACT_LAZY_TYPE
+  ) {
+    if (type._init === readChunk) {
+      // This is a lazy node created by Flight. It is probably a client reference.
+      // We use the "use client" string to indicate that this is the boundary into
+      // the client. There will only be one for any given owner chain.
+      return '"use client"';
+    }
+    // We don't want to eagerly initialize the initializer in DEV mode so we can't
+    // call it to extract the type so we don't know the type of this component.
+    return '<...>';
+  }
+  try {
+    const name = getComponentNameFromType(type);
+    return name ? '<' + name + '>' : '<...>';
+  } catch (x) {
+    return '<...>';
+  }
+}
+
 function createElement(
   type: mixed,
   key: mixed,
@@ -647,11 +687,28 @@ function createElement(
         writable: true,
         value: stack,
       });
+
+      let task: null | ConsoleTask = null;
+      if (supportsCreateTask && stack !== null) {
+        const createTaskFn = (console: any).createTask.bind(
+          console,
+          getTaskName(type),
+        );
+        const callStack = buildFakeCallStack(stack, createTaskFn);
+        // This owner should ideally have already been initialized to avoid getting
+        // user stack frames on the stack.
+        const ownerTask = owner === null ? null : initializeFakeTask(owner);
+        if (ownerTask === null) {
+          task = callStack();
+        } else {
+          task = ownerTask.run(callStack);
+        }
+      }
       Object.defineProperty(element, '_debugTask', {
         configurable: false,
         enumerable: false,
         writable: true,
-        value: null,
+        value: task,
       });
     }
     // TODO: We should be freezing the element but currently, we might write into
@@ -1582,6 +1639,118 @@ function resolveHint<Code: HintCode>(
   dispatchHint(code, hintModel);
 }
 
+// eslint-disable-next-line react-internal/no-production-logging
+const supportsCreateTask =
+  __DEV__ && enableOwnerStacks && !!(console: any).createTask;
+
+const taskCache: null | WeakMap<
+  ReactComponentInfo | ReactAsyncInfo,
+  ConsoleTask,
+> = supportsCreateTask ? new WeakMap() : null;
+
+type FakeFunction<T> = (FakeFunction<T>) => T;
+const fakeFunctionCache: Map<string, FakeFunction<any>> = __DEV__
+  ? new Map()
+  : (null: any);
+
+function createFakeFunction<T>(
+  name: string,
+  filename: string,
+  line: number,
+  col: number,
+): FakeFunction<T> {
+  // This creates a fake copy of a Server Module. It represents a module that has already
+  // executed on the server but we re-execute a blank copy for its stack frames on the client.
+
+  const comment =
+    '/* This module was rendered by a Server Component. Turn on Source Maps to see the server source. */';
+
+  // We generate code where the call is at the line and column of the server executed code.
+  // This allows us to use the original source map as the source map of this fake file to
+  // point to the original source.
+  let code;
+  if (line <= 1) {
+    code = '_=>' + ' '.repeat(col < 4 ? 0 : col - 4) + '_()\n' + comment + '\n';
+  } else {
+    code =
+      comment +
+      '\n'.repeat(line - 2) +
+      '_=>\n' +
+      ' '.repeat(col < 1 ? 0 : col - 1) +
+      '_()\n';
+  }
+
+  if (filename) {
+    code += '//# sourceURL=' + filename;
+  }
+
+  // eslint-disable-next-line no-eval
+  const fn: FakeFunction<T> = (0, eval)(code);
+  // $FlowFixMe[cannot-write]
+  Object.defineProperty(fn, 'name', {value: name || '(anonymous)'});
+  // $FlowFixMe[prop-missing]
+  fn.displayName = name;
+  return fn;
+}
+
+const frameRegExp =
+  /^ {3} at (?:(.+) \(([^\)]+):(\d+):(\d+)\)|([^\)]+):(\d+):(\d+))$/;
+
+function buildFakeCallStack<T>(stack: string, innerCall: () => T): () => T {
+  const frames = stack.split('\n');
+  let callStack = innerCall;
+  for (let i = 0; i < frames.length; i++) {
+    const frame = frames[i];
+    let fn = fakeFunctionCache.get(frame);
+    if (fn === undefined) {
+      const parsed = frameRegExp.exec(frame);
+      if (!parsed) {
+        // We assume the server returns a V8 compatible stack trace.
+        continue;
+      }
+      const name = parsed[1] || '';
+      const filename = parsed[2] || parsed[5] || '';
+      const line = +(parsed[3] || parsed[6]);
+      const col = +(parsed[4] || parsed[7]);
+      fn = createFakeFunction(name, filename, line, col);
+    }
+    callStack = fn.bind(null, callStack);
+  }
+  return callStack;
+}
+
+function initializeFakeTask(
+  debugInfo: ReactComponentInfo | ReactAsyncInfo,
+): null | ConsoleTask {
+  if (taskCache === null || typeof debugInfo.stack !== 'string') {
+    return null;
+  }
+  const componentInfo: ReactComponentInfo = (debugInfo: any); // Refined
+  const stack: string = debugInfo.stack;
+  const cachedEntry = taskCache.get((componentInfo: any));
+  if (cachedEntry !== undefined) {
+    return cachedEntry;
+  }
+
+  const ownerTask =
+    componentInfo.owner == null
+      ? null
+      : initializeFakeTask(componentInfo.owner);
+
+  // eslint-disable-next-line react-internal/no-production-logging
+  const createTaskFn = (console: any).createTask.bind(
+    console,
+    getServerComponentTaskName(componentInfo),
+  );
+  const callStack = buildFakeCallStack(stack, createTaskFn);
+
+  if (ownerTask === null) {
+    return callStack();
+  } else {
+    return ownerTask.run(callStack);
+  }
+}
+
 function resolveDebugInfo(
   response: Response,
   id: number,
@@ -1594,6 +1763,10 @@ function resolveDebugInfo(
       'resolveDebugInfo should never be called in production mode. This is a bug in React.',
     );
   }
+  // We eagerly initialize the fake task because this resolving happens outside any
+  // render phase so we're not inside a user space stack at this point. If we waited
+  // to initialize it when we need it, we might be inside user code.
+  initializeFakeTask(debugInfo);
   const chunk = getChunk(response, id);
   const chunkDebugInfo: ReactDebugInfo =
     chunk._debugInfo || (chunk._debugInfo = []);
@@ -1615,12 +1788,28 @@ function resolveConsoleEntry(
   const payload: [string, string, null | ReactComponentInfo, string, mixed] =
     parseModel(response, value);
   const methodName = payload[0];
-  // TODO: Restore the fake stack before logging.
-  // const stackTrace = payload[1];
-  // const owner = payload[2];
+  const stackTrace = payload[1];
+  const owner = payload[2];
   const env = payload[3];
   const args = payload.slice(4);
-  printToConsole(methodName, args, env);
+  if (!enableOwnerStacks) {
+    // Printing with stack isn't really limited to owner stacks but
+    // we gate it behind the same flag for now while iterating.
+    printToConsole(methodName, args, env);
+    return;
+  }
+  const callStack = buildFakeCallStack(
+    stackTrace,
+    printToConsole.bind(null, methodName, args, env),
+  );
+  if (owner != null) {
+    const task = initializeFakeTask(owner);
+    if (task !== null) {
+      task.run(callStack);
+      return;
+    }
+  }
+  callStack();
 }
 
 function mergeBuffer(

packages/react-server/src/ReactFlightServer.js

@@ -241,14 +241,7 @@ function patchConsole(consoleInst: typeof console, methodName: string) {
         // Extract the stack. Not all console logs print the full stack but they have at
         // least the line it was called from. We could optimize transfer by keeping just
         // one stack frame but keeping it simple for now and include all frames.
-        let stack = filterDebugStack(new Error('react-stack-top-frame'));
-        const firstLine = stack.indexOf('\n');
-        if (firstLine === -1) {
-          stack = '';
-        } else {
-          // Skip the console wrapper itself.
-          stack = stack.slice(firstLine + 1);
-        }
+        const stack = filterDebugStack(new Error('react-stack-top-frame'));
         request.pendingChunks++;
         // We don't currently use this id for anything but we emit it so that we can later
         // refer to previous logs in debug info to associate them with a component.