From 237a966da058fa063abd2e1365e90f2a26aa4e51 Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@gmail.com>
Date: Thu, 21 Nov 2019 14:10:26 +0000
Subject: [PATCH] [Fresh] Fix an infinite loop in an edge case (#17414)

* [Fresh] Fix an infinite loop in an edge case

* Make it work in IE11
---
 .../react-refresh/src/ReactFreshRuntime.js    | 33 +++++++++--
 .../src/__tests__/ReactFresh-test.js          | 55 +++++++++++++++++++
 2 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/packages/react-refresh/src/ReactFreshRuntime.js b/packages/react-refresh/src/ReactFreshRuntime.js
index 1b40d1fa6d..65a8f73657 100644
--- a/packages/react-refresh/src/ReactFreshRuntime.js
+++ b/packages/react-refresh/src/ReactFreshRuntime.js
@@ -154,6 +154,22 @@ function resolveFamily(type) {
   return updatedFamiliesByType.get(type);
 }
 
+// If we didn't care about IE11, we could use new Map/Set(iterable).
+function cloneMap<K, V>(map: Map<K, V>): Map<K, V> {
+  let clone = new Map();
+  map.forEach((value, key) => {
+    clone.set(key, value);
+  });
+  return clone;
+}
+function cloneSet<T>(set: Set<T>): Set<T> {
+  let clone = new Set();
+  set.forEach(value => {
+    clone.add(value);
+  });
+  return clone;
+}
+
 export function performReactRefresh(): RefreshUpdate | null {
   if (__DEV__) {
     if (pendingUpdates.length === 0) {
@@ -195,8 +211,17 @@ export function performReactRefresh(): RefreshUpdate | null {
 
     let didError = false;
     let firstError = null;
-    failedRoots.forEach((element, root) => {
-      const helpers = helpersByRoot.get(root);
+
+    // We snapshot maps and sets that are mutated during commits.
+    // If we don't do this, there is a risk they will be mutated while
+    // we iterate over them. For example, trying to recover a failed root
+    // may cause another root to be added to the failed list -- an infinite loop.
+    let failedRootsSnapshot = cloneMap(failedRoots);
+    let mountedRootsSnapshot = cloneSet(mountedRoots);
+    let helpersByRootSnapshot = cloneMap(helpersByRoot);
+
+    failedRootsSnapshot.forEach((element, root) => {
+      const helpers = helpersByRootSnapshot.get(root);
       if (helpers === undefined) {
         throw new Error(
           'Could not find helpers for a root. This is a bug in React Refresh.',
@@ -212,8 +237,8 @@ export function performReactRefresh(): RefreshUpdate | null {
         // Keep trying other roots.
       }
     });
-    mountedRoots.forEach(root => {
-      const helpers = helpersByRoot.get(root);
+    mountedRootsSnapshot.forEach(root => {
+      const helpers = helpersByRootSnapshot.get(root);
       if (helpers === undefined) {
         throw new Error(
           'Could not find helpers for a root. This is a bug in React Refresh.',
diff --git a/packages/react-refresh/src/__tests__/ReactFresh-test.js b/packages/react-refresh/src/__tests__/ReactFresh-test.js
index ae3e598a1b..046cd28669 100644
--- a/packages/react-refresh/src/__tests__/ReactFresh-test.js
+++ b/packages/react-refresh/src/__tests__/ReactFresh-test.js
@@ -2909,6 +2909,61 @@ describe('ReactFresh', () => {
     }
   });
 
+  it('regression test: does not get into an infinite loop', () => {
+    if (__DEV__) {
+      let containerA = document.createElement('div');
+      let containerB = document.createElement('div');
+
+      // Initially, nothing interesting.
+      let RootAV1 = () => {
+        return 'A1';
+      };
+      $RefreshReg$(RootAV1, 'RootA');
+      let RootBV1 = () => {
+        return 'B1';
+      };
+      $RefreshReg$(RootBV1, 'RootB');
+
+      act(() => {
+        ReactDOM.render(<RootAV1 />, containerA);
+        ReactDOM.render(<RootBV1 />, containerB);
+      });
+      expect(containerA.innerHTML).toBe('A1');
+      expect(containerB.innerHTML).toBe('B1');
+
+      // Then make the first root fail.
+      let RootAV2 = () => {
+        throw new Error('A2!');
+      };
+      $RefreshReg$(RootAV2, 'RootA');
+      expect(() => ReactFreshRuntime.performReactRefresh()).toThrow('A2!');
+      expect(containerA.innerHTML).toBe('');
+      expect(containerB.innerHTML).toBe('B1');
+
+      // Then patch the first root, but make it fail in the commit phase.
+      // This used to trigger an infinite loop due to a list of failed roots
+      // being mutated while it was being iterated on.
+      let RootAV3 = () => {
+        React.useLayoutEffect(() => {
+          throw new Error('A3!');
+        }, []);
+        return 'A3';
+      };
+      $RefreshReg$(RootAV3, 'RootA');
+      expect(() => ReactFreshRuntime.performReactRefresh()).toThrow('A3!');
+      expect(containerA.innerHTML).toBe('');
+      expect(containerB.innerHTML).toBe('B1');
+
+      let RootAV4 = () => {
+        return 'A4';
+      };
+      $RefreshReg$(RootAV4, 'RootA');
+      ReactFreshRuntime.performReactRefresh();
+      expect(containerA.innerHTML).toBe('A4');
+      expect(containerB.innerHTML).toBe('B1');
+    }
+  });
+
   it('remounts classes on every edit', () => {
     if (__DEV__) {
       let HelloV1 = render(() => {