[DevTools] Track Server Environment Names of Each SuspenseNode (#34605)

Tracks the environment names of the I/O in each SuspenseNode and sent it
to the front end when the suspenders change.

In the front end, every child boundary should really be treated as it
has all environment names of the parents too since they're blocked by
the parent too. We could do this tracking on backend but if there's ever
one added on the root would need to be send for every child.

This lets us highlight which subtrees are blocked by content on the
server.

---------

Co-authored-by: Sebastian "Sebbie" Silbermann <[email protected]>

Author: sebmarkbage

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

@@ -299,6 +299,7 @@ type SuspenseNode = {
   nextSibling: null | 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.
   // 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,
@@ -327,6 +328,7 @@ function createSuspenseNode(
     nextSibling: null,
     rects: null,
     suspendedBy: new Map(),
+    environments: new Map(),
     hasUniqueSuspenders: false,
     hasUnknownSuspenders: false,
   });
@@ -2220,6 +2222,10 @@ export function attach(
         }
         operations[i++] = fiberIdWithChanges;
         operations[i++] = suspense.hasUniqueSuspenders ? 1 : 0;
+        operations[i++] = suspense.environments.size;
+        suspense.environments.forEach((count, env) => {
+          operations[i++] = getStringID(env);
+        });
       });
     }
 
@@ -2725,6 +2731,13 @@ export function attach(
       return;
     }
 
+    // TODO: Just enqueue the operations here instead of stashing by id.
+
+    // Ensure each environment gets recorded in the string table since it is emitted
+    // before we loop it over again later during flush.
+    suspenseNode.environments.forEach((count, env) => {
+      getStringID(env);
+    });
     pendingSuspenderChanges.add(fiberInstance.id);
   }
 
@@ -2807,7 +2820,20 @@ export function attach(
     let suspendedBySet = suspenseNodeSuspendedBy.get(ioInfo);
     if (suspendedBySet === undefined) {
       suspendedBySet = new Set();
-      suspenseNodeSuspendedBy.set(asyncInfo.awaited, suspendedBySet);
+      suspenseNodeSuspendedBy.set(ioInfo, suspendedBySet);
+      // We've added a dependency. We must increment the ref count of the environment.
+      const env = ioInfo.env;
+      if (env != null) {
+        const environmentCounts = parentSuspenseNode.environments;
+        const count = environmentCounts.get(env);
+        if (count === undefined || count === 0) {
+          environmentCounts.set(env, 1);
+          // We've discovered a new environment for this SuspenseNode. We'll to update the node.
+          recordSuspenseSuspenders(parentSuspenseNode);
+        } else {
+          environmentCounts.set(env, count + 1);
+        }
+      }
     }
     // The child of the Suspense boundary that was suspended on this, or null if suspended at the root.
     // This is used to keep track of how many dependents are still alive and also to get information
@@ -2897,6 +2923,7 @@ export function attach(
         : instance.suspenseNode;
     if (previousSuspendedBy !== null && suspenseNode !== null) {
       const nextSuspendedBy = instance.suspendedBy;
+      let changedEnvironment = false;
       for (let i = 0; i < previousSuspendedBy.length; i++) {
         const asyncInfo = previousSuspendedBy[i];
         if (
@@ -2935,7 +2962,26 @@ export function attach(
             }
           }
           if (suspendedBySet !== undefined && suspendedBySet.size === 0) {
-            suspenseNode.suspendedBy.delete(asyncInfo.awaited);
+            suspenseNode.suspendedBy.delete(ioInfo);
+            // Successfully removed all dependencies. We can decrement the ref count of the environment.
+            const env = ioInfo.env;
+            if (env != null) {
+              const environmentCounts = suspenseNode.environments;
+              const count = environmentCounts.get(env);
+              if (count === undefined || count === 0) {
+                throw new Error(
+                  'We are removing an environment but it was not in the set. ' +
+                    'This is a bug in React.',
+                );
+              }
+              if (count === 1) {
+                environmentCounts.delete(env);
+                // Last one. We've now change the set of environments. We'll need to update the node.
+                changedEnvironment = true;
+              } else {
+                environmentCounts.set(env, count - 1);
+              }
+            }
           }
           if (
             suspenseNode.hasUniqueSuspenders &&
@@ -2948,6 +2994,9 @@ export function attach(
           }
         }
       }
+      if (changedEnvironment) {
+        recordSuspenseSuspenders(suspenseNode);
+      }
     }
   }
 

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

@@ -1759,12 +1759,22 @@ export default class Store extends EventEmitter<{
           break;
         }
         case SUSPENSE_TREE_OPERATION_SUSPENDERS: {
-          const changeLength = operations[i + 1];
-          i += 2;
+          i++;
+          const changeLength = operations[i++];
 
           for (let changeIndex = 0; changeIndex < changeLength; changeIndex++) {
-            const id = operations[i];
-            const hasUniqueSuspenders = operations[i + 1] === 1;
+            const id = operations[i++];
+            const hasUniqueSuspenders = operations[i++] === 1;
+            const environmentNamesLength = operations[i++];
+            const environmentNames = [];
+            for (
+              let envIndex = 0;
+              envIndex < environmentNamesLength;
+              envIndex++
+            ) {
+              const environmentNameStringID = operations[i++];
+              environmentNames.push(stringTable[environmentNameStringID]);
+            }
             const suspense = this._idToSuspense.get(id);
 
             if (suspense === undefined) {
@@ -1777,8 +1787,6 @@ export default class Store extends EventEmitter<{
               break;
             }
 
-            i += 2;
-
             if (__DEBUG__) {
               const previousHasUniqueSuspenders = suspense.hasUniqueSuspenders;
               debug(
@@ -1788,6 +1796,7 @@ export default class Store extends EventEmitter<{
             }
 
             suspense.hasUniqueSuspenders = hasUniqueSuspenders;
+            // TODO: Recompute the environment names.
           }
 
           hasSuspenseTreeChanged = true;

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

@@ -454,14 +454,22 @@ function updateTree(
       }
 
       case SUSPENSE_TREE_OPERATION_SUSPENDERS: {
-        const changesLength = ((operations[i + 1]: any): number);
-
-        if (__DEBUG__) {
-          const changes = operations.slice(i + 2, i + 2 + changesLength * 2);
-          debug('Suspender changes', `[${changes.join(',')}]`);
+        i++;
+        const changeLength = ((operations[i++]: any): number);
+
+        for (let changeIndex = 0; changeIndex < changeLength; changeIndex++) {
+          const suspenseNodeId = operations[i++];
+          const hasUniqueSuspenders = operations[i++] === 1;
+          const environmentNamesLength = operations[i++];
+          i += environmentNamesLength;
+          if (__DEBUG__) {
+            debug(
+              'Suspender changes',
+              `Suspense node ${suspenseNodeId} unique suspenders set to ${String(hasUniqueSuspenders)} with ${String(environmentNamesLength)} environments`,
+            );
+          }
         }
 
-        i += 2 + changesLength * 2;
         break;
       }
 

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

@@ -426,12 +426,18 @@ export function printOperationsArray(operations: Array<number>) {
         break;
       }
       case SUSPENSE_TREE_OPERATION_SUSPENDERS: {
-        const changeLength = operations[i + 1];
-        i += 2;
-        const changes = operations.slice(i, i + changeLength * 2);
-        i += changeLength;
-
-        logs.push(`Suspense node suspender changes ${changes.join(',')}`);
+        i++;
+        const changeLength = ((operations[i++]: any): number);
+
+        for (let changeIndex = 0; changeIndex < changeLength; changeIndex++) {
+          const id = operations[i++];
+          const hasUniqueSuspenders = operations[i++] === 1;
+          const environmentNamesLength = operations[i++];
+          i += environmentNamesLength;
+          logs.push(
+            `Suspense node ${id} unique suspenders set to ${String(hasUniqueSuspenders)} with ${String(environmentNamesLength)} environments`,
+          );
+        }
 
         break;
       }