DevTools: Fix Fiber leak caused by Refresh hot reloading

packages/react-devtools-shared/src/__tests__/FastRefreshDevToolsIntegration-test.js

@@ -0,0 +1,236 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+describe('Fast Refresh', () => {
+  let React;
+  let ReactDOM;
+  let ReactFreshRuntime;
+  let act;
+  let babel;
+  let container;
+  let exportsObj;
+  let freshPlugin;
+  let store;
+  let withErrorsOrWarningsIgnored;
+
+  afterEach(() => {
+    jest.resetModules();
+  });
+
+  beforeEach(() => {
+    exportsObj = undefined;
+    container = document.createElement('div');
+
+    babel = require('@babel/core');
+    freshPlugin = require('react-refresh/babel');
+
+    store = global.store;
+
+    React = require('react');
+
+    ReactFreshRuntime = require('react-refresh/runtime');
+    ReactFreshRuntime.injectIntoGlobalHook(global);
+
+    ReactDOM = require('react-dom');
+
+    const utils = require('./utils');
+    act = utils.act;
+    withErrorsOrWarningsIgnored = utils.withErrorsOrWarningsIgnored;
+  });
+
+  function execute(source) {
+    const compiled = babel.transform(source, {
+      babelrc: false,
+      presets: ['@babel/react'],
+      plugins: [
+        [freshPlugin, {skipEnvCheck: true}],
+        '@babel/plugin-transform-modules-commonjs',
+        '@babel/plugin-transform-destructuring',
+      ].filter(Boolean),
+    }).code;
+    exportsObj = {};
+    // eslint-disable-next-line no-new-func
+    new Function(
+      'global',
+      'React',
+      'exports',
+      '$RefreshReg$',
+      '$RefreshSig$',
+      compiled,
+    )(global, React, exportsObj, $RefreshReg$, $RefreshSig$);
+    // Module systems will register exports as a fallback.
+    // This is useful for cases when e.g. a class is exported,
+    // and we don't want to propagate the update beyond this module.
+    $RefreshReg$(exportsObj.default, 'exports.default');
+    return exportsObj.default;
+  }
+
+  function render(source) {
+    const Component = execute(source);
+    act(() => {
+      ReactDOM.render(<Component />, container);
+    });
+    // Module initialization shouldn't be counted as a hot update.
+    expect(ReactFreshRuntime.performReactRefresh()).toBe(null);
+  }
+
+  function patch(source) {
+    const prevExports = exportsObj;
+    execute(source);
+    const nextExports = exportsObj;
+
+    // Check if exported families have changed.
+    // (In a real module system we'd do this for *all* exports.)
+    // For example, this can happen if you convert a class to a function.
+    // Or if you wrap something in a HOC.
+    const didExportsChange =
+      ReactFreshRuntime.getFamilyByType(prevExports.default) !==
+      ReactFreshRuntime.getFamilyByType(nextExports.default);
+    if (didExportsChange) {
+      // In a real module system, we would propagate such updates upwards,
+      // and re-execute modules that imported this one. (Just like if we edited them.)
+      // This makes adding/removing/renaming exports re-render references to them.
+      // Here, we'll just force a re-render using the newer type to emulate this.
+      const NextComponent = nextExports.default;
+      act(() => {
+        ReactDOM.render(<NextComponent />, container);
+      });
+    }
+    act(() => {
+      const result = ReactFreshRuntime.performReactRefresh();
+      if (!didExportsChange) {
+        // Normally we expect that some components got updated in our tests.
+        expect(result).not.toBe(null);
+      } else {
+        // However, we have tests where we convert functions to classes,
+        // and in those cases it's expected nothing would get updated.
+        // (Instead, the export change branch above would take care of it.)
+      }
+    });
+    expect(ReactFreshRuntime._getMountedRootCount()).toBe(1);
+  }
+
+  function $RefreshReg$(type, id) {
+    ReactFreshRuntime.register(type, id);
+  }
+
+  function $RefreshSig$() {
+    return ReactFreshRuntime.createSignatureFunctionForTransform();
+  }
+
+  it('should not break the DevTools store', () => {
+    render(`
+      function Parent() {
+        return <Child key="A" />;
+      };
+
+      function Child() {
+        return null;
+      };
+
+      export default Parent;
+    `);
+
+    expect(store).toMatchInlineSnapshot(`
+      [root]
+        ▾ <Parent>
+            <Child key="A">
+    `);
+
+    patch(`
+      function Parent() {
+        return <Child key="B" />;
+      };
+
+      function Child() {
+        return null;
+      };
+
+      export default Parent;
+    `);
+    expect(store).toMatchInlineSnapshot(`
+      [root]
+        ▾ <Parent>
+            <Child key="B">
+    `);
+  });
+
+  it('should not break when there are warnings in between patching', () => {
+    withErrorsOrWarningsIgnored(['Expected warning during render'], () => {
+      render(`
+      const {useState} = React;
+
+      export default function Component() {
+        const [state, setState] = useState(1);
+
+        console.warn("Expected warning during render");
+
+        return <div />;
+      }
+    `);
+    });
+    expect(store).toMatchInlineSnapshot(`
+      ✕ 0, ⚠ 1
+      [root]
+          <Component> ⚠
+    `);
+
+    let element = container.firstChild;
+
+    withErrorsOrWarningsIgnored(['Expected warning during render'], () => {
+      patch(`
+      const {useState} = React;
+
+      export default function Component() {
+        const [state, setState] = useState(1);
+
+        console.warn("Expected warning during render");
+
+        return <div id="one" />;
+      }
+    `);
+    });
+
+    // This is the same component type, so the warning count carries over.
+    expect(store).toMatchInlineSnapshot(`
+      ✕ 0, ⚠ 2
+      [root]
+          <Component> ⚠
+    `);
+
+    // State is preserved; this verifies that Fast Refresh is wired up.
+    expect(container.firstChild).toBe(element);
+    element = container.firstChild;
+
+    withErrorsOrWarningsIgnored(['Expected warning during render'], () => {
+      patch(`
+      const {useEffect, useState} = React;
+
+      export default function Component() {
+        const [state, setState] = useState(3);
+        useEffect(() => {});
+
+        console.warn("Expected warning during render");
+
+        return <div id="one" />;
+      }
+    `);
+    });
+
+    // This is a new component type, so the warning count has been reset.
+    expect(store).toMatchInlineSnapshot(`
+      ✕ 0, ⚠ 1
+      [root]
+          <Component> ⚠
+    `);
+
+    // State is reset because hooks changed.
+    expect(container.firstChild).not.toBe(element);
+  });
+});

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

@@ -1007,6 +1007,11 @@ export function attach(
   const getProfilingData = () => {
     throw new Error('getProfilingData not supported by this renderer');
   };
+  const handleClonedForForceRemount = () => {
+    throw new Error(
+      'handleClonedForForceRemount not supported by this renderer',
+    );
+  };
   const handleCommitFiberRoot = () => {
     throw new Error('handleCommitFiberRoot not supported by this renderer');
   };
@@ -1084,6 +1089,7 @@ export function attach(
     getOwnersList,
     getPathForElement,
     getProfilingData,
+    handleClonedForForceRemount,
     handleCommitFiberRoot,
     handleCommitFiberUnmount,
     handlePostCommitFiberRoot,

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

@@ -638,6 +638,10 @@ export function attach(
     // We should clean up Fibers like this when flushing; see recordPendingErrorsAndWarnings().
     const fiberID = getFiberID(getPrimaryFiber(fiber));
 
+    if (__DEBUG__) {
+      debug('onErrorOrWarning', fiber, null, `${type}: "${message}"`);
+    }
+
     // Mark this Fiber as needed its warning/error count updated during the next flush.
     fibersWithChangedErrorOrWarningCounts.add(fiberID);
 
@@ -702,16 +706,15 @@ export function attach(
     if (__DEBUG__) {
       const displayName =
         fiber.tag + ':' + (getDisplayNameForFiber(fiber) || 'null');
-      const id = getFiberID(fiber);
+
+      const id = getFiberIDSafe(fiber) || '-';
       const parentDisplayName = parentFiber
         ? parentFiber.tag +
           ':' +
           (getDisplayNameForFiber(parentFiber) || 'null')
         : '';
-      const parentID = parentFiber ? getFiberID(parentFiber) : '';
-      // NOTE: calling getFiberID or getPrimaryFiber is unsafe here
-      // because it will put them in the map. For now, we'll omit them.
-      // TODO: better debugging story for this.
+      const parentID = parentFiber ? getFiberIDSafe(parentFiber) || '-' : '';
+
       console.log(
         `[renderer] %c${name} %c${displayName} (${id}) %c${
           parentFiber ? `${parentDisplayName} (${parentID})` : ''
@@ -979,6 +982,23 @@ export function attach(
   // When a mount or update is in progress, this value tracks the root that is being operated on.
   let currentRootID: number = -1;
 
+  function getFiberIDSafe(fiber: Fiber): number | null {
+    let primaryFiber = null;
+    if (primaryFibers.has(fiber)) {
+      primaryFiber = fiber;
+    }
+    const {alternate} = fiber;
+    if (alternate != null && primaryFibers.has(alternate)) {
+      primaryFiber = alternate;
+    }
+
+    if (primaryFiber && fiberToIDMap.has(primaryFiber)) {
+      return ((fiberToIDMap.get(primaryFiber): any): number);
+    }
+
+    return null;
+  }
+
   function getFiberID(primaryFiber: Fiber): number {
     if (!fiberToIDMap.has(primaryFiber)) {
       const id = getUID();
@@ -1641,6 +1661,7 @@ export function attach(
         pendingRealUnmountedIDs.push(id);
       }
     }
+
     fiberToIDMap.delete(primaryFiber);
     idToFiberMap.delete(id);
     primaryFibers.delete(primaryFiber);
@@ -3815,6 +3836,41 @@ export function attach(
     traceUpdatesEnabled = isEnabled;
   }
 
+  function handleClonedForForceRemount(oldFiber: Fiber, newFiber: Fiber): void {
+    if (__DEBUG__) {
+      console.log(
+        '[renderer] handleClonedForForceRemount',
+        getDisplayNameForFiber(oldFiber),
+        '(' + getFiberID(getPrimaryFiber(oldFiber)),
+        '->',
+        getFiberID(getPrimaryFiber(newFiber)) + ')',
+      );
+    }
+
+    let primaryFiber = null;
+    if (primaryFibers.has(oldFiber)) {
+      primaryFiber = oldFiber;
+    }
+    const {alternate} = oldFiber;
+    if (alternate != null && primaryFibers.has(alternate)) {
+      primaryFiber = alternate;
+    }
+
+    // Fast Refresh is about to (synchronously) clone and replace this part of the tree.
+    // In order to avoid these Fibers from leaking on the DevTools side,
+    // and possibly remaining visible in the Components tab as well,
+    // queue up unmount operations to send on the next commit.
+    if (primaryFiber) {
+      recordUnmount(primaryFiber, false);
+      unmountFiberChildrenRecursively(primaryFiber);
+
+      const fiberID = ((fiberToIDMap.get(primaryFiber): any): number);
+      fiberToIDMap.delete(primaryFiber);
+      idToFiberMap.delete(fiberID);
+      primaryFibers.delete(primaryFiber);
+    }
+  }
+
   return {
     cleanup,
     clearErrorsAndWarnings,
@@ -3831,6 +3887,7 @@ export function attach(
     getOwnersList,
     getPathForElement,
     getProfilingData,
+    handleClonedForForceRemount,
     handleCommitFiberRoot,
     handleCommitFiberUnmount,
     handlePostCommitFiberRoot,

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

@@ -323,6 +323,7 @@ export type RendererInterface = {
   getProfilingData(): ProfilingDataBackend,
   getOwnersList: (id: number) => Array<SerializedElement> | null,
   getPathForElement: (id: number) => Array<PathFrame> | null,
+  handleClonedForForceRemount: (oldFiber: Fiber, newFiber: Fiber) => void,
   handleCommitFiberRoot: (fiber: Object, commitPriority?: number) => void,
   handleCommitFiberUnmount: (fiber: Object) => void,
   handlePostCommitFiberRoot: (fiber: Object) => void,
@@ -396,5 +397,12 @@ export type DevToolsHook = {
     // Added in v16.9 to support Fast Refresh
     didError?: boolean,
   ) => void,
+
+  // Added in v17.x to improve Fast Refresh + DevTools integration
+  onClonedForForceRemount: (
+    rendererID: RendererID,
+    oldFiber: Fiber,
+    newFiber: Fiber,
+  ) => void,
   ...
 };