Add suspenseCallback feature for runtime tracing of loading states

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -30,6 +30,7 @@ import {
   enableSuspenseServerRenderer,
   enableFlareAPI,
   enableFundamentalAPI,
+  enableSuspenseCallback,
 } from 'shared/ReactFeatureFlags';
 import {
   FunctionComponent,
@@ -1321,6 +1322,18 @@ function commitSuspenseComponent(finishedWork: Fiber) {
   if (supportsMutation && primaryChildParent !== null) {
     hideOrUnhideAllChildren(primaryChildParent, newDidTimeout);
   }
+
+  if (enableSuspenseCallback && newState !== null) {
+    const suspenseCallback = finishedWork.memoizedProps.suspenseCallback;
+    if (typeof suspenseCallback === 'function') {
+      const thenables: Set<Thenable> | null = (finishedWork.updateQueue: any);
+      if (thenables !== null) {
+        suspenseCallback(new Set(thenables));
+      }
+    } else if (__DEV__) {
+      warning(false, 'Unexpected type for suspenseCallback.');
+    }
+  }
 }
 
 function attachSuspenseRetryListeners(finishedWork: Fiber) {

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -115,6 +115,7 @@ import {
 } from './ReactFiberHydrationContext';
 import {
   enableSchedulerTracing,
+  enableSuspenseCallback,
   enableSuspenseServerRenderer,
   enableFlareAPI,
   enableFundamentalAPI,
@@ -884,6 +885,14 @@ function completeWork(
           workInProgress.effectTag |= Update;
         }
       }
+      if (
+        enableSuspenseCallback &&
+        workInProgress.updateQueue !== null &&
+        workInProgress.memoizedProps.suspenseCallback != null
+      ) {
+        // Always notify the callback
+        workInProgress.effectTag |= Update;
+      }
       break;
     }
     case Fragment:

packages/react-reconciler/src/__tests__/ReactSuspenseCallback-test.internal.js

@@ -0,0 +1,217 @@
+/**
+ * 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.
+ *
+ * @emails react-core
+ * @jest-environment node
+ */
+'use strict';
+
+let React;
+let ReactFeatureFlags;
+let ReactNoop;
+let Scheduler;
+
+describe('ReactSuspense', () => {
+  beforeEach(() => {
+    jest.resetModules();
+    ReactFeatureFlags = require('shared/ReactFeatureFlags');
+    ReactFeatureFlags.enableSuspenseCallback = true;
+
+    React = require('react');
+    ReactNoop = require('react-noop-renderer');
+    Scheduler = require('scheduler');
+  });
+
+  function text(t) {
+    return {text: t, hidden: false};
+  }
+
+  function createThenable() {
+    let completed = false;
+    const resolveRef = {current: null};
+    let promise = {
+      then(resolve, reject) {
+        resolveRef.current = () => {
+          completed = true;
+          resolve();
+        };
+      },
+    };
+
+    const PromiseComp = () => {
+      if (!completed) {
+        throw promise;
+      }
+      return 'Done';
+    };
+    return {promise, resolveRef, PromiseComp};
+  }
+
+  it('1 then 0 suspense callback', () => {
+    const {promise, resolveRef, PromiseComp} = createThenable();
+
+    let ops = [];
+    const suspenseCallback = thenables => {
+      ops.push(thenables);
+    };
+
+    const element = (
+      <React.Suspense suspenseCallback={suspenseCallback} fallback={'Waiting'}>
+        <PromiseComp />
+      </React.Suspense>
+    );
+
+    ReactNoop.render(element);
+    expect(Scheduler).toFlushWithoutYielding();
+    expect(ReactNoop.getChildren()).toEqual([text('Waiting')]);
+    expect(ops).toEqual([new Set([promise])]);
+    ops = [];
+
+    resolveRef.current();
+    expect(Scheduler).toFlushWithoutYielding();
+    expect(ReactNoop.getChildren()).toEqual([text('Done')]);
+    expect(ops).toEqual([]);
+  });
+
+  it('2 then 1 then 0 suspense callback', () => {
+    const {
+      promise: promise1,
+      resolveRef: resolveRef1,
+      PromiseComp: PromiseComp1,
+    } = createThenable();
+    const {
+      promise: promise2,
+      resolveRef: resolveRef2,
+      PromiseComp: PromiseComp2,
+    } = createThenable();
+
+    let ops = [];
+    const suspenseCallback1 = thenables => {
+      ops.push(thenables);
+    };
+
+    const element = (
+      <React.Suspense
+        suspenseCallback={suspenseCallback1}
+        fallback={'Waiting Tier 1'}>
+        <PromiseComp1 />
+        <PromiseComp2 />
+      </React.Suspense>
+    );
+
+    ReactNoop.render(element);
+    expect(Scheduler).toFlushWithoutYielding();
+    expect(ReactNoop.getChildren()).toEqual([text('Waiting Tier 1')]);
+    expect(ops).toEqual([new Set([promise1, promise2])]);
+    ops = [];
+
+    resolveRef1.current();
+    ReactNoop.render(element);
+    expect(Scheduler).toFlushWithoutYielding();
+    expect(ReactNoop.getChildren()).toEqual([text('Waiting Tier 1')]);
+    expect(ops).toEqual([new Set([promise2])]);
+    ops = [];
+
+    resolveRef2.current();
+    ReactNoop.render(element);
+    expect(Scheduler).toFlushWithoutYielding();
+    expect(ReactNoop.getChildren()).toEqual([text('Done'), text('Done')]);
+    expect(ops).toEqual([]);
+  });
+
+  it('nested suspense promises are reported only for their tier', () => {
+    const {promise, PromiseComp} = createThenable();
+
+    let ops1 = [];
+    const suspenseCallback1 = thenables => {
+      ops1.push(thenables);
+    };
+    let ops2 = [];
+    const suspenseCallback2 = thenables => {
+      ops2.push(thenables);
+    };
+
+    const element = (
+      <React.Suspense
+        suspenseCallback={suspenseCallback1}
+        fallback={'Waiting Tier 1'}>
+        <React.Suspense
+          suspenseCallback={suspenseCallback2}
+          fallback={'Waiting Tier 2'}>
+          <PromiseComp />
+        </React.Suspense>
+      </React.Suspense>
+    );
+
+    ReactNoop.render(element);
+    expect(Scheduler).toFlushWithoutYielding();
+    expect(ReactNoop.getChildren()).toEqual([text('Waiting Tier 2')]);
+    expect(ops1).toEqual([]);
+    expect(ops2).toEqual([new Set([promise])]);
+  });
+
+  it('competing suspense promises', () => {
+    const {
+      promise: promise1,
+      resolveRef: resolveRef1,
+      PromiseComp: PromiseComp1,
+    } = createThenable();
+    const {
+      promise: promise2,
+      resolveRef: resolveRef2,
+      PromiseComp: PromiseComp2,
+    } = createThenable();
+
+    let ops1 = [];
+    const suspenseCallback1 = thenables => {
+      ops1.push(thenables);
+    };
+    let ops2 = [];
+    const suspenseCallback2 = thenables => {
+      ops2.push(thenables);
+    };
+
+    const element = (
+      <React.Suspense
+        suspenseCallback={suspenseCallback1}
+        fallback={'Waiting Tier 1'}>
+        <React.Suspense
+          suspenseCallback={suspenseCallback2}
+          fallback={'Waiting Tier 2'}>
+          <PromiseComp2 />
+        </React.Suspense>
+        <PromiseComp1 />
+      </React.Suspense>
+    );
+
+    ReactNoop.render(element);
+    expect(Scheduler).toFlushWithoutYielding();
+    expect(ReactNoop.getChildren()).toEqual([text('Waiting Tier 1')]);
+    expect(ops1).toEqual([new Set([promise1])]);
+    expect(ops2).toEqual([]);
+    ops1 = [];
+    ops2 = [];
+
+    resolveRef1.current();
+    ReactNoop.render(element);
+    expect(Scheduler).toFlushWithoutYielding();
+    expect(ReactNoop.getChildren()).toEqual([
+      text('Waiting Tier 2'),
+      text('Done'),
+    ]);
+    expect(ops1).toEqual([]);
+    expect(ops2).toEqual([new Set([promise2])]);
+    ops1 = [];
+    ops2 = [];
+
+    resolveRef2.current();
+    ReactNoop.render(element);
+    expect(Scheduler).toFlushWithoutYielding();
+    expect(ReactNoop.getChildren()).toEqual([text('Done'), text('Done')]);
+    expect(ops1).toEqual([]);
+    expect(ops2).toEqual([]);
+  });
+});

packages/shared/ReactFeatureFlags.js

@@ -78,3 +78,8 @@ export const revertPassiveEffectsChange = false;
 // but without making them discrete. The flag exists in case it causes
 // starvation problems.
 export const enableUserBlockingEvents = false;
+
+// Add a callback property to suspense to notify which promises are currently
+// in the update queue. This allows reporting and tracing of what is causing
+// the user to see a loading state.
+export const enableSuspenseCallback = false;

packages/shared/forks/ReactFeatureFlags.native-fb.js

@@ -37,6 +37,7 @@ export const enableJSXTransformAPI = false;
 export const warnAboutMissingMockScheduler = true;
 export const revertPassiveEffectsChange = false;
 export const enableUserBlockingEvents = false;
+export const enableSuspenseCallback = false;
 
 // Only used in www builds.
 export function addUserTimingListener() {

packages/shared/forks/ReactFeatureFlags.native-oss.js

@@ -32,6 +32,7 @@ export const enableJSXTransformAPI = false;
 export const warnAboutMissingMockScheduler = false;
 export const revertPassiveEffectsChange = false;
 export const enableUserBlockingEvents = false;
+export const enableSuspenseCallback = false;
 
 // Only used in www builds.
 export function addUserTimingListener() {

packages/shared/forks/ReactFeatureFlags.persistent.js

@@ -32,6 +32,7 @@ export const enableJSXTransformAPI = false;
 export const warnAboutMissingMockScheduler = true;
 export const revertPassiveEffectsChange = false;
 export const enableUserBlockingEvents = false;
+export const enableSuspenseCallback = false;
 
 // Only used in www builds.
 export function addUserTimingListener() {

packages/shared/forks/ReactFeatureFlags.test-renderer.js

@@ -32,6 +32,7 @@ export const enableJSXTransformAPI = false;
 export const warnAboutMissingMockScheduler = false;
 export const revertPassiveEffectsChange = false;
 export const enableUserBlockingEvents = false;
+export const enableSuspenseCallback = false;
 
 // Only used in www builds.
 export function addUserTimingListener() {

packages/shared/forks/ReactFeatureFlags.test-renderer.www.js

@@ -32,6 +32,7 @@ export const enableFundamentalAPI = false;
 export const enableJSXTransformAPI = true;
 export const warnAboutMissingMockScheduler = true;
 export const enableUserBlockingEvents = false;
+export const enableSuspenseCallback = true;
 
 // Only used in www builds.
 export function addUserTimingListener() {

packages/shared/forks/ReactFeatureFlags.www.js

@@ -76,6 +76,8 @@ export const enableJSXTransformAPI = true;
 
 export const warnAboutMissingMockScheduler = true;
 
+export const enableSuspenseCallback = true;
+
 // Flow magic to verify the exports of this file match the original version.
 // eslint-disable-next-line no-unused-vars
 type Check<_X, Y: _X, X: Y = _X> = null;