Clear profiling data for root when the root is unmounted

Fixes https://github.com/bvaughn/react-devtools-experimental/issues/232

src/__tests__/__snapshots__/profiler-test.js.snap

@@ -0,0 +1,31 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`Profiler should start and stop profiling, handle root unmounting: 1: mount 1`] = `
+[root]
+  ▸ <Parent key="A">
+[root]
+  ▸ <Parent key="B">
+`;
+
+exports[`Profiler should start and stop profiling, handle root unmounting: 2: profiling started 1`] = `
+[root]
+  ▸ <Parent key="A">
+[root]
+  ▸ <Parent key="B">
+`;
+
+exports[`Profiler should start and stop profiling, handle root unmounting: 3: update 1`] = `
+[root]
+  ▸ <Parent key="A">
+[root]
+  ▸ <Parent key="B">
+`;
+
+exports[`Profiler should start and stop profiling, handle root unmounting: 4: unmount B 1`] = `
+[root]
+  ▸ <Parent key="A">
+`;
+
+exports[`Profiler should start and stop profiling, handle root unmounting: 5: unmount A 1`] = ``;
+
+exports[`Profiler should start and stop profiling, handle root unmounting: 6: profiling stopped 1`] = ``;

src/__tests__/profiler-test.js

@@ -0,0 +1,60 @@
+// @flow
+
+describe('Profiler', () => {
+  let React;
+  let ReactDOM;
+  let TestUtils;
+  let store;
+
+  const act = (callback: Function) => {
+    TestUtils.act(() => {
+      callback();
+    });
+    jest.runAllTimers(); // Flush Bridge operations
+  };
+
+  beforeEach(() => {
+    store = global.store;
+
+    React = require('react');
+    ReactDOM = require('react-dom');
+    TestUtils = require('react-dom/test-utils');
+  });
+
+  it('should start and stop profiling, handle root unmounting', async () => {
+    const Parent = ({ count }) =>
+      new Array(count).fill(true).map((_, index) => <Child key={index} />);
+    const Child = () => <div>Hi!</div>;
+
+    const containerA = document.createElement('div');
+    const containerB = document.createElement('div');
+
+    act(() => {
+      ReactDOM.render(<Parent key="A" count={3} />, containerA);
+      ReactDOM.render(<Parent key="B" count={2} />, containerB);
+    });
+    expect(store).toMatchSnapshot('1: mount');
+
+    act(() => {
+      store.startProfiling();
+    });
+    expect(store).toMatchSnapshot('2: profiling started');
+
+    act(() => {
+      ReactDOM.render(<Parent key="A" count={4} />, containerA);
+      ReactDOM.render(<Parent key="B" count={1} />, containerB);
+    });
+    expect(store).toMatchSnapshot('3: update');
+
+    act(() => ReactDOM.unmountComponentAtNode(containerB));
+    expect(store).toMatchSnapshot('4: unmount B');
+
+    act(() => ReactDOM.unmountComponentAtNode(containerA));
+    expect(store).toMatchSnapshot('5: unmount A');
+
+    act(() => {
+      store.stopProfiling();
+    });
+    expect(store).toMatchSnapshot('6: profiling stopped');
+  });
+});

src/devtools/store.js

@@ -185,27 +185,31 @@ export default class Store extends EventEmitter {
 
   assertEmptyMaps() {
     // This is only used in tests to avoid memory leaks.
-    if (this._idToElement.size !== 0) {
-      throw new Error('Expected _idToElement to be empty.');
-    }
-    if (this._ownersMap.size !== 0) {
-      throw new Error('Expected _ownersMap to be empty.');
-    }
-    if (this._profilingOperationsByRootID.size !== 0) {
-      throw new Error('Expected _profilingOperationsByRootID to be empty.');
-    }
-    if (this._profilingScreenshotsByRootID.size !== 0) {
-      throw new Error('Expected _profilingScreenshotsByRootID to be empty.');
-    }
-    if (this._profilingSnapshotsByElementID.size !== 0) {
-      throw new Error('Expected _profilingSnapshotsByElementID to be empty.');
-    }
-    if (this._rootIDToCapabilities.size !== 0) {
-      throw new Error('Expected _rootIDToCapabilities to be empty.');
-    }
-    if (this._rootIDToRendererID.size !== 0) {
-      throw new Error('Expected _rootIDToRendererID to be empty.');
-    }
+    const assertOneMap = (mapName, map) => {
+      if (map.size !== 0) {
+        throw new Error(
+          `Expected ${mapName} to be empty, got ${
+            map.size
+          }: ${require('util').inspect(this, { depth: 20 })}`
+        );
+      }
+    };
+    assertOneMap('_idToElement', this._idToElement);
+    assertOneMap('_ownersMap', this._ownersMap);
+    assertOneMap(
+      '_profilingOperationsByRootID',
+      this._profilingOperationsByRootID
+    );
+    assertOneMap(
+      '_profilingScreenshotsByRootID',
+      this._profilingScreenshotsByRootID
+    );
+    assertOneMap(
+      '_profilingSnapshotsByElementID',
+      this._profilingSnapshotsByElementID
+    );
+    assertOneMap('_rootIDToCapabilities', this._rootIDToCapabilities);
+    assertOneMap('_rootIDToRendererID', this._rootIDToRendererID);
   }
 
   get captureScreenshots(): boolean {
@@ -665,6 +669,15 @@ export default class Store extends EventEmitter {
     }
   };
 
+  _clearProfilingSnapshotRecursive = (elementID: number) => {
+    const element = this.getElementByID(elementID);
+    if (element !== null) {
+      this._profilingSnapshotsByElementID.delete(elementID);
+
+      element.children.forEach(this._clearProfilingSnapshotRecursive);
+    }
+  };
+
   _adjustParentTreeWeight = (
     parentElement: Element | null,
     weightDelta: number
@@ -876,6 +889,11 @@ export default class Store extends EventEmitter {
               throw new Error(`Node ${id} was removed before its children.`);
             }
 
+            // 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._idToElement.delete(id);
 
             let parentElement = null;
@@ -888,6 +906,9 @@ export default class Store extends EventEmitter {
               this._rootIDToRendererID.delete(id);
               this._rootIDToCapabilities.delete(id);
 
+              this._profilingOperationsByRootID.delete(id);
+              this._profilingScreenshotsByRootID.delete(id);
+
               haveRootsChanged = true;
             } else {
               if (__DEBUG__) {