Store profiler snapshot data by root (and clear on root unmount)

src/backend/renderer.js

@@ -651,7 +651,8 @@ export function attach(
   let pendingSimulatedUnmountedIDs: Array<number> = [];
   let pendingOperationsQueue: Array<Uint32Array> | null = [];
   let pendingStringTable: Map<string, number> = new Map();
-  let pendingStringTableLength = 0;
+  let pendingStringTableLength: number = 0;
+  let pendingUnmountedRootID: number | null = null;
 
   function pushOperation(op: number): void {
     if (__DEV__) {
@@ -669,7 +670,8 @@ export function attach(
     if (
       pendingOperations.length === 0 &&
       pendingRealUnmountedIDs.length === 0 &&
-      pendingSimulatedUnmountedIDs.length === 0
+      pendingSimulatedUnmountedIDs.length === 0 &&
+      pendingUnmountedRootID === null
     ) {
       // If we're currently profiling, send an "operations" method even if there are no mutations to the tree.
       // The frontend needs this no-op info to know how to reconstruct the tree for each commit,
@@ -679,17 +681,21 @@ export function attach(
       }
     }
 
+    const numUnmountIDs =
+      pendingRealUnmountedIDs.length +
+      pendingSimulatedUnmountedIDs.length +
+      (pendingUnmountedRootID === null ? 0 : 1);
+
     const ops = new Uint32Array(
       // Identify which renderer this update is coming from.
       2 + // [rendererID, rootFiberID]
       // How big is the string table?
       1 + // [stringTableLength]
-      // Then goes the actual string table.
-      pendingStringTableLength +
-      // All unmounts are batched in a single message.
-      2 + // [TREE_OPERATION_REMOVE, removedIDLength]
-        pendingRealUnmountedIDs.length +
-        pendingSimulatedUnmountedIDs.length +
+        // Then goes the actual string table.
+        pendingStringTableLength +
+        // All unmounts are batched in a single message.
+        // [TREE_OPERATION_REMOVE, removedIDLength, ...ids]
+        (numUnmountIDs > 0 ? 2 + numUnmountIDs : 0) +
         // Regular operations
         pendingOperations.length
     );
@@ -699,7 +705,7 @@ export function attach(
     // Which in turn enables fiber props, states, and hooks to be inspected.
     let i = 0;
     ops[i++] = rendererID;
-    ops[i++] = getFiberID(getPrimaryFiber(root.current));
+    ops[i++] = currentRootID; // Use this ID in case the root was unmounted!
 
     // Now fill in the string table.
     // [stringTableLength, str1Length, ...str1, str2Length, ...str2, ...]
@@ -710,24 +716,30 @@ export function attach(
       i += key.length;
     });
 
-    // All unmounts except roots are batched in a single message.
-    ops[i++] = TREE_OPERATION_REMOVE;
-    // The first number is how many unmounted IDs we're gonna send.
-    ops[i++] =
-      pendingRealUnmountedIDs.length + pendingSimulatedUnmountedIDs.length;
-    // Fill in the real unmounts in the reverse order.
-    // They were inserted parents-first by React, but we want children-first.
-    // So we traverse our array backwards.
-    for (let j = pendingRealUnmountedIDs.length - 1; j >= 0; j--) {
-      ops[i++] = pendingRealUnmountedIDs[j];
+    if (numUnmountIDs > 0) {
+      // All unmounts except roots are batched in a single message.
+      ops[i++] = TREE_OPERATION_REMOVE;
+      // The first number is how many unmounted IDs we're gonna send.
+      ops[i++] = numUnmountIDs;
+      // Fill in the real unmounts in the reverse order.
+      // They were inserted parents-first by React, but we want children-first.
+      // So we traverse our array backwards.
+      for (let j = pendingRealUnmountedIDs.length - 1; j >= 0; j--) {
+        ops[i++] = pendingRealUnmountedIDs[j];
+      }
+      // Fill in the simulated unmounts (hidden Suspense subtrees) in their order.
+      // (We want children to go before parents.)
+      // They go *after* the real unmounts because we know for sure they won't be
+      // children of already pushed "real" IDs. If they were, we wouldn't be able
+      // to discover them during the traversal, as they would have been deleted.
+      ops.set(pendingSimulatedUnmountedIDs, i);
+      i += pendingSimulatedUnmountedIDs.length;
+      // The root ID should always be unmounted last.
+      if (pendingUnmountedRootID !== null) {
+        ops[i] = pendingUnmountedRootID;
+        i++;
+      }
     }
-    // Fill in the simulated unmounts (hidden Suspense subtrees) in their order.
-    // (We want children to go before parents.)
-    // They go *after* the real unmounts because we know for sure they won't be
-    // children of already pushed "real" IDs. If they were, we wouldn't be able
-    // to discover them during the traversal, as they would have been deleted.
-    ops.set(pendingSimulatedUnmountedIDs, i);
-    i += pendingSimulatedUnmountedIDs.length;
     // Fill in the rest of the operations.
     ops.set(pendingOperations, i);
 
@@ -747,6 +759,7 @@ export function attach(
     pendingOperations.length = 0;
     pendingRealUnmountedIDs.length = 0;
     pendingSimulatedUnmountedIDs.length = 0;
+    pendingUnmountedRootID = null;
     pendingStringTable.clear();
     pendingStringTableLength = 0;
   }
@@ -858,11 +871,9 @@ export function attach(
     }
     const id = getFiberID(primaryFiber);
     if (isRoot) {
-      // Removing a root needs to happen at the end
-      // so we don't batch it with other unmounts.
-      pushOperation(TREE_OPERATION_REMOVE);
-      pushOperation(1); // Remove one item
-      pushOperation(id);
+      // Roots must be removed only after all children (pending and simultated) have been removed.
+      // So we track it separately.
+      pendingUnmountedRootID = id;
     } else if (!shouldFilterFiber(fiber)) {
       // To maintain child-first ordering,
       // we'll push it into one of these queues,

src/devtools/store.js

@@ -3,6 +3,7 @@
 import EventEmitter from 'events';
 import memoize from 'memoize-one';
 import throttle from 'lodash.throttle';
+import { inspect } from 'util';
 import {
   TREE_OPERATION_ADD,
   TREE_OPERATION_REMOVE,
@@ -106,9 +107,9 @@ export default class Store extends EventEmitter {
   // Once profiling is finished, this snapshot can be used along with "operations" messages emitted during profiling,
   // to reconstruct the state of each root for each commit.
   // It's okay to use a single root to store this information because node IDs are unique across all roots.
-  _profilingSnapshotsByElementID: Map<
+  _profilingSnapshotsByRootID: Map<
     number,
-    ProfilingSnapshotNode
+    Map<number, ProfilingSnapshotNode>
   > = new Map();
 
   // Incremented each time the store is mutated.
@@ -203,8 +204,8 @@ export default class Store extends EventEmitter {
       '_profilingScreenshotsByRootID'
     );
     this.assertEmptyMap(
-      this._profilingSnapshotsByElementID,
-      '_profilingSnapshotsByElementID'
+      this._profilingSnapshotsByRootID,
+      '_profilingSnapshotsByRootID'
     );
     this.assertEmptyMap(this._rootIDToCapabilities, '_rootIDToCapabilities');
     this.assertEmptyMap(this._rootIDToRendererID, '_rootIDToRendererID');
@@ -214,9 +215,9 @@ export default class Store extends EventEmitter {
   assertEmptyMap(map: Map<any, any>, mapName: string) {
     if (map.size !== 0) {
       throw new Error(
-        `Expected ${mapName} to be empty, got ${
-          map.size
-        }: ${require('util').inspect(this, { depth: 20 })}`
+        `Expected ${mapName} to be empty, got ${map.size}: ${inspect(map, {
+          depth: 20,
+        })}`
       );
     }
   }
@@ -290,7 +291,7 @@ export default class Store extends EventEmitter {
     this._importedProfilingData = value;
     this._profilingOperationsByRootID = new Map();
     this._profilingScreenshotsByRootID = new Map();
-    this._profilingSnapshotsByElementID = new Map();
+    this._profilingSnapshotsByRootID = new Map();
     this._profilingCache.invalidate();
 
     this.emit('importedProfilingData');
@@ -316,8 +317,8 @@ export default class Store extends EventEmitter {
     return this._profilingScreenshotsByRootID;
   }
 
-  get profilingSnapshot(): Map<number, ProfilingSnapshotNode> {
-    return this._profilingSnapshotsByElementID;
+  get profilingSnapshots(): Map<number, Map<number, ProfilingSnapshotNode>> {
+    return this._profilingSnapshotsByRootID;
   }
 
   get revision(): number {
@@ -348,7 +349,7 @@ export default class Store extends EventEmitter {
     this._importedProfilingData = null;
     this._profilingOperationsByRootID = new Map();
     this._profilingScreenshotsByRootID = new Map();
-    this._profilingSnapshotsByElementID = new Map();
+    this._profilingSnapshotsByRootID = new Map();
 
     // Invalidate suspense cache if profiling data is being (re-)recorded.
     // Note that we clear now because any existing data is "stale".
@@ -415,10 +416,8 @@ export default class Store extends EventEmitter {
 
   getElementByID(id: number): Element | null {
     const element = this._idToElement.get(id);
-
     if (element == null) {
       console.warn(`No element found with id "${id}"`);
-
       return null;
     }
 
@@ -687,26 +686,22 @@ export default class Store extends EventEmitter {
     THROTTLE_CAPTURE_SCREENSHOT_DURATION
   );
 
-  _takeProfilingSnapshotRecursive = (elementID: number) => {
+  _takeProfilingSnapshotRecursive = (
+    elementID: number,
+    profilingSnapshot: Map<number, ProfilingSnapshotNode>
+  ) => {
     const element = this.getElementByID(elementID);
     if (element !== null) {
-      this._profilingSnapshotsByElementID.set(elementID, {
+      profilingSnapshot.set(elementID, {
         id: elementID,
         children: element.children.slice(0),
         displayName: element.displayName,
         key: element.key,
       });
 
-      element.children.forEach(this._takeProfilingSnapshotRecursive);
-    }
-  };
-
-  _clearProfilingSnapshotRecursive = (elementID: number) => {
-    const element = this.getElementByID(elementID);
-    if (element !== null) {
-      this._profilingSnapshotsByElementID.delete(elementID);
-
-      element.children.forEach(this._clearProfilingSnapshotRecursive);
+      element.children.forEach(childID =>
+        this._takeProfilingSnapshotRecursive(childID, profilingSnapshot)
+      );
     }
   };
 
@@ -921,6 +916,8 @@ export default class Store extends EventEmitter {
               throw new Error(`Node ${id} was removed before its children.`);
             }
 
+            this._idToElement.delete(id);
+
             let parentElement = null;
             if (parentID === 0) {
               if (__DEBUG__) {
@@ -933,11 +930,7 @@ export default class Store extends EventEmitter {
 
               this._profilingOperationsByRootID.delete(id);
               this._profilingScreenshotsByRootID.delete(id);
-
-              // The following call depends on `getElementByID`
-              // which depends on the element being in `_idToElement`,
-              // so we have to do it before removing the element from `_idToElement`.
-              this._clearProfilingSnapshotRecursive(id);
+              this._profilingSnapshotsByRootID.delete(id);
 
               haveRootsChanged = true;
             } else {
@@ -954,8 +947,6 @@ export default class Store extends EventEmitter {
               parentElement.children.splice(index, 1);
             }
 
-            this._idToElement.delete(id);
-
             this._adjustParentTreeWeight(parentElement, -weight);
             removedElementIDs.set(id, parentID);
 
@@ -1051,8 +1042,12 @@ export default class Store extends EventEmitter {
       this._importedProfilingData = null;
       this._profilingOperationsByRootID = new Map();
       this._profilingScreenshotsByRootID = new Map();
-      this._profilingSnapshotsByElementID = new Map();
-      this.roots.forEach(this._takeProfilingSnapshotRecursive);
+      this._profilingSnapshotsByRootID = new Map();
+      this.roots.forEach(rootID => {
+        const profilingSnapshot = new Map();
+        this._profilingSnapshotsByRootID.set(rootID, profilingSnapshot);
+        this._takeProfilingSnapshotRecursive(rootID, profilingSnapshot);
+      });
     }
 
     if (this._isProfiling !== isProfiling) {

src/devtools/views/Profiler/CommitTreeBuilder.js

@@ -15,6 +15,7 @@ import type { ElementType } from 'src/types';
 import type {
   CommitTreeFrontend,
   CommitTreeNodeFrontend,
+  ProfilingSnapshotNode,
   ProfilingSummaryFrontend,
 } from 'src/devtools/views/Profiler/types';
 
@@ -66,13 +67,23 @@ export function getCommitTree({
   if (commitIndex === 0) {
     const nodes = new Map();
 
+    const { importedProfilingData } = store;
+    const profilingSnapshot =
+      importedProfilingData != null
+        ? importedProfilingData.profilingSnapshots.get(rootID)
+        : store.profilingSnapshots.get(rootID);
+
+    if (profilingSnapshot == null) {
+      throw Error(`Could not find profiling snapshot for root "${rootID}"`);
+    }
+
     // Construct the initial tree.
     recursivelyInitializeTree(
       rootID,
       0,
       nodes,
       profilingSummary.initialTreeBaseDurations,
-      store
+      profilingSnapshot
     );
 
     // Mutate the tree
@@ -122,13 +133,8 @@ function recursivelyInitializeTree(
   parentID: number,
   nodes: Map<number, CommitTreeNodeFrontend>,
   initialTreeBaseDurations: Map<number, number>,
-  store: Store
+  profilingSnapshot: Map<number, ProfilingSnapshotNode>
 ): void {
-  const { importedProfilingData } = store;
-  const profilingSnapshot =
-    importedProfilingData != null
-      ? importedProfilingData.profilingSnapshot
-      : store.profilingSnapshot;
   const node = profilingSnapshot.get(id);
   if (node != null) {
     nodes.set(id, {
@@ -146,7 +152,7 @@ function recursivelyInitializeTree(
         id,
         nodes,
         initialTreeBaseDurations,
-        store
+        profilingSnapshot
       )
     );
   }

src/devtools/views/Profiler/ProfilingImportExportButtons.js

@@ -37,7 +37,7 @@ export default function ProfilingImportExportButtons() {
     const queue = [rootID];
     while (queue.length) {
       const id = queue.pop();
-      profilingSnapshotForRoot.push([id, store.profilingSnapshot.get(id)]);
+      profilingSnapshotForRoot.push([id, store.profilingSnapshots.get(id)]);
     }
 
     bridge.send('exportProfilingSummary', {
@@ -51,7 +51,7 @@ export default function ProfilingImportExportButtons() {
     rendererID,
     rootID,
     store.profilingOperations,
-    store.profilingSnapshot,
+    store.profilingSnapshots,
   ]);
 
   const uploadData = useCallback(() => {

src/devtools/views/Profiler/types.js

@@ -65,7 +65,7 @@ export type ProfilingSnapshotNode = {|
 export type ImportedProfilingData = {|
   version: number,
   profilingOperations: Map<number, Array<Uint32Array>>,
-  profilingSnapshot: Map<number, ProfilingSnapshotNode>,
+  profilingSnapshots: Map<number, Map<number, ProfilingSnapshotNode>>,
   commitDetails: CommitDetailsFrontend,
   interactions: InteractionsFrontend,
   profilingSummary: ProfilingSummaryFrontend,