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.

Author: sebmarkbage

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,