[Flight] Optimize Async Stack Collection (#33727)

We need to optimize the collection of debug info for dev mode. This is
an incredibly hot path since it instruments all I/O and Promises in the
app.

These optimizations focus primarily on the collection of stack traces.
They are expensive to collect because we need to eagerly collect the
stacks since they can otherwise cause memory leaks. We also need to do
some of the processing of them up front. We also end up only using a few
of them in the end but we don't know which ones we'll use.

The first compromise here is that I now only collect the stacks of
"awaits" if they were in a specific request's render. In some cases it's
useful to collect them even outside of this if they're part of a
sequence that started early. I still collect stacks for the created
Promises outside of this though which can still provide some context.

The other optimization to awaits, is that since we'll only use the inner
most one that had an await directly in userspace, we can stop collecting
stacks on a chain of awaits after we find one. This requires a quick
filter on a single callsite to determine. Since we now only collect
stacks from awaits that belongs to a specific Request we can use that
request's specific filter option. Technically this might not be quite
correct if that same thing ends up deduped across Requests but that's an
edge case.

Additionally, I now stop collecting stack for I/O nodes. They're almost
always superseded by the Promise that wraps them anyway. Even if you
write mostly Promise free code, you'll likely end up with a Promise at
the root of the component eventually anyway and then you end up using
its stack anyway. You have to really contort the code to end up with
zero Promises at which point it's not very useful anyway. At best it's
maybe mostly useful for giving a name to the I/O when the rest is just
stuff like `new Promise`.

However, a possible alternative optimization could be to *only* collect
the stack of spawned I/O and not the stack of Promises. The issue with
Promises (not awaits) is that we never know what will end up resolving
them in the end when they're created so we have to always eagerly
collect stacks. This could be an issue when you have a lot of
abstractions that end up not actually be related to I/O at all. The
issue with collecting stacks only for I/O is that the actual I/O can be
pooled or batched so you end up not having the stack when the conceptual
start of each operation within the batch started. Which is why I decided
to keep the Promise stack.

packages/react-server/src/ReactFlightServerConfigDebugNode.js

@@ -7,6 +7,8 @@
  * @flow
  */
 
+import type {ReactStackTrace} from 'shared/ReactTypes';
+
 import type {
   AsyncSequence,
   IONode,
@@ -24,6 +26,7 @@ import {
   UNRESOLVED_AWAIT_NODE,
 } from './ReactFlightAsyncSequence';
 import {resolveOwner} from './flight/ReactFlightCurrentOwner';
+import {resolveRequest, isAwaitInUserspace} from './ReactFlightServer';
 import {createHook, executionAsyncId, AsyncResource} from 'async_hooks';
 import {enableAsyncDebugInfo} from 'shared/ReactFeatureFlags';
 import {parseStackTrace} from './ReactFlightServerConfig';
@@ -66,6 +69,8 @@ function resolvePromiseOrAwaitNode(
   return resolvedNode;
 }
 
+const emptyStack: ReactStackTrace = [];
+
 // Initialize the tracing of async operations.
 // We do this globally since the async work can potentially eagerly
 // start before the first request and once requests start they can interleave.
@@ -110,10 +115,33 @@ export function initAsyncDebugInfo(): void {
             }
             // If the thing we're waiting on is another Await we still track that sequence
             // so that we can later pick the best stack trace in user space.
+            let stack = null;
+            if (
+              trigger.stack !== null &&
+              (trigger.tag === AWAIT_NODE ||
+                trigger.tag === UNRESOLVED_AWAIT_NODE)
+            ) {
+              // We already had a stack for an await. In a chain of awaits we'll only need one good stack.
+              // We mark it with an empty stack to signal to any await on this await that we have a stack.
+              stack = emptyStack;
+            } else {
+              const request = resolveRequest();
+              if (request === null) {
+                // We don't collect stacks for awaits that weren't in the scope of a specific render.
+              } else {
+                stack = parseStackTrace(new Error(), 5);
+                if (!isAwaitInUserspace(request, stack)) {
+                  // If this await was not done directly in user space, then clear the stack. We won't use it
+                  // anyway. This lets future awaits on this await know that we still need to get their stacks
+                  // until we find one in user space.
+                  stack = null;
+                }
+              }
+            }
             node = ({
               tag: UNRESOLVED_AWAIT_NODE,
               owner: resolveOwner(),
-              stack: parseStackTrace(new Error(), 5),
+              stack: stack,
               start: performance.now(),
               end: -1.1, // set when resolved.
               promise: new WeakRef((resource: Promise<any>)),
@@ -145,7 +173,7 @@ export function initAsyncDebugInfo(): void {
             node = ({
               tag: IO_NODE,
               owner: resolveOwner(),
-              stack: parseStackTrace(new Error(), 3), // This is only used if no native promises are used.
+              stack: null,
               start: performance.now(),
               end: -1.1, // Only set when pinged.
               promise: null,
@@ -160,7 +188,7 @@ export function initAsyncDebugInfo(): void {
             node = ({
               tag: IO_NODE,
               owner: resolveOwner(),
-              stack: parseStackTrace(new Error(), 3),
+              stack: null,
               start: performance.now(),
               end: -1.1, // Only set when pinged.
               promise: null,