[Flight] Cache the value if we visit the same I/O or Promise multiple times along different paths (facebook#35005)

We avoid visiting the same async node twice but if we see it again we
returned "null" indicating that there's no I/O there.

This means that if you have two different Promises both resolving from
the same I/O node then we only show one of them. However, in general we
treat that as two different I/O entries to allow for things like
batching to still show up separately.

This fixes that by caching the return value for multiple visits. So if
we found I/O (but no user space await) in one path and then we visit
that path through a different Promise chain, then we'll still emit it
twice.

However, if we visit the same exact Promise that we emitted an await on
then we skip it. Because there's no need to emit two awaits on the same
thing. It only matters when the path ends up informing whether it has
I/O or not.

Author: sebmarkbage

packages/react-server/src/ReactFlightServer.js

@@ -2316,15 +2316,37 @@ function visitAsyncNode(
   request: Request,
   task: Task,
   node: AsyncSequence,
-  visited: Set<AsyncSequence | ReactDebugInfo>,
+  visited: Map<
+    AsyncSequence | ReactDebugInfo,
+    void | null | PromiseNode | IONode,
+  >,
   cutOff: number,
 ): void | null | PromiseNode | IONode {
   if (visited.has(node)) {
     // It's possible to visit them same node twice when it's part of both an "awaited" path
     // and a "previous" path. This also gracefully handles cycles which would be a bug.
-    return null;
+    return visited.get(node);
+  }
+  // Set it as visited early in case we see ourselves before returning.
+  visited.set(node, null);
+  const result = visitAsyncNodeImpl(request, task, node, visited, cutOff);
+  if (result !== null) {
+    // If we ended up with a value, let's use that value for future visits.
+    visited.set(node, result);
   }
-  visited.add(node);
+  return result;
+}
+
+function visitAsyncNodeImpl(
+  request: Request,
+  task: Task,
+  node: AsyncSequence,
+  visited: Map<
+    AsyncSequence | ReactDebugInfo,
+    void | null | PromiseNode | IONode,
+  >,
+  cutOff: number,
+): void | null | PromiseNode | IONode {
   if (node.end >= 0 && node.end <= request.timeOrigin) {
     // This was already resolved when we started this render. It must have been either something
     // that's part of a start up sequence or externally cached data. We exclude that information.
@@ -2416,7 +2438,7 @@ function visitAsyncNode(
       if (promise !== undefined) {
         const debugInfo = promise._debugInfo;
         if (debugInfo != null && !visited.has(debugInfo)) {
-          visited.add(debugInfo);
+          visited.set(debugInfo, null);
           forwardDebugInfo(request, task, debugInfo);
         }
       }
@@ -2483,6 +2505,10 @@ function visitAsyncNode(
               // Promise that was ultimately awaited by the user space await.
               serializeIONode(request, ioNode, awaited.promise);
 
+              // If we ever visit this I/O node again, skip it because we already emitted this
+              // exact entry and we don't need two awaits on the same thing.
+              visited.set(ioNode, null);
+
               // Ensure the owner is already outlined.
               if (node.owner != null) {
                 outlineComponentInfo(request, node.owner);
@@ -2521,7 +2547,7 @@ function visitAsyncNode(
       if (promise !== undefined) {
         const debugInfo = promise._debugInfo;
         if (debugInfo != null && !visited.has(debugInfo)) {
-          visited.add(debugInfo);
+          visited.set(debugInfo, null);
           forwardDebugInfo(request, task, debugInfo);
         }
       }
@@ -2542,9 +2568,12 @@ function emitAsyncSequence(
   owner: null | ReactComponentInfo,
   stack: null | Error,
 ): void {
-  const visited: Set<AsyncSequence | ReactDebugInfo> = new Set();
+  const visited: Map<
+    AsyncSequence | ReactDebugInfo,
+    void | null | PromiseNode | IONode,
+  > = new Map();
   if (__DEV__ && alreadyForwardedDebugInfo) {
-    visited.add(alreadyForwardedDebugInfo);
+    visited.set(alreadyForwardedDebugInfo, null);
   }
   const awaitedNode = visitAsyncNode(request, task, node, visited, task.time);
   if (awaitedNode === undefined) {