[Fresh] Retry failed roots on refresh

packages/react-reconciler/src/ReactFiberDevToolsHook.js

@@ -15,6 +15,7 @@ import type {Fiber} from './ReactFiber';
 import type {FiberRoot} from './ReactFiberRoot';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 
+import {DidCapture} from 'shared/ReactSideEffectTags';
 import warningWithoutStack from 'shared/warningWithoutStack';
 
 declare var __REACT_DEVTOOLS_GLOBAL_HOOK__: Object | void;
@@ -55,15 +56,16 @@ export function injectInternals(internals: Object): boolean {
     // We have successfully injected, so now it is safe to set up hooks.
     onCommitFiberRoot = (root, expirationTime) => {
       try {
+        const didError = (root.current.effectTag & DidCapture) === DidCapture;
         if (enableProfilerTimer) {
           const currentTime = requestCurrentTime();
           const priorityLevel = inferPriorityFromExpirationTime(
             currentTime,
             expirationTime,
           );
-          hook.onCommitFiberRoot(rendererID, root, priorityLevel);
+          hook.onCommitFiberRoot(rendererID, root, priorityLevel, didError);
         } else {
-          hook.onCommitFiberRoot(rendererID, root);
+          hook.onCommitFiberRoot(rendererID, root, undefined, didError);
         }
       } catch (err) {
         if (__DEV__ && !hasLoggedError) {

packages/react-reconciler/src/ReactFiberHotReloading.js

@@ -11,12 +11,15 @@ import type {ReactElement} from 'shared/ReactElementType';
 import type {Fiber} from './ReactFiber';
 import type {FiberRoot} from './ReactFiberRoot';
 import type {Instance} from './ReactFiberHostConfig';
+import type {ReactNodeList} from 'shared/ReactTypes';
 
 import {
   flushSync,
   scheduleWork,
   flushPassiveEffects,
 } from './ReactFiberWorkLoop';
+import {updateContainerAtExpirationTime} from './ReactFiberReconciler';
+import {emptyContextObject} from './ReactFiberContext';
 import {Sync} from './ReactFiberExpirationTime';
 import {
   ClassComponent,
@@ -49,6 +52,7 @@ type RefreshHandler = any => Family | void;
 // Used by React Refresh runtime through DevTools Global Hook.
 export type SetRefreshHandler = (handler: RefreshHandler | null) => void;
 export type ScheduleRefresh = (root: FiberRoot, update: RefreshUpdate) => void;
+export type ScheduleRoot = (root: FiberRoot, element: ReactNodeList) => void;
 export type FindHostInstancesForRefresh = (
   root: FiberRoot,
   families: Array<Family>,
@@ -242,6 +246,22 @@ export let scheduleRefresh: ScheduleRefresh = (
   }
 };
 
+export let scheduleRoot: ScheduleRoot = (
+  root: FiberRoot,
+  element: ReactNodeList,
+): void => {
+  if (__DEV__) {
+    if (root.context !== emptyContextObject) {
+      // Super edge case: root has a legacy _renderSubtree context
+      // but we don't know the parentComponent so we can't pass it.
+      // Just ignore. We'll delete this with _renderSubtree code path later.
+      return;
+    }
+    flushPassiveEffects();
+    updateContainerAtExpirationTime(element, root, null, Sync, null);
+  }
+};
+
 function scheduleFibersWithFamiliesRecursively(
   fiber: Fiber,
   updatedFamilies: Set<Family>,

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -72,6 +72,7 @@ import {revertPassiveEffectsChange} from 'shared/ReactFeatureFlags';
 import {requestCurrentSuspenseConfig} from './ReactFiberSuspenseConfig';
 import {
   scheduleRefresh,
+  scheduleRoot,
   setRefreshHandler,
   findHostInstancesForRefresh,
 } from './ReactFiberHotReloading';
@@ -498,6 +499,7 @@ export function injectIntoDevTools(devToolsConfig: DevToolsConfig): boolean {
     // React Refresh
     findHostInstancesForRefresh: __DEV__ ? findHostInstancesForRefresh : null,
     scheduleRefresh: __DEV__ ? scheduleRefresh : null,
+    scheduleRoot: __DEV__ ? scheduleRoot : null,
     setRefreshHandler: __DEV__ ? setRefreshHandler : null,
   });
 }

packages/react-refresh/src/ReactFreshRuntime.js

@@ -13,9 +13,11 @@ import type {
   Family,
   RefreshUpdate,
   ScheduleRefresh,
+  ScheduleRoot,
   FindHostInstancesForRefresh,
   SetRefreshHandler,
 } from 'react-reconciler/src/ReactFiberHotReloading';
+import type {ReactNodeList} from 'shared/ReactTypes';
 
 import {REACT_MEMO_TYPE, REACT_FORWARD_REF_TYPE} from 'shared/ReactSymbols';
 import warningWithoutStack from 'shared/warningWithoutStack';
@@ -57,9 +59,13 @@ let pendingUpdates: Array<[Family, any]> = [];
 // This is injected by the renderer via DevTools global hook.
 let setRefreshHandler: null | SetRefreshHandler = null;
 let scheduleRefresh: null | ScheduleRefresh = null;
+let scheduleRoot: null | ScheduleRoot = null;
 let findHostInstancesForRefresh: null | FindHostInstancesForRefresh = null;
 
-let mountedRoots = new Set();
+// We keep track of mounted roots so we can schedule updates.
+let mountedRoots: Set<FiberRoot> = new Set();
+// If a root captures an error, we add its element to this Map so we can retry on edit.
+let failedRoots: Map<FiberRoot, ReactNodeList> = new Map();
 
 function computeFullKey(signature: Signature): string {
   if (signature.fullKey !== null) {
@@ -196,14 +202,36 @@ export function performReactRefresh(): RefreshUpdate | null {
       );
       return null;
     }
+    if (typeof scheduleRoot !== 'function') {
+      warningWithoutStack(
+        false,
+        'Could not find the scheduleRoot() implementation. ' +
+          'This likely means that injectIntoGlobalHook() was either ' +
+          'called before the global DevTools hook was set up, or after the ' +
+          'renderer has already initialized. Please file an issue with a reproducing case.',
+      );
+      return null;
+    }
     const scheduleRefreshForRoot = scheduleRefresh;
+    const scheduleRenderForRoot = scheduleRoot;
 
     // Even if there are no roots, set the handler on first update.
     // This ensures that if *new* roots are mounted, they'll use the resolve handler.
     setRefreshHandler(resolveFamily);
 
     let didError = false;
     let firstError = null;
+    failedRoots.forEach((element, root) => {
+      try {
+        scheduleRenderForRoot(root, element);
+      } catch (err) {
+        if (!didError) {
+          didError = true;
+          firstError = err;
+        }
+        // Keep trying other roots.
+      }
+    });
     mountedRoots.forEach(root => {
       try {
         scheduleRefreshForRoot(root, update);
@@ -245,7 +273,7 @@ export function register(type: any, id: string): void {
 
     // Create family or remember to update it.
     // None of this bookkeeping affects reconciliation
-    // until the first prepareUpdate() call above.
+    // until the first performReactRefresh() call above.
     let family = allFamiliesByID.get(id);
     if (family === undefined) {
       family = {current: type};
@@ -362,7 +390,12 @@ export function injectIntoGlobalHook(globalObject: any): void {
       globalObject.__REACT_DEVTOOLS_GLOBAL_HOOK__ = hook = {
         supportsFiber: true,
         inject() {},
-        onCommitFiberRoot(id: mixed, root: FiberRoot) {},
+        onCommitFiberRoot(
+          id: mixed,
+          root: FiberRoot,
+          maybePriorityLevel: mixed,
+          didError: boolean,
+        ) {},
         onCommitFiberUnmount() {},
       };
     }
@@ -373,14 +406,20 @@ export function injectIntoGlobalHook(globalObject: any): void {
       findHostInstancesForRefresh = ((injected: any)
         .findHostInstancesForRefresh: FindHostInstancesForRefresh);
       scheduleRefresh = ((injected: any).scheduleRefresh: ScheduleRefresh);
+      scheduleRoot = ((injected: any).scheduleRoot: ScheduleRoot);
       setRefreshHandler = ((injected: any)
         .setRefreshHandler: SetRefreshHandler);
       return oldInject.apply(this, arguments);
     };
 
     // We also want to track currently mounted roots.
     const oldOnCommitFiberRoot = hook.onCommitFiberRoot;
-    hook.onCommitFiberRoot = function(id: mixed, root: FiberRoot) {
+    hook.onCommitFiberRoot = function(
+      id: mixed,
+      root: FiberRoot,
+      maybePriorityLevel: mixed,
+      didError: boolean,
+    ) {
       const current = root.current;
       const alternate = current.alternate;
 
@@ -399,12 +438,18 @@ export function injectIntoGlobalHook(globalObject: any): void {
         if (!wasMounted && isMounted) {
           // Mount a new root.
           mountedRoots.add(root);
+          failedRoots.delete(root);
         } else if (wasMounted && isMounted) {
           // Update an existing root.
           // This doesn't affect our mounted root Set.
         } else if (wasMounted && !isMounted) {
           // Unmount an existing root.
           mountedRoots.delete(root);
+          if (didError) {
+            // We'll remount it on future edits.
+            // Remember what was rendered so we can restore it.
+            failedRoots.set(root, alternate.memoizedState.element);
+          }
         }
       } else {
         // Mount a new root.

packages/react-refresh/src/__tests__/ReactFresh-test.js

@@ -2707,6 +2707,96 @@ describe('ReactFresh', () => {
     }
   });
 
+  it('remounts a failed root on update', () => {
+    if (__DEV__) {
+      render(() => {
+        function Hello() {
+          return <h1>Hi</h1>;
+        }
+        $RefreshReg$(Hello, 'Hello');
+
+        return Hello;
+      });
+      expect(container.innerHTML).toBe('<h1>Hi</h1>');
+
+      // Perform a hot update that fails.
+      // This removes the root.
+      expect(() => {
+        patch(() => {
+          function Hello() {
+            throw new Error('No');
+          }
+          $RefreshReg$(Hello, 'Hello');
+        });
+      }).toThrow('No');
+      expect(container.innerHTML).toBe('');
+
+      // A bad retry
+      expect(() => {
+        patch(() => {
+          function Hello() {
+            throw new Error('Not yet');
+          }
+          $RefreshReg$(Hello, 'Hello');
+        });
+      }).toThrow('Not yet');
+      expect(container.innerHTML).toBe('');
+
+      // Perform a hot update that fixes the error.
+      patch(() => {
+        function Hello() {
+          return <h1>Fixed!</h1>;
+        }
+        $RefreshReg$(Hello, 'Hello');
+      });
+      // This should remount the root.
+      expect(container.innerHTML).toBe('<h1>Fixed!</h1>');
+
+      // Verify next hot reload doesn't remount anything.
+      let helloNode = container.firstChild;
+      patch(() => {
+        function Hello() {
+          return <h1>Nice.</h1>;
+        }
+        $RefreshReg$(Hello, 'Hello');
+      });
+      expect(container.firstChild).toBe(helloNode);
+      expect(helloNode.textContent).toBe('Nice.');
+
+      // Break again.
+      expect(() => {
+        patch(() => {
+          function Hello() {
+            throw new Error('Oops');
+          }
+          $RefreshReg$(Hello, 'Hello');
+        });
+      }).toThrow('Oops');
+      expect(container.innerHTML).toBe('');
+
+      // Perform a hot update that fixes the error.
+      patch(() => {
+        function Hello() {
+          return <h1>At last.</h1>;
+        }
+        $RefreshReg$(Hello, 'Hello');
+      });
+      // This should remount the root.
+      expect(container.innerHTML).toBe('<h1>At last.</h1>');
+
+      // Check we don't attempt to reverse an intentional unmount.
+      ReactDOM.unmountComponentAtNode(container);
+      expect(container.innerHTML).toBe('');
+      patch(() => {
+        function Hello() {
+          return <h1>Never mind me!</h1>;
+        }
+        $RefreshReg$(Hello, 'Hello');
+      });
+      expect(container.innerHTML).toBe('');
+    }
+  });
+
   it('remounts classes on every edit', () => {
     if (__DEV__) {
       let HelloV1 = render(() => {