[Fizz] Track postpones in fallbacks (#27421)

This fixes so that you can postpone in a fallback. This postpones the
parent boundary. I track the fallbacks in a separate replay node so that
when we resume, we can replay the fallback itself and finish the
fallback and then possibly later the content. By doing this we also
ensure we don't complete the parent too early since now it has a render
task on it.

There is one case that this surfaces that isn't limited to
prerender/resume but also render/hydrateRoot. I left todos in the tests
for this.

If you postpone in a fallback, and suspend in the content but eventually
don't postpone in the content then we should be able to just skip
postponing since the content rendered and we no longer need the
fallback. This is a bit of a weird edge case though since fallbacks are
supposed to be very minimal.

This happens because in both cases the fallback starts rendering early
as soon as the content suspends. This also ensures that the parent
doesn't complete early by increasing the blocking tasks. Unfortunately,
the fallback will irreversibly postpone its parent boundary as soon as
it hits a postpone.

When you suspend, the same thing happens but we typically deal with this
by doing a "soft" abort on the fallback since we don't need it anymore
which unblocks the parent boundary. We can't do that with postpone right
now though since it's considered a terminal state.

I think I'll just leave this as is for now since it's an edge case but
it's an annoying exception in the model. Makes me feel I haven't quite
nailed it just yet.

DiffTrain build for [bff6be8](bff6be8)

Author: sebmarkbage

compiled/facebook-www/REVISION

@@ -1 +1 @@
-7f6201889e8e628eeb53e05d8850ddffa3c2e74a
+bff6be8eb1d77980c13f3e01be63cb813a377058

compiled/facebook-www/ReactDOMServer-dev.classic.js

@@ -19,7 +19,7 @@ if (__DEV__) {
 var React = require("react");
 var ReactDOM = require("react-dom");
 
-var ReactVersion = "18.3.0-www-classic-8314963c";
+var ReactVersion = "18.3.0-www-classic-87fc393f";
 
 // This refers to a WWW module.
 var warningWWW = require("warning");
@@ -9994,7 +9994,7 @@ function pingTask(request, task) {
   }
 }
 
-function createSuspenseBoundary(request, fallbackAbortableTasks, keyPath) {
+function createSuspenseBoundary(request, fallbackAbortableTasks) {
   return {
     status: PENDING,
     rootSegmentID: -1,
@@ -10005,7 +10005,8 @@ function createSuspenseBoundary(request, fallbackAbortableTasks, keyPath) {
     fallbackAbortableTasks: fallbackAbortableTasks,
     errorDigest: null,
     resources: createBoundaryResources(),
-    keyPath: keyPath
+    trackedContentKeyPath: null,
+    trackedFallbackNode: null
   };
 }
 
@@ -10271,7 +10272,12 @@ function renderSuspenseBoundary(request, someTask, keyPath, props) {
   var fallback = props.fallback;
   var content = props.children;
   var fallbackAbortSet = new Set();
-  var newBoundary = createSuspenseBoundary(request, fallbackAbortSet, keyPath);
+  var newBoundary = createSuspenseBoundary(request, fallbackAbortSet);
+
+  if (request.trackedPostpones !== null) {
+    newBoundary.trackedContentKeyPath = keyPath;
+  }
+
   var insertionIndex = parentSegment.chunks.length; // The children of the boundary segment is actually the fallback.
 
   var boundarySegment = createPendingSegment(
@@ -10363,6 +10369,25 @@ function renderSuspenseBoundary(request, someTask, keyPath, props) {
     task.blockedBoundary = parentBoundary;
     task.blockedSegment = parentSegment;
     task.keyPath = prevKeyPath;
+  }
+
+  var fallbackKeyPath = [keyPath[0], "Suspense Fallback", keyPath[2]];
+  var trackedPostpones = request.trackedPostpones;
+
+  if (trackedPostpones !== null) {
+    // We create a detached replay node to track any postpones inside the fallback.
+    var fallbackReplayNode = [fallbackKeyPath[1], fallbackKeyPath[2], [], null];
+    trackedPostpones.workingMap.set(fallbackKeyPath, fallbackReplayNode);
+
+    if (newBoundary.status === POSTPONED) {
+      // This must exist now.
+      var boundaryReplayNode = trackedPostpones.workingMap.get(keyPath);
+      boundaryReplayNode[4] = fallbackReplayNode;
+    } else {
+      // We might not inject it into the postponed tree, unless the content actually
+      // postpones too. We need to keep track of it until that happpens.
+      newBoundary.trackedFallbackNode = fallbackReplayNode;
+    }
   } // We create suspended task for the fallback because we don't want to actually work
   // on it yet in case we finish the main content, so we queue for later.
 
@@ -10373,8 +10398,8 @@ function renderSuspenseBoundary(request, someTask, keyPath, props) {
     -1,
     parentBoundary,
     boundarySegment,
-    fallbackAbortSet, // TODO: Should distinguish key path of fallback and primary tasks
-    keyPath,
+    fallbackAbortSet,
+    fallbackKeyPath,
     task.formatContext,
     task.legacyContext,
     task.context,
@@ -10397,19 +10422,18 @@ function replaySuspenseBoundary(
   props,
   id,
   childNodes,
-  childSlots
+  childSlots,
+  fallbackNodes,
+  fallbackSlots
 ) {
   pushBuiltInComponentStackInDEV(task, "Suspense");
   var prevKeyPath = task.keyPath;
   var previousReplaySet = task.replay;
   var parentBoundary = task.blockedBoundary;
   var content = props.children;
+  var fallback = props.fallback;
   var fallbackAbortSet = new Set();
-  var resumedBoundary = createSuspenseBoundary(
-    request,
-    fallbackAbortSet,
-    task.keyPath
-  );
+  var resumedBoundary = createSuspenseBoundary(request, fallbackAbortSet);
   resumedBoundary.parentFlushed = true; // We restore the same id of this boundary as was used during prerender.
 
   resumedBoundary.rootSegmentID = id; // We can reuse the current context and task to render the content immediately without
@@ -10438,14 +10462,6 @@ function replaySuspenseBoundary(
       renderNode(request, task, content, -1);
     }
 
-    if (
-      resumedBoundary.pendingTasks === 0 &&
-      resumedBoundary.status === PENDING
-    ) {
-      resumedBoundary.status = COMPLETED;
-      request.completedBoundaries.push(resumedBoundary);
-    }
-
     if (task.replay.pendingTasks === 1 && task.replay.nodes.length > 0) {
       throw new Error(
         "Couldn't find all resumable slots by key/index during replaying. " +
@@ -10454,6 +10470,19 @@ function replaySuspenseBoundary(
     }
 
     task.replay.pendingTasks--;
+
+    if (
+      resumedBoundary.pendingTasks === 0 &&
+      resumedBoundary.status === PENDING
+    ) {
+      resumedBoundary.status = COMPLETED;
+      request.completedBoundaries.push(resumedBoundary); // This must have been the last segment we were waiting on. This boundary is now complete.
+      // Therefore we won't need the fallback. We early return so that we don't have to create
+      // the fallback.
+
+      popComponentStackInDEV(task);
+      return;
+    }
   } catch (error) {
     resumedBoundary.status = CLIENT_RENDERED;
     var errorDigest;
@@ -10484,7 +10513,66 @@ function replaySuspenseBoundary(
     task.blockedBoundary = parentBoundary;
     task.replay = previousReplaySet;
     task.keyPath = prevKeyPath;
-  } // TODO: Should this be in the finally?
+  }
+
+  var fallbackKeyPath = [keyPath[0], "Suspense Fallback", keyPath[2]];
+  var suspendedFallbackTask; // We create suspended task for the fallback because we don't want to actually work
+  // on it yet in case we finish the main content, so we queue for later.
+
+  if (typeof fallbackSlots === "number") {
+    // Resuming directly in the fallback.
+    var resumedSegment = createPendingSegment(
+      request,
+      0,
+      null,
+      task.formatContext,
+      false,
+      false
+    );
+    resumedSegment.id = fallbackSlots;
+    resumedSegment.parentFlushed = true;
+    suspendedFallbackTask = createRenderTask(
+      request,
+      null,
+      fallback,
+      -1,
+      parentBoundary,
+      resumedSegment,
+      fallbackAbortSet,
+      fallbackKeyPath,
+      task.formatContext,
+      task.legacyContext,
+      task.context,
+      task.treeContext
+    );
+  } else {
+    var fallbackReplay = {
+      nodes: fallbackNodes,
+      slots: fallbackSlots,
+      pendingTasks: 0
+    };
+    suspendedFallbackTask = createReplayTask(
+      request,
+      null,
+      fallbackReplay,
+      fallback,
+      -1,
+      parentBoundary,
+      fallbackAbortSet,
+      fallbackKeyPath,
+      task.formatContext,
+      task.legacyContext,
+      task.context,
+      task.treeContext
+    );
+  }
+
+  {
+    suspendedFallbackTask.componentStack = task.componentStack;
+  } // TODO: This should be queued at a separate lower priority queue so that we only work
+  // on preparing fallbacks if we don't have any more main content to task on.
+
+  request.pingedTasks.push(suspendedFallbackTask); // TODO: Should this be in the finally?
 
   popComponentStackInDEV(task);
 }
@@ -11392,9 +11480,11 @@ function replayElement(
         task,
         keyPath,
         props,
-        node[4],
+        node[5],
         node[2],
-        node[3]
+        node[3],
+        node[4] === null ? [] : node[4][2],
+        node[4] === null ? null : node[4][3]
       );
     } // We finished rendering this node, so now we can consume this
     // slot. This must happen after in case we rerender this task.
@@ -12104,11 +12194,7 @@ function abortRemainingSuspenseBoundary(
   error,
   errorDigest
 ) {
-  var resumedBoundary = createSuspenseBoundary(
-    request,
-    new Set(),
-    null // The keyPath doesn't matter at this point so we don't bother rebuilding it.
-  );
+  var resumedBoundary = createSuspenseBoundary(request, new Set());
   resumedBoundary.parentFlushed = true; // We restore the same id of this boundary as was used during prerender.
 
   resumedBoundary.rootSegmentID = rootSegmentID;
@@ -12163,7 +12249,7 @@ function abortRemainingReplayNodes(
       );
     } else {
       var boundaryNode = node;
-      var rootSegmentID = boundaryNode[4];
+      var rootSegmentID = boundaryNode[5];
       abortRemainingSuspenseBoundary(
         request,
         rootSegmentID,
@@ -12950,9 +13036,7 @@ function flushCompletedQueues(request, destination) {
             destination,
             request.resumableState,
             request.renderState,
-            request.allPendingTasks === 0 &&
-              (request.trackedPostpones === null ||
-                request.trackedPostpones.workingMap.size === 0)
+            request.allPendingTasks === 0 && request.trackedPostpones === null
           );
         }