Add "hydrationOptions" behind the enableSuspenseCallback flag (#16434)

This gets invoked when a boundary is either hydrated or if it is deleted
because it updated or got deleted before it mounted.

packages/react-art/src/ReactART.js

@@ -66,7 +66,7 @@ class Surface extends React.Component {
 
     this._surface = Mode.Surface(+width, +height, this._tagRef);
 
-    this._mountNode = createContainer(this._surface, LegacyRoot, false);
+    this._mountNode = createContainer(this._surface, LegacyRoot, false, null);
     updateContainer(this.props.children, this._mountNode, this);
   }
 

packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js

@@ -24,6 +24,7 @@ describe('ReactDOMServerPartialHydration', () => {
 
     ReactFeatureFlags = require('shared/ReactFeatureFlags');
     ReactFeatureFlags.enableSuspenseServerRenderer = true;
+    ReactFeatureFlags.enableSuspenseCallback = true;
 
     React = require('react');
     ReactDOM = require('react-dom');
@@ -92,6 +93,153 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(ref.current).toBe(span);
   });
 
+  it('calls the hydration callbacks after hydration or deletion', async () => {
+    let suspend = false;
+    let resolve;
+    let promise = new Promise(resolvePromise => (resolve = resolvePromise));
+    function Child() {
+      if (suspend) {
+        throw promise;
+      } else {
+        return 'Hello';
+      }
+    }
+
+    let suspend2 = false;
+    let promise2 = new Promise(() => {});
+    function Child2() {
+      if (suspend2) {
+        throw promise2;
+      } else {
+        return 'World';
+      }
+    }
+
+    function App({value}) {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <Child />
+          </Suspense>
+          <Suspense fallback="Loading...">
+            <Child2 value={value} />
+          </Suspense>
+        </div>
+      );
+    }
+
+    // First we render the final HTML. With the streaming renderer
+    // this may have suspense points on the server but here we want
+    // to test the completed HTML. Don't suspend on the server.
+    suspend = false;
+    suspend2 = false;
+    let finalHTML = ReactDOMServer.renderToString(<App />);
+
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    let hydrated = [];
+    let deleted = [];
+
+    // On the client we don't have all data yet but we want to start
+    // hydrating anyway.
+    suspend = true;
+    suspend2 = true;
+    let root = ReactDOM.unstable_createRoot(container, {
+      hydrate: true,
+      hydrationOptions: {
+        onHydrated(node) {
+          hydrated.push(node);
+        },
+        onDeleted(node) {
+          deleted.push(node);
+        },
+      },
+    });
+    act(() => {
+      root.render(<App />);
+    });
+
+    expect(hydrated.length).toBe(0);
+    expect(deleted.length).toBe(0);
+
+    await act(async () => {
+      // Resolving the promise should continue hydration
+      suspend = false;
+      resolve();
+      await promise;
+    });
+
+    expect(hydrated.length).toBe(1);
+    expect(deleted.length).toBe(0);
+
+    // Performing an update should force it to delete the boundary
+    root.render(<App value={true} />);
+
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+
+    expect(hydrated.length).toBe(1);
+    expect(deleted.length).toBe(1);
+  });
+
+  it('calls the onDeleted hydration callback if the parent gets deleted', async () => {
+    let suspend = false;
+    let promise = new Promise(() => {});
+    function Child() {
+      if (suspend) {
+        throw promise;
+      } else {
+        return 'Hello';
+      }
+    }
+
+    function App({deleted}) {
+      if (deleted) {
+        return null;
+      }
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <Child />
+          </Suspense>
+        </div>
+      );
+    }
+
+    suspend = false;
+    let finalHTML = ReactDOMServer.renderToString(<App />);
+
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    let deleted = [];
+
+    // On the client we don't have all data yet but we want to start
+    // hydrating anyway.
+    suspend = true;
+    let root = ReactDOM.unstable_createRoot(container, {
+      hydrate: true,
+      hydrationOptions: {
+        onDeleted(node) {
+          deleted.push(node);
+        },
+      },
+    });
+    act(() => {
+      root.render(<App />);
+    });
+
+    expect(deleted.length).toBe(0);
+
+    act(() => {
+      root.render(<App deleted={true} />);
+    });
+
+    // The callback should have been invoked.
+    expect(deleted.length).toBe(1);
+  });
+
   it('warns and replaces the boundary content in legacy mode', async () => {
     let suspend = false;
     let resolve;

packages/react-dom/src/client/ReactDOM.js

@@ -367,15 +367,26 @@ ReactWork.prototype._onCommit = function(): void {
 function ReactSyncRoot(
   container: DOMContainer,
   tag: RootTag,
-  hydrate: boolean,
+  options: void | RootOptions,
 ) {
   // Tag is either LegacyRoot or Concurrent Root
-  const root = createContainer(container, tag, hydrate);
+  const hydrate = options != null && options.hydrate === true;
+  const hydrationCallbacks =
+    (options != null && options.hydrationOptions) || null;
+  const root = createContainer(container, tag, hydrate, hydrationCallbacks);
   this._internalRoot = root;
 }
 
-function ReactRoot(container: DOMContainer, hydrate: boolean) {
-  const root = createContainer(container, ConcurrentRoot, hydrate);
+function ReactRoot(container: DOMContainer, options: void | RootOptions) {
+  const hydrate = options != null && options.hydrate === true;
+  const hydrationCallbacks =
+    (options != null && options.hydrationOptions) || null;
+  const root = createContainer(
+    container,
+    ConcurrentRoot,
+    hydrate,
+    hydrationCallbacks,
+  );
   this._internalRoot = root;
 }
 
@@ -532,7 +543,15 @@ function legacyCreateRootFromDOMContainer(
   }
 
   // Legacy roots are not batched.
-  return new ReactSyncRoot(container, LegacyRoot, shouldHydrate);
+  return new ReactSyncRoot(
+    container,
+    LegacyRoot,
+    shouldHydrate
+      ? {
+          hydrate: true,
+        }
+      : undefined,
+  );
 }
 
 function legacyRenderSubtreeIntoContainer(
@@ -824,6 +843,10 @@ const ReactDOM: Object = {
 
 type RootOptions = {
   hydrate?: boolean,
+  hydrationOptions?: {
+    onHydrated?: (suspenseNode: Comment) => void,
+    onDeleted?: (suspenseNode: Comment) => void,
+  },
 };
 
 function createRoot(
@@ -839,8 +862,7 @@ function createRoot(
     functionName,
   );
   warnIfReactDOMContainerInDEV(container);
-  const hydrate = options != null && options.hydrate === true;
-  return new ReactRoot(container, hydrate);
+  return new ReactRoot(container, options);
 }
 
 function createSyncRoot(
@@ -856,8 +878,7 @@ function createSyncRoot(
     functionName,
   );
   warnIfReactDOMContainerInDEV(container);
-  const hydrate = options != null && options.hydrate === true;
-  return new ReactSyncRoot(container, BatchedRoot, hydrate);
+  return new ReactSyncRoot(container, BatchedRoot, options);
 }
 
 function warnIfReactDOMContainerInDEV(container) {

packages/react-native-renderer/src/ReactFabric.js

@@ -144,7 +144,7 @@ const ReactFabric: ReactFabricType = {
     if (!root) {
       // TODO (bvaughn): If we decide to keep the wrapper component,
       // We could create a wrapper for containerTag as well to reduce special casing.
-      root = createContainer(containerTag, LegacyRoot, false);
+      root = createContainer(containerTag, LegacyRoot, false, null);
       roots.set(containerTag, root);
     }
     updateContainer(element, root, null, callback);

packages/react-native-renderer/src/ReactNativeRenderer.js

@@ -141,7 +141,7 @@ const ReactNativeRenderer: ReactNativeType = {
     if (!root) {
       // TODO (bvaughn): If we decide to keep the wrapper component,
       // We could create a wrapper for containerTag as well to reduce special casing.
-      root = createContainer(containerTag, LegacyRoot, false);
+      root = createContainer(containerTag, LegacyRoot, false, null);
       roots.set(containerTag, root);
     }
     updateContainer(element, root, null, callback);

packages/react-noop-renderer/src/createReactNoop.js

@@ -908,7 +908,7 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
       if (!root) {
         const container = {rootID: rootID, pendingChildren: [], children: []};
         rootContainers.set(rootID, container);
-        root = NoopRenderer.createContainer(container, tag, false);
+        root = NoopRenderer.createContainer(container, tag, false, null);
         roots.set(rootID, root);
       }
       return root.current.stateNode.containerInfo;
@@ -925,6 +925,7 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
         container,
         ConcurrentRoot,
         false,
+        null,
       );
       return {
         _Scheduler: Scheduler,
@@ -950,6 +951,7 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
         container,
         BatchedRoot,
         false,
+        null,
       );
       return {
         _Scheduler: Scheduler,

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -607,7 +607,12 @@ function commitLifeCycles(
       }
       return;
     }
-    case SuspenseComponent:
+    case SuspenseComponent: {
+      if (enableSuspenseCallback) {
+        commitSuspenseHydrationCallbacks(finishedRoot, finishedWork);
+      }
+      return;
+    }
     case SuspenseListComponent:
     case IncompleteClassComponent:
     case FundamentalComponent:
@@ -644,7 +649,8 @@ function hideOrUnhideAllChildren(finishedWork, isHidden) {
         }
       } else if (
         node.tag === SuspenseComponent &&
-        node.memoizedState !== null
+        node.memoizedState !== null &&
+        node.memoizedState.dehydrated === null
       ) {
         // Found a nested Suspense component that timed out. Skip over the
         // primary child fragment, which should remain hidden.
@@ -719,6 +725,7 @@ function commitDetachRef(current: Fiber) {
 // deletion, so don't let them throw. Host-originating errors should
 // interrupt deletion, so it's okay
 function commitUnmount(
+  finishedRoot: FiberRoot,
   current: Fiber,
   renderPriorityLevel: ReactPriorityLevel,
 ): void {
@@ -801,7 +808,7 @@ function commitUnmount(
       // We are also not using this parent because
       // the portal will get pushed immediately.
       if (supportsMutation) {
-        unmountHostComponents(current, renderPriorityLevel);
+        unmountHostComponents(finishedRoot, current, renderPriorityLevel);
       } else if (supportsPersistence) {
         emptyPortalContainer(current);
       }
@@ -815,11 +822,24 @@ function commitUnmount(
           current.stateNode = null;
         }
       }
+      return;
+    }
+    case DehydratedFragment: {
+      if (enableSuspenseCallback) {
+        const hydrationCallbacks = finishedRoot.hydrationCallbacks;
+        if (hydrationCallbacks !== null) {
+          const onDeleted = hydrationCallbacks.onDeleted;
+          if (onDeleted) {
+            onDeleted((current.stateNode: SuspenseInstance));
+          }
+        }
+      }
     }
   }
 }
 
 function commitNestedUnmounts(
+  finishedRoot: FiberRoot,
   root: Fiber,
   renderPriorityLevel: ReactPriorityLevel,
 ): void {
@@ -830,7 +850,7 @@ function commitNestedUnmounts(
   // we do an inner loop while we're still inside the host node.
   let node: Fiber = root;
   while (true) {
-    commitUnmount(node, renderPriorityLevel);
+    commitUnmount(finishedRoot, node, renderPriorityLevel);
     // Visit children because they may contain more composite or host nodes.
     // Skip portals because commitUnmount() currently visits them recursively.
     if (
@@ -1081,7 +1101,11 @@ function commitPlacement(finishedWork: Fiber): void {
   }
 }
 
-function unmountHostComponents(current, renderPriorityLevel): void {
+function unmountHostComponents(
+  finishedRoot,
+  current,
+  renderPriorityLevel,
+): void {
   // We only have the top Fiber that was deleted but we need to recurse down its
   // children to find all the terminal nodes.
   let node: Fiber = current;
@@ -1129,7 +1153,7 @@ function unmountHostComponents(current, renderPriorityLevel): void {
     }
 
     if (node.tag === HostComponent || node.tag === HostText) {
-      commitNestedUnmounts(node, renderPriorityLevel);
+      commitNestedUnmounts(finishedRoot, node, renderPriorityLevel);
       // After all the children have unmounted, it is now safe to remove the
       // node from the tree.
       if (currentParentIsContainer) {
@@ -1146,7 +1170,7 @@ function unmountHostComponents(current, renderPriorityLevel): void {
       // Don't visit children because we already visited them.
     } else if (enableFundamentalAPI && node.tag === FundamentalComponent) {
       const fundamentalNode = node.stateNode.instance;
-      commitNestedUnmounts(node, renderPriorityLevel);
+      commitNestedUnmounts(finishedRoot, node, renderPriorityLevel);
       // After all the children have unmounted, it is now safe to remove the
       // node from the tree.
       if (currentParentIsContainer) {
@@ -1164,6 +1188,16 @@ function unmountHostComponents(current, renderPriorityLevel): void {
       enableSuspenseServerRenderer &&
       node.tag === DehydratedFragment
     ) {
+      if (enableSuspenseCallback) {
+        const hydrationCallbacks = finishedRoot.hydrationCallbacks;
+        if (hydrationCallbacks !== null) {
+          const onDeleted = hydrationCallbacks.onDeleted;
+          if (onDeleted) {
+            onDeleted((node.stateNode: SuspenseInstance));
+          }
+        }
+      }
+
       // Delete the dehydrated suspense boundary and all of its content.
       if (currentParentIsContainer) {
         clearSuspenseBoundaryFromContainer(
@@ -1188,7 +1222,7 @@ function unmountHostComponents(current, renderPriorityLevel): void {
         continue;
       }
     } else {
-      commitUnmount(node, renderPriorityLevel);
+      commitUnmount(finishedRoot, node, renderPriorityLevel);
       // Visit children because we may find more host components below.
       if (node.child !== null) {
         node.child.return = node;
@@ -1216,16 +1250,17 @@ function unmountHostComponents(current, renderPriorityLevel): void {
 }
 
 function commitDeletion(
+  finishedRoot: FiberRoot,
   current: Fiber,
   renderPriorityLevel: ReactPriorityLevel,
 ): void {
   if (supportsMutation) {
     // Recursively delete all host nodes from the parent.
     // Detach refs and call componentWillUnmount() on the whole subtree.
-    unmountHostComponents(current, renderPriorityLevel);
+    unmountHostComponents(finishedRoot, current, renderPriorityLevel);
   } else {
     // Detach refs and call componentWillUnmount() on the whole subtree.
-    commitNestedUnmounts(current, renderPriorityLevel);
+    commitNestedUnmounts(finishedRoot, current, renderPriorityLevel);
   }
   detachFiber(current);
 }
@@ -1382,6 +1417,30 @@ function commitSuspenseComponent(finishedWork: Fiber) {
   }
 }
 
+function commitSuspenseHydrationCallbacks(
+  finishedRoot: FiberRoot,
+  finishedWork: Fiber,
+) {
+  if (enableSuspenseCallback) {
+    const hydrationCallbacks = finishedRoot.hydrationCallbacks;
+    if (hydrationCallbacks !== null) {
+      const onHydrated = hydrationCallbacks.onHydrated;
+      if (onHydrated) {
+        const newState: SuspenseState | null = finishedWork.memoizedState;
+        if (newState === null) {
+          const current = finishedWork.alternate;
+          if (current !== null) {
+            const prevState: SuspenseState | null = current.memoizedState;
+            if (prevState !== null && prevState.dehydrated !== null) {
+              onHydrated(prevState.dehydrated);
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
 function attachSuspenseRetryListeners(finishedWork: Fiber) {
   // If this boundary just timed out, then it will have a set of thenables.
   // For each thenable, attach a listener so that when it resolves, React

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -859,6 +859,10 @@ function completeWork(
             if ((workInProgress.effectTag & DidCapture) === NoEffect) {
               // This boundary did not suspend so it's now hydrated and unsuspended.
               workInProgress.memoizedState = null;
+              if (enableSuspenseCallback) {
+                // Notify the callback.
+                workInProgress.effectTag |= Update;
+              }
             } else {
               // Something suspended. Schedule an effect to attach retry listeners.
               workInProgress.effectTag |= Update;

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -20,6 +20,7 @@ import {FundamentalComponent} from 'shared/ReactWorkTags';
 import type {ReactNodeList} from 'shared/ReactTypes';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 import type {SuspenseConfig} from './ReactFiberSuspenseConfig';
+import type {SuspenseHydrationCallbacks} from './ReactFiberSuspenseComponent';
 
 import {
   findCurrentHostFiber,
@@ -294,8 +295,9 @@ export function createContainer(
   containerInfo: Container,
   tag: RootTag,
   hydrate: boolean,
+  hydrationCallbacks: null | SuspenseHydrationCallbacks,
 ): OpaqueRoot {
-  return createFiberRoot(containerInfo, tag, hydrate);
+  return createFiberRoot(containerInfo, tag, hydrate, hydrationCallbacks);
 }
 
 export function updateContainer(

packages/react-reconciler/src/ReactFiberRoot.js

@@ -13,11 +13,15 @@ import type {RootTag} from 'shared/ReactRootTags';
 import type {TimeoutHandle, NoTimeout} from './ReactFiberHostConfig';
 import type {Thenable} from './ReactFiberWorkLoop';
 import type {Interaction} from 'scheduler/src/Tracing';
+import type {SuspenseHydrationCallbacks} from './ReactFiberSuspenseComponent';
 
 import {noTimeout} from './ReactFiberHostConfig';
 import {createHostRootFiber} from './ReactFiber';
 import {NoWork} from './ReactFiberExpirationTime';
-import {enableSchedulerTracing} from 'shared/ReactFeatureFlags';
+import {
+  enableSchedulerTracing,
+  enableSuspenseCallback,
+} from 'shared/ReactFeatureFlags';
 import {unstable_getThreadID} from 'scheduler/tracing';
 
 // TODO: This should be lifted into the renderer.
@@ -83,6 +87,11 @@ type ProfilingOnlyFiberRootProperties = {|
   pendingInteractionMap: PendingInteractionMap,
 |};
 
+// The follow fields are only used by enableSuspenseCallback for hydration.
+type SuspenseCallbackOnlyFiberRootProperties = {|
+  hydrationCallbacks: null | SuspenseHydrationCallbacks,
+|};
+
 // Exported FiberRoot type includes all properties,
 // To avoid requiring potentially error-prone :any casts throughout the project.
 // Profiling properties are only safe to access in profiling builds (when enableSchedulerTracing is true).
@@ -91,6 +100,7 @@ type ProfilingOnlyFiberRootProperties = {|
 export type FiberRoot = {
   ...BaseFiberRootProperties,
   ...ProfilingOnlyFiberRootProperties,
+  ...SuspenseCallbackOnlyFiberRootProperties,
 };
 
 function FiberRootNode(containerInfo, tag, hydrate) {
@@ -117,14 +127,21 @@ function FiberRootNode(containerInfo, tag, hydrate) {
     this.memoizedInteractions = new Set();
     this.pendingInteractionMap = new Map();
   }
+  if (enableSuspenseCallback) {
+    this.hydrationCallbacks = null;
+  }
 }
 
 export function createFiberRoot(
   containerInfo: any,
   tag: RootTag,
   hydrate: boolean,
+  hydrationCallbacks: null | SuspenseHydrationCallbacks,
 ): FiberRoot {
   const root: FiberRoot = (new FiberRootNode(containerInfo, tag, hydrate): any);
+  if (enableSuspenseCallback) {
+    root.hydrationCallbacks = hydrationCallbacks;
+  }
 
   // Cyclic construction. This cheats the type system right now because
   // stateNode is any.

packages/react-reconciler/src/ReactFiberSuspenseComponent.js

@@ -17,6 +17,11 @@ import {
   isSuspenseInstanceFallback,
 } from './ReactFiberHostConfig';
 
+export type SuspenseHydrationCallbacks = {
+  onHydrated?: (suspenseInstance: SuspenseInstance) => void,
+  onDeleted?: (suspenseInstance: SuspenseInstance) => void,
+};
+
 // A null SuspenseState represents an unsuspended normal Suspense boundary.
 // A non-null SuspenseState means that it is blocked for one reason or another.
 // - A non-null dehydrated field means it's blocked pending hydration.

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -1638,6 +1638,7 @@ function commitRootImpl(root, renderPriorityLevel) {
           null,
           commitMutationEffects,
           null,
+          root,
           renderPriorityLevel,
         );
         if (hasCaughtError()) {
@@ -1648,7 +1649,7 @@ function commitRootImpl(root, renderPriorityLevel) {
         }
       } else {
         try {
-          commitMutationEffects(renderPriorityLevel);
+          commitMutationEffects(root, renderPriorityLevel);
         } catch (error) {
           invariant(nextEffect !== null, 'Should be working on an effect.');
           captureCommitPhaseError(nextEffect, error);
@@ -1837,7 +1838,7 @@ function commitBeforeMutationEffects() {
   }
 }
 
-function commitMutationEffects(renderPriorityLevel) {
+function commitMutationEffects(root: FiberRoot, renderPriorityLevel) {
   // TODO: Should probably move the bulk of this function to commitWork.
   while (nextEffect !== null) {
     setCurrentDebugFiberInDEV(nextEffect);
@@ -1888,7 +1889,7 @@ function commitMutationEffects(renderPriorityLevel) {
         break;
       }
       case Deletion: {
-        commitDeletion(nextEffect, renderPriorityLevel);
+        commitDeletion(root, nextEffect, renderPriorityLevel);
         break;
       }
     }

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

@@ -58,6 +58,7 @@ describe('ReactFiberHostContext', () => {
       /* root: */ null,
       ConcurrentRoot,
       false,
+      null,
     );
     Renderer.updateContainer(
       <a>
@@ -110,6 +111,7 @@ describe('ReactFiberHostContext', () => {
       rootContext,
       ConcurrentRoot,
       false,
+      null,
     );
     Renderer.updateContainer(
       <a>

packages/react-test-renderer/src/ReactTestRenderer.js

@@ -442,6 +442,7 @@ const ReactTestRendererFiber = {
       container,
       isConcurrent ? ConcurrentRoot : LegacyRoot,
       false,
+      null,
     );
     invariant(root != null, 'something went wrong');
     updateContainer(element, root, null, null);

packages/shared/ReactFeatureFlags.js

@@ -84,6 +84,8 @@ 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.
+// Also allows hydration callbacks to fire when a dehydrated boundary gets
+// hydrated or deleted.
 export const enableSuspenseCallback = false;
 
 // Part of the simplification of React.createElement so we can eventually move