Fix stacks for unresolved I/O during abort

Author: unstubbable

packages/react-server/src/ReactFlightServer.js

@@ -75,6 +75,7 @@ import type {
   AsyncSequence,
   IONode,
   PromiseNode,
+  UnresolvedAwaitNode,
   UnresolvedPromiseNode,
 } from './ReactFlightAsyncSequence';
 
@@ -95,6 +96,7 @@ import {
   markAsyncSequenceRootTask,
   getCurrentAsyncSequence,
   getAsyncSequenceFromPromise,
+  getInternalAwaitNode,
   parseStackTrace,
   parseStackTracePrivate,
   supportsComponentStorage,
@@ -2305,7 +2307,18 @@ function visitAsyncNode(
             // We aborted this render. If this Promise spanned the abort time it was probably the
             // Promise that was aborted. This won't necessarily have I/O associated with it but
             // it's a point of interest.
-            match = node;
+            // However, if the Promise and IO node have the same owner, it likely means a sync component
+            // created both the Promise and initiated the I/O. Prefer the IO node for more specific info.
+            if (
+              ioNode !== null &&
+              ioNode.tag === IO_NODE &&
+              node.owner !== null &&
+              ioNode.owner === node.owner
+            ) {
+              match = ioNode;
+            } else {
+              match = node;
+            }
           }
         } else if (ioNode !== null) {
           // This Promise was blocked on I/O. That's a signal that this Promise is interesting to log.
@@ -4417,7 +4430,7 @@ function outlineIOInfo(request: Request, ioInfo: ReactIOInfo): void {
 
 function serializeIONode(
   request: Request,
-  ioNode: IONode | PromiseNode | UnresolvedPromiseNode,
+  ioNode: IONode | PromiseNode | UnresolvedPromiseNode | UnresolvedAwaitNode,
   promiseRef: null | WeakRef<Promise<mixed>>,
 ): string {
   const existingRef = request.writtenDebugObjects.get(ioNode);
@@ -5383,26 +5396,29 @@ function forwardDebugInfoFromAbortedTask(request: Request, task: Task): void {
           // See if any of the dependencies are resolved yet.
           node = node.awaited;
         }
+        // For unresolved Promises, check if we have an internal await node that shows what
+        // the async function is currently blocked on. For resolved Promises, the regular
+        // awaited field already contains the necessary information.
         if (node.tag === UNRESOLVED_PROMISE_NODE) {
-          // We don't know what Promise will eventually end up resolving this Promise and if it
-          // was I/O at all. However, we assume that it was some kind of I/O since it didn't
-          // complete in time before aborting.
-          // The best we can do is try to emit the stack of where this Promise was created.
+          const internalAwait = getInternalAwaitNode(node);
+          if (internalAwait !== null) {
+            node = internalAwait;
+          }
+        }
+        if (node.tag === UNRESOLVED_AWAIT_NODE) {
+          // We found the await that's blocking. Use its stack to show where the component is stuck.
           serializeIONode(request, node, null);
           request.pendingChunks++;
           const env = (0, request.environmentName)();
           const asyncInfo: ReactAsyncInfo = {
-            awaited: ((node: any): ReactIOInfo), // This is deduped by this reference.
+            awaited: ((node: any): ReactIOInfo),
             env: env,
           };
-          // We don't have a start time for this await but in case there was no start time emitted
-          // we need to include something. TODO: We should maybe ideally track the time when we
-          // called .then() but without updating the task.time field since that's used for the cutoff.
           advanceTaskTime(request, task, task.time);
           emitDebugChunk(request, task.id, asyncInfo);
         } else {
-          // We have a resolved Promise. Its debug info can include both awaited data and rejected
-          // promises after the abort.
+          // We have a resolved or unresolved Promise. Its debug info can include both awaited
+          // data and rejected promises after the abort.
           emitAsyncSequence(request, task, sequence, debugInfo, null, null);
         }
       }

packages/react-server/src/ReactFlightServerConfigDebugNode.js

@@ -7,7 +7,7 @@
  * @flow
  */
 
-import type {ReactStackTrace} from 'shared/ReactTypes';
+import type {ReactStackTrace, ReactComponentInfo} from 'shared/ReactTypes';
 
 import type {
   AsyncSequence,
@@ -40,6 +40,35 @@ const pendingOperations: Map<number, AsyncSequence> =
 // Keep the last resolved await as a workaround for async functions missing data.
 let lastRanAwait: null | AwaitNode = null;
 
+// These two maps work together to track what async functions are blocked on when aborting:
+//
+// 1. unresolvedPromiseNodesByOwner: Maps owner -> Promise (to find the Promise to link)
+//    When a Promise is created, we track it by its owner. This typically captures async
+//    function return Promises. Sync components may also have Promises tracked here, but
+//    they won't be linked since sync functions can't have awaits with matching owners.
+//
+// 2. internalAwaitNodesByPromise: Maps Promise -> await (stores the actual link)
+//    When an await happens with the same owner as a tracked Promise, we link that Promise
+//    to the await. This shows what the async function is currently blocked on.
+//
+// By storing the links separately from the regular awaited field, we can use this information
+// only during abort scenarios without affecting normal rendering.
+const unresolvedPromiseNodesByOwner: WeakMap<
+  ReactComponentInfo,
+  UnresolvedPromiseNode,
+> = new WeakMap();
+const internalAwaitNodesByPromise: WeakMap<
+  UnresolvedPromiseNode | PromiseNode,
+  UnresolvedAwaitNode | AwaitNode,
+> = new WeakMap();
+
+export function getInternalAwaitNode(
+  promiseNode: UnresolvedPromiseNode | PromiseNode,
+): null | UnresolvedAwaitNode | AwaitNode {
+  const awaitNode = internalAwaitNodesByPromise.get(promiseNode);
+  return awaitNode === undefined ? null : awaitNode;
+}
+
 function resolvePromiseOrAwaitNode(
   unresolvedNode: UnresolvedAwaitNode | UnresolvedPromiseNode,
   endTime: number,
@@ -114,16 +143,27 @@ export function initAsyncDebugInfo(): void {
               }
             }
             const current = pendingOperations.get(currentAsyncId);
+            const owner = resolveOwner();
             node = ({
               tag: UNRESOLVED_AWAIT_NODE,
-              owner: resolveOwner(),
+              owner: owner,
               stack: stack,
               start: performance.now(),
               end: -1.1, // set when resolved.
               promise: promiseRef,
               awaited: trigger, // The thing we're awaiting on. Might get overrriden when we resolve.
               previous: current === undefined ? null : current, // The path that led us here.
             }: UnresolvedAwaitNode);
+            // Link the owner's Promise to this await so we can track what it's blocked on.
+            // This only links when the await and Promise have the same owner (i.e., async functions
+            // awaiting within themselves). Promises from sync components won't match any awaits.
+            // We store this in a separate WeakMap to avoid affecting normal rendering.
+            if (owner !== null) {
+              const ownerPromiseNode = unresolvedPromiseNodesByOwner.get(owner);
+              if (ownerPromiseNode !== undefined) {
+                internalAwaitNodesByPromise.set(ownerPromiseNode, node);
+              }
+            }
           } else {
             const owner = resolveOwner();
             node = ({
@@ -140,6 +180,17 @@ export function initAsyncDebugInfo(): void {
                   : trigger,
               previous: null,
             }: UnresolvedPromiseNode);
+            // Track Promises by owner so awaits with matching owners can link to them.
+            // Only track the first Promise per owner. This typically captures async function
+            // return Promises, but may also track Promises from sync components - those won't
+            // be linked since sync functions can't have awaits with matching owners.
+            if (
+              owner !== null &&
+              trigger === undefined &&
+              !unresolvedPromiseNodesByOwner.has(owner)
+            ) {
+              unresolvedPromiseNodesByOwner.set(owner, node);
+            }
           }
         } else if (
           type !== 'Microtask' &&