Make root.unmount() synchronous

When you unmount a root, the internal state that React stores on the
DOM node is immediately cleared. So, we should also synchronously
delete the React tree. You should be able to create a new root using
the same container.

Author: acdlite

packages/react-dom/src/__tests__/ReactDOMRoot-test.js

@@ -14,6 +14,7 @@ let ReactDOM = require('react-dom');
 let ReactDOMServer = require('react-dom/server');
 let Scheduler = require('scheduler');
 let act;
+let useEffect;
 
 describe('ReactDOMRoot', () => {
   let container;
@@ -26,6 +27,7 @@ describe('ReactDOMRoot', () => {
     ReactDOMServer = require('react-dom/server');
     Scheduler = require('scheduler');
     act = require('jest-react').act;
+    useEffect = React.useEffect;
   });
 
   it('renders children', () => {
@@ -342,4 +344,62 @@ describe('ReactDOMRoot', () => {
     });
     expect(container.textContent).toEqual('b');
   });
+
+  it('unmount is synchronous', async () => {
+    const root = ReactDOM.createRoot(container);
+    await act(async () => {
+      root.render('Hi');
+    });
+    expect(container.textContent).toEqual('Hi');
+
+    await act(async () => {
+      root.unmount();
+      // Should have already unmounted
+      expect(container.textContent).toEqual('');
+    });
+  });
+
+  it('throws if an unmounted root is updated', async () => {
+    const root = ReactDOM.createRoot(container);
+    await act(async () => {
+      root.render('Hi');
+    });
+    expect(container.textContent).toEqual('Hi');
+
+    root.unmount();
+
+    expect(() => root.render("I'm back")).toThrow(
+      'Cannot update an unmounted root.',
+    );
+  });
+
+  it('warns if root is unmounted inside an effect', async () => {
+    const container1 = document.createElement('div');
+    const root1 = ReactDOM.createRoot(container1);
+    const container2 = document.createElement('div');
+    const root2 = ReactDOM.createRoot(container2);
+
+    function App({step}) {
+      useEffect(() => {
+        if (step === 2) {
+          root2.unmount();
+        }
+      }, [step]);
+      return 'Hi';
+    }
+
+    await act(async () => {
+      root1.render(<App step={1} />);
+    });
+    expect(container1.textContent).toEqual('Hi');
+
+    expect(() => {
+      ReactDOM.flushSync(() => {
+        root1.render(<App step={2} />);
+      });
+    }).toErrorDev(
+      'Attempted to synchronously unmount a root while React was ' +
+        'already rendering.',
+    );
+  });
 });

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

@@ -14,7 +14,7 @@ import type {FiberRoot} from 'react-reconciler/src/ReactInternalTypes';
 export type RootType = {
   render(children: ReactNodeList): void,
   unmount(): void,
-  _internalRoot: FiberRoot,
+  _internalRoot: FiberRoot | null,
   ...
 };
 
@@ -62,17 +62,23 @@ import {
   updateContainer,
   findHostInstanceWithNoPortals,
   registerMutableSourceForHydration,
+  flushSync,
+  isAlreadyRendering,
 } from 'react-reconciler/src/ReactFiberReconciler';
 import invariant from 'shared/invariant';
 import {ConcurrentRoot} from 'react-reconciler/src/ReactRootTags';
 import {allowConcurrentByDefault} from 'shared/ReactFeatureFlags';
 
-function ReactDOMRoot(internalRoot) {
+function ReactDOMRoot(internalRoot: FiberRoot) {
   this._internalRoot = internalRoot;
 }
 
 ReactDOMRoot.prototype.render = function(children: ReactNodeList): void {
   const root = this._internalRoot;
+  if (root === null) {
+    invariant(false, 'Cannot update an unmounted root.');
+  }
+
   if (__DEV__) {
     if (typeof arguments[1] === 'function') {
       console.error(
@@ -109,10 +115,23 @@ ReactDOMRoot.prototype.unmount = function(): void {
     }
   }
   const root = this._internalRoot;
-  const container = root.containerInfo;
-  updateContainer(null, root, null, () => {
+  if (root !== null) {
+    this._internalRoot = null;
+    const container = root.containerInfo;
+    if (__DEV__) {
+      if (isAlreadyRendering()) {
+        console.error(
+          'Attempted to synchronously unmount a root while React was already ' +
+            'rendering. React cannot finish unmounting the root until the ' +
+            'current render has completed, which may lead to a race condition.',
+        );
+      }
+    }
+    flushSync(() => {
+      updateContainer(null, root, null, null);
+    });
     unmarkContainerAsRoot(container);
-  });
+  }
 };
 
 export function createRoot(

scripts/error-codes/codes.json

@@ -396,5 +396,6 @@
   "405": "hydrateRoot(...): Target container is not a DOM element.",
   "406": "act(...) is not supported in production builds of React.",
   "407": "Missing getServerSnapshot, which is required for server-rendered content. Will revert to client rendering.",
-  "408": "Missing getServerSnapshot, which is required for server-rendered content."
+  "408": "Missing getServerSnapshot, which is required for server-rendered content.",
+  "409": "Cannot update an unmounted root."
 }