Allow DevTools to toggle Suspense fallbacks (#15232)

* Allow DevTools to toggle Suspense state

* Change API to overrideSuspense

This lets detect support for overriding Suspense from DevTools.

* Add ConcurrentMode test

* Newlines

* Remove unnecessary change

* Naming changes

Author: gaearon

packages/react-debug-tools/src/__tests__/ReactDevToolsHooksIntegration-test.js

@@ -14,13 +14,18 @@ describe('React hooks DevTools integration', () => {
   let React;
   let ReactDebugTools;
   let ReactTestRenderer;
+  let Scheduler;
   let act;
   let overrideHookState;
+  let scheduleUpdate;
+  let setSuspenseHandler;
 
   beforeEach(() => {
     global.__REACT_DEVTOOLS_GLOBAL_HOOK__ = {
       inject: injected => {
         overrideHookState = injected.overrideHookState;
+        scheduleUpdate = injected.scheduleUpdate;
+        setSuspenseHandler = injected.setSuspenseHandler;
       },
       supportsFiber: true,
       onCommitFiberRoot: () => {},
@@ -32,6 +37,7 @@ describe('React hooks DevTools integration', () => {
     React = require('react');
     ReactDebugTools = require('react-debug-tools');
     ReactTestRenderer = require('react-test-renderer');
+    Scheduler = require('scheduler');
 
     act = ReactTestRenderer.act;
   });
@@ -173,4 +179,112 @@ describe('React hooks DevTools integration', () => {
       });
     }
   });
+
+  it('should support overriding suspense in sync mode', () => {
+    if (__DEV__) {
+      // Lock the first render
+      setSuspenseHandler(() => true);
+    }
+
+    function MyComponent() {
+      return 'Done';
+    }
+
+    const renderer = ReactTestRenderer.create(
+      <div>
+        <React.Suspense fallback={'Loading'}>
+          <MyComponent />
+        </React.Suspense>
+      </div>,
+    );
+    const fiber = renderer.root._currentFiber().child;
+    if (__DEV__) {
+      // First render was locked
+      expect(renderer.toJSON().children).toEqual(['Loading']);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Loading']);
+
+      // Release the lock
+      setSuspenseHandler(() => false);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Done']);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Done']);
+
+      // Lock again
+      setSuspenseHandler(() => true);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Loading']);
+
+      // Release the lock again
+      setSuspenseHandler(() => false);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Done']);
+
+      // Ensure it checks specific fibers.
+      setSuspenseHandler(f => f === fiber || f === fiber.alternate);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Loading']);
+      setSuspenseHandler(f => f !== fiber && f !== fiber.alternate);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Done']);
+    } else {
+      expect(renderer.toJSON().children).toEqual(['Done']);
+    }
+  });
+
+  it('should support overriding suspense in concurrent mode', () => {
+    if (__DEV__) {
+      // Lock the first render
+      setSuspenseHandler(() => true);
+    }
+
+    function MyComponent() {
+      return 'Done';
+    }
+
+    const renderer = ReactTestRenderer.create(
+      <div>
+        <React.Suspense fallback={'Loading'}>
+          <MyComponent />
+        </React.Suspense>
+      </div>,
+      {unstable_isConcurrent: true},
+    );
+    expect(Scheduler).toFlushAndYield([]);
+    const fiber = renderer.root._currentFiber().child;
+    if (__DEV__) {
+      // First render was locked
+      expect(renderer.toJSON().children).toEqual(['Loading']);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Loading']);
+
+      // Release the lock
+      setSuspenseHandler(() => false);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Done']);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Done']);
+
+      // Lock again
+      setSuspenseHandler(() => true);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Loading']);
+
+      // Release the lock again
+      setSuspenseHandler(() => false);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Done']);
+
+      // Ensure it checks specific fibers.
+      setSuspenseHandler(f => f === fiber || f === fiber.alternate);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Loading']);
+      setSuspenseHandler(f => f !== fiber && f !== fiber.alternate);
+      scheduleUpdate(fiber); // Re-render
+      expect(renderer.toJSON().children).toEqual(['Done']);
+    } else {
+      expect(renderer.toJSON().children).toEqual(['Done']);
+    }
+  });
 });

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -96,6 +96,7 @@ import {
   registerSuspenseInstanceRetry,
 } from './ReactFiberHostConfig';
 import type {SuspenseInstance} from './ReactFiberHostConfig';
+import {shouldSuspend} from './ReactFiberReconciler';
 import {
   pushHostContext,
   pushHostContainer,
@@ -1392,6 +1393,12 @@ function updateSuspenseComponent(
   const mode = workInProgress.mode;
   const nextProps = workInProgress.pendingProps;
 
+  if (__DEV__) {
+    if (shouldSuspend(workInProgress)) {
+      workInProgress.effectTag |= DidCapture;
+    }
+  }
+
   // We should attempt to render the primary children unless this boundary
   // already suspended during this render (`alreadyCaptured` is true).
   let nextState: SuspenseState | null = workInProgress.memoizedState;

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -340,8 +340,16 @@ export function findHostInstanceWithNoPortals(
   return hostFiber.stateNode;
 }
 
+let shouldSuspendImpl = fiber => false;
+
+export function shouldSuspend(fiber: Fiber): boolean {
+  return shouldSuspendImpl(fiber);
+}
+
 let overrideHookState = null;
 let overrideProps = null;
+let scheduleUpdate = null;
+let setSuspenseHandler = null;
 
 if (__DEV__) {
   const copyWithSetImpl = (
@@ -409,6 +417,15 @@ if (__DEV__) {
     }
     scheduleWork(fiber, Sync);
   };
+
+  scheduleUpdate = (fiber: Fiber) => {
+    flushPassiveEffects();
+    scheduleWork(fiber, Sync);
+  };
+
+  setSuspenseHandler = (newShouldSuspendImpl: Fiber => boolean) => {
+    shouldSuspendImpl = newShouldSuspendImpl;
+  };
 }
 
 export function injectIntoDevTools(devToolsConfig: DevToolsConfig): boolean {
@@ -419,6 +436,8 @@ export function injectIntoDevTools(devToolsConfig: DevToolsConfig): boolean {
     ...devToolsConfig,
     overrideHookState,
     overrideProps,
+    setSuspenseHandler,
+    scheduleUpdate,
     currentDispatcherRef: ReactCurrentDispatcher,
     findHostInstanceByFiber(fiber: Fiber): Instance | TextInstance | null {
       const hostFiber = findCurrentHostFiber(fiber);