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,
@@ -4511,7 +4513,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);
@@ -5452,9 +5454,16 @@ function forwardDebugInfoFromCurrentContext(
     }
   }
   if (enableProfilerTimer && enableAsyncDebugInfo) {
-    const sequence = getCurrentAsyncSequence();
-    if (sequence !== null) {
-      emitAsyncSequence(request, task, sequence, debugInfo, null, null);
+    if (request.status === ABORTING) {
+      // When aborting, skip forwarding debug info here.
+      // forwardDebugInfoFromAbortedTask will handle it more accurately.
+    } else {
+      // For normal resolve/reject, use the current execution context, as it
+      // shows what actually caused the Promise to settle.
+      const sequence = getCurrentAsyncSequence();
+      if (sequence !== null) {
+        emitAsyncSequence(request, task, sequence, debugInfo, null, null);
+      }
     }
   }
 }
@@ -5492,11 +5501,18 @@ function forwardDebugInfoFromAbortedTask(request: Request, task: Task): void {
           // See if any of the dependencies are resolved yet.
           node = node.awaited;
         }
+        // For unresolved promise nodes, check if we have an internal await node
+        // that shows what the async function is blocked on.
         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_PROMISE_NODE ||
+          node.tag === UNRESOLVED_AWAIT_NODE
+        ) {
           serializeIONode(request, node, null);
           request.pendingChunks++;
           const env = (0, request.environmentName)();

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,28 @@ 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();
+
 function resolvePromiseOrAwaitNode(
   unresolvedNode: UnresolvedAwaitNode | UnresolvedPromiseNode,
   endTime: number,
@@ -114,16 +136,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 +173,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' &&
@@ -356,3 +400,10 @@ export function getAsyncSequenceFromPromise(
   }
   return node;
 }
+
+export function getInternalAwaitNode(
+  promiseNode: UnresolvedPromiseNode | PromiseNode,
+): null | UnresolvedAwaitNode | AwaitNode {
+  const awaitNode = internalAwaitNodesByPromise.get(promiseNode);
+  return awaitNode === undefined ? null : awaitNode;
+}

packages/react-server/src/ReactFlightServerConfigDebugNoop.js

@@ -7,7 +7,13 @@
  * @flow
  */
 
-import type {AsyncSequence} from './ReactFlightAsyncSequence';
+import type {
+  AsyncSequence,
+  AwaitNode,
+  PromiseNode,
+  UnresolvedAwaitNode,
+  UnresolvedPromiseNode,
+} from './ReactFlightAsyncSequence';
 
 // Exported for runtimes that don't support Promise instrumentation for async debugging.
 export function initAsyncDebugInfo(): void {}
@@ -20,3 +26,8 @@ export function getAsyncSequenceFromPromise(
 ): null | AsyncSequence {
   return null;
 }
+export function getInternalAwaitNode(
+  promiseNode: UnresolvedPromiseNode | PromiseNode,
+): null | UnresolvedAwaitNode | AwaitNode {
+  return null;
+}