[DevTools] Sort suspense timeline by end time instead of just document order (#35011)

Right now it's possible for things like server environments to appear
before other content in the timeline just because it's in a different
document order.

Ofc the order in production is not guaranteed but we can at least use
the timing information we have as a hint towards the actual order.

Unfortunately since the end time of the RSC stream itself is always
after the content that resolved to produce it, it becomes kind of
determined by the chunking. Similarly since for a clean refresh, the
scripts and styles will typically load after the server content they
appear later. Similarly SSR typically finishes after the RSC parts.
Therefore a hack here is that I artificially delay everything with a
non-null environment (RSC) so that RSC always comes after client-side
(Suspense). This is also consistent with how we color things that have
an environment even if children are just Suspense.

To ensure that we never show a child before a parent, in the timeline,
each child has a minimum time of its parent.

packages/react-devtools-shared/src/backend/fiber/renderer.js

@@ -301,6 +301,7 @@ type SuspenseNode = {
   rects: null | Array<Rect>, // The bounding rects of content children.
   suspendedBy: Map<ReactIOInfo, Set<DevToolsInstance>>, // Tracks which data we're suspended by and the children that suspend it.
   environments: Map<string, number>, // Tracks the Flight environment names that suspended this. I.e. if the server blocked this.
+  endTime: number, // Track a short cut to the maximum end time value within the suspendedBy set.
   // Track whether any of the items in suspendedBy are unique this this Suspense boundaries or if they're all
   // also in the parent sets. This determine whether this could contribute in the loading sequence.
   hasUniqueSuspenders: boolean,
@@ -330,6 +331,7 @@ function createSuspenseNode(
     rects: null,
     suspendedBy: new Map(),
     environments: new Map(),
+    endTime: 0,
     hasUniqueSuspenders: false,
     hasUnknownSuspenders: false,
   });
@@ -2156,8 +2158,8 @@ export function attach(
         // Regular operations
         pendingOperations.length +
         // All suspender changes are batched in a single message.
-        // [SUSPENSE_TREE_OPERATION_SUSPENDERS, suspenderChangesLength, ...[id, hasUniqueSuspenders, isSuspended]]
+        // [SUSPENSE_TREE_OPERATION_SUSPENDERS, suspenderChangesLength, ...[id, hasUniqueSuspenders, endTime, isSuspended]]
-        (numSuspenderChanges > 0 ? 2 + numSuspenderChanges * 3 : 0),
+        (numSuspenderChanges > 0 ? 2 + numSuspenderChanges * 4 : 0),
     );
 
     // Identify which renderer this update is coming from.
@@ -2242,6 +2244,7 @@ export function attach(
         }
         operations[i++] = fiberIdWithChanges;
         operations[i++] = suspense.hasUniqueSuspenders ? 1 : 0;
+        operations[i++] = Math.round(suspense.endTime * 1000);
         const instance = suspense.instance;
         const isSuspended =
           // TODO: Track if other SuspenseNode like SuspenseList rows are suspended.
@@ -2912,12 +2915,19 @@ export function attach(
     // like owner instances to link down into the tree.
     if (!suspendedBySet.has(parentInstance)) {
       suspendedBySet.add(parentInstance);
+      const virtualEndTime = getVirtualEndTime(ioInfo);
       if (
         !parentSuspenseNode.hasUniqueSuspenders &&
         !ioExistsInSuspenseAncestor(parentSuspenseNode, ioInfo)
       ) {
         // This didn't exist in the parent before, so let's mark this boundary as having a unique suspender.
         parentSuspenseNode.hasUniqueSuspenders = true;
+        if (parentSuspenseNode.endTime < virtualEndTime) {
+          parentSuspenseNode.endTime = virtualEndTime;
+        }
+        recordSuspenseSuspenders(parentSuspenseNode);
+      } else if (parentSuspenseNode.endTime < virtualEndTime) {
+        parentSuspenseNode.endTime = virtualEndTime;
         recordSuspenseSuspenders(parentSuspenseNode);
       }
     }
@@ -2979,6 +2989,26 @@ export function attach(
     }
   }
 
+  function getVirtualEndTime(ioInfo: ReactIOInfo): number {
+    if (ioInfo.env != null) {
+      // Sort client side content first so that scripts and streams don't
+      // cover up the effect of server time.
+      return ioInfo.end + 1000000;
+    }
+    return ioInfo.end;
+  }
+
+  function computeEndTime(suspenseNode: SuspenseNode) {
+    let maxEndTime = 0;
+    suspenseNode.suspendedBy.forEach((set, ioInfo) => {
+      const virtualEndTime = getVirtualEndTime(ioInfo);
+      if (virtualEndTime > maxEndTime) {
+        maxEndTime = virtualEndTime;
+      }
+    });
+    return maxEndTime;
+  }
+
   function removePreviousSuspendedBy(
     instance: DevToolsInstance,
     previousSuspendedBy: null | Array<ReactAsyncInfo>,
@@ -2996,6 +3026,7 @@ export function attach(
     if (previousSuspendedBy !== null && suspenseNode !== null) {
       const nextSuspendedBy = instance.suspendedBy;
       let changedEnvironment = false;
+      let mayHaveChangedEndTime = false;
       for (let i = 0; i < previousSuspendedBy.length; i++) {
         const asyncInfo = previousSuspendedBy[i];
         if (
@@ -3009,6 +3040,11 @@ export function attach(
           const ioInfo = asyncInfo.awaited;
           const suspendedBySet = suspenseNode.suspendedBy.get(ioInfo);
 
+          if (suspenseNode.endTime === getVirtualEndTime(ioInfo)) {
+            // This may be the only remaining entry at this end time. Recompute the end time.
+            mayHaveChangedEndTime = true;
+          }
+
           if (
             suspendedBySet === undefined ||
             !suspendedBySet.delete(instance)
@@ -3066,7 +3102,11 @@ export function attach(
           }
         }
       }
-      if (changedEnvironment) {
+      const newEndTime = mayHaveChangedEndTime
+        ? computeEndTime(suspenseNode)
+        : suspenseNode.endTime;
+      if (changedEnvironment || newEndTime !== suspenseNode.endTime) {
+        suspenseNode.endTime = newEndTime;
         recordSuspenseSuspenders(suspenseNode);
       }
     }

packages/react-devtools-shared/src/devtools/store.js

@@ -925,7 +925,7 @@ export default class Store extends EventEmitter<{
    */
   getSuspendableDocumentOrderSuspense(
     uniqueSuspendersOnly: boolean,
-  ): $ReadOnlyArray<SuspenseTimelineStep> {
+  ): Array<SuspenseTimelineStep> {
     const target: Array<SuspenseTimelineStep> = [];
     const roots = this.roots;
     let rootStep: null | SuspenseTimelineStep = null;
@@ -949,17 +949,25 @@ export default class Store extends EventEmitter<{
           rootStep = {
             id: suspense.id,
             environment: environmentName,
+            endTime: suspense.endTime,
           };
           target.push(rootStep);
-        } else if (rootStep.environment === null) {
+        } else {
-          // If any root has an environment name, then let's use it.
+          if (rootStep.environment === null) {
-          rootStep.environment = environmentName;
+            // If any root has an environment name, then let's use it.
+            rootStep.environment = environmentName;
+          }
+          if (suspense.endTime > rootStep.endTime) {
+            // If any root has a higher end time, let's use that.
+            rootStep.endTime = suspense.endTime;
+          }
         }
         this.pushTimelineStepsInDocumentOrder(
           suspense.children,
           target,
           uniqueSuspendersOnly,
           environments,
+          0, // Don't pass a minimum end time at the root. The root is always first so doesn't matter.
         );
       }
     }
@@ -972,6 +980,7 @@ export default class Store extends EventEmitter<{
     target: Array<SuspenseTimelineStep>,
     uniqueSuspendersOnly: boolean,
     parentEnvironments: Array<string>,
+    parentEndTime: number,
   ): void {
     for (let i = 0; i < children.length; i++) {
       const child = this.getSuspenseByID(children[i]);
@@ -996,10 +1005,15 @@ export default class Store extends EventEmitter<{
         unionEnvironments.length > 0
           ? unionEnvironments[unionEnvironments.length - 1]
           : null;
+      // The end time of a child boundary can in effect never be earlier than its parent even if
+      // everything unsuspended before that.
+      const maxEndTime =
+        parentEndTime > child.endTime ? parentEndTime : child.endTime;
       if (hasRects && (!uniqueSuspendersOnly || child.hasUniqueSuspenders)) {
         target.push({
           id: child.id,
           environment: environmentName,
+          endTime: maxEndTime,
         });
       }
       this.pushTimelineStepsInDocumentOrder(
@@ -1007,10 +1021,28 @@ export default class Store extends EventEmitter<{
         target,
         uniqueSuspendersOnly,
         unionEnvironments,
+        maxEndTime,
       );
     }
   }
 
+  getEndTimeOrDocumentOrderSuspense(
+    uniqueSuspendersOnly: boolean,
+  ): $ReadOnlyArray<SuspenseTimelineStep> {
+    const timeline =
+      this.getSuspendableDocumentOrderSuspense(uniqueSuspendersOnly);
+    if (timeline.length === 0) {
+      return timeline;
+    }
+    const root = timeline[0];
+    // We mutate in place since we assume we've got a fresh array.
+    timeline.sort((a, b) => {
+      // Root is always first
+      return a === root ? -1 : b === root ? 1 : a.endTime - b.endTime;
+    });
+    return timeline;
+  }
+
   getRendererIDForElement(id: number): number | null {
     let current = this._idToElement.get(id);
     while (current !== undefined) {
@@ -1688,6 +1720,7 @@ export default class Store extends EventEmitter<{
             hasUniqueSuspenders: false,
             isSuspended: isSuspended,
             environments: [],
+            endTime: 0,
           });
 
           hasSuspenseTreeChanged = true;
@@ -1884,6 +1917,7 @@ export default class Store extends EventEmitter<{
           for (let changeIndex = 0; changeIndex < changeLength; changeIndex++) {
             const id = operations[i++];
             const hasUniqueSuspenders = operations[i++] === 1;
+            const endTime = operations[i++] / 1000;
             const isSuspended = operations[i++] === 1;
             const environmentNamesLength = operations[i++];
             const environmentNames = [];
@@ -1919,6 +1953,7 @@ export default class Store extends EventEmitter<{
             }
 
             suspense.hasUniqueSuspenders = hasUniqueSuspenders;
+            suspense.endTime = endTime;
             suspense.isSuspended = isSuspended;
             suspense.environments = environmentNames;
           }

packages/react-devtools-shared/src/devtools/views/Profiler/CommitTreeBuilder.js

@@ -460,13 +460,14 @@ function updateTree(
         for (let changeIndex = 0; changeIndex < changeLength; changeIndex++) {
           const suspenseNodeId = operations[i++];
           const hasUniqueSuspenders = operations[i++] === 1;
+          const endTime = operations[i++] / 1000;
           const isSuspended = operations[i++] === 1;
           const environmentNamesLength = operations[i++];
           i += environmentNamesLength;
           if (__DEBUG__) {
             debug(
               'Suspender changes',
-              `Suspense node ${suspenseNodeId} unique suspenders set to ${String(hasUniqueSuspenders)} is suspended set to ${String(isSuspended)} with ${String(environmentNamesLength)} environments`,
+              `Suspense node ${suspenseNodeId} unique suspenders set to ${String(hasUniqueSuspenders)} ending at ${String(endTime)} is suspended set to ${String(isSuspended)} with ${String(environmentNamesLength)} environments`,
             );
           }
         }

packages/react-devtools-shared/src/devtools/views/SuspenseTab/SuspenseTab.js

@@ -74,7 +74,7 @@ function ToggleUniqueSuspenders() {
   function handleToggleUniqueSuspenders() {
     const nextUniqueSuspendersOnly = !uniqueSuspendersOnly;
     // TODO: Handle different timeline modes (e.g. random order)
-    const nextTimeline = store.getSuspendableDocumentOrderSuspense(
+    const nextTimeline = store.getEndTimeOrDocumentOrderSuspense(
       nextUniqueSuspendersOnly,
     );
     suspenseTreeDispatch({

packages/react-devtools-shared/src/devtools/views/SuspenseTab/SuspenseTreeContext.js

@@ -111,7 +111,7 @@ type Props = {
 function getInitialState(store: Store): SuspenseTreeState {
   const uniqueSuspendersOnly = true;
   const timeline =
-    store.getSuspendableDocumentOrderSuspense(uniqueSuspendersOnly);
+    store.getEndTimeOrDocumentOrderSuspense(uniqueSuspendersOnly);
   const timelineIndex = timeline.length - 1;
   const selectedSuspenseID =
     timelineIndex === -1 ? null : timeline[timelineIndex].id;
@@ -182,7 +182,7 @@ function SuspenseTreeContextController({children}: Props): React.Node {
             }
 
             // TODO: Handle different timeline modes (e.g. random order)
-            const nextTimeline = store.getSuspendableDocumentOrderSuspense(
+            const nextTimeline = store.getEndTimeOrDocumentOrderSuspense(
               state.uniqueSuspendersOnly,
             );
 

packages/react-devtools-shared/src/frontend/types.js

@@ -196,6 +196,7 @@ export type Rect = {
 export type SuspenseTimelineStep = {
   id: SuspenseNode['id'], // TODO: Will become a group.
   environment: null | string,
+  endTime: number,
 };
 
 export type SuspenseNode = {
@@ -207,6 +208,7 @@ export type SuspenseNode = {
   hasUniqueSuspenders: boolean,
   isSuspended: boolean,
   environments: Array<string>,
+  endTime: number,
 };
 
 // Serialized version of ReactIOInfo

packages/react-devtools-shared/src/utils.js

@@ -432,11 +432,12 @@ export function printOperationsArray(operations: Array<number>) {
         for (let changeIndex = 0; changeIndex < changeLength; changeIndex++) {
           const id = operations[i++];
           const hasUniqueSuspenders = operations[i++] === 1;
+          const endTime = operations[i++] / 1000;
           const isSuspended = operations[i++] === 1;
           const environmentNamesLength = operations[i++];
           i += environmentNamesLength;
           logs.push(
-            `Suspense node ${id} unique suspenders set to ${String(hasUniqueSuspenders)} is suspended set to ${String(isSuspended)} with ${String(environmentNamesLength)} environments`,
+            `Suspense node ${id} unique suspenders set to ${String(hasUniqueSuspenders)} ending at ${String(endTime)} is suspended set to ${String(isSuspended)} with ${String(environmentNamesLength)} environments`,
           );
         }