Unmount roots correctly in React 18+

Author: alexeyr

node_package/src/clientStartup.ts

@@ -5,21 +5,25 @@ import type {
   ReactOnRails as ReactOnRailsType,
   RegisteredComponent,
   RenderFunction,
-} from './types/index';
+  Root,
+} from './types';
 
 import createReactOutput from './createReactOutput';
 import { isServerRenderHash } from './isServerRenderResult';
 import reactHydrateOrRender from './reactHydrateOrRender';
+import { supportsRootApi } from './reactApis';
 
 declare global {
   interface Window {
     ReactOnRails: ReactOnRailsType;
     __REACT_ON_RAILS_EVENT_HANDLERS_RAN_ONCE__?: boolean;
+    roots: Root[];
   }
 
   namespace NodeJS {
     interface Global {
       ReactOnRails: ReactOnRailsType;
+      roots: Root[];
     }
   }
   namespace Turbolinks {
@@ -169,7 +173,10 @@ function render(el: Element, railsContext: RailsContext): void {
 You returned a server side type of react-router error: ${JSON.stringify(reactElementOrRouterResult)}
 You should return a React.Component always for the client side entry point.`);
       } else {
-        reactHydrateOrRender(domNode, reactElementOrRouterResult as ReactElement, shouldHydrate);
+        const rootOrElement = reactHydrateOrRender(domNode, reactElementOrRouterResult as ReactElement, shouldHydrate);
+        if (supportsRootApi) {
+          context.roots.push(rootOrElement as Root);
+        }
       }
     }
   } catch (e: any) {
@@ -203,6 +210,9 @@ export function reactOnRailsPageLoaded(): void {
   if (!railsContext) return;
 
   forEachStore(railsContext);
+  if (supportsRootApi) {
+    findContext().roots = [];
+  }
   forEachReactOnRailsComponentInitialize(render, railsContext);
 }
 
@@ -222,9 +232,15 @@ function unmount(el: Element): void {
 
 function reactOnRailsPageUnloaded(): void {
   debugTurbolinks('reactOnRailsPageUnloaded');
-  const els = reactOnRailsHtmlElements();
-  for (let i = 0; i < els.length; i += 1) {
-    unmount(els[i]);
+  if (supportsRootApi) {
+    for (const root of findContext().roots) {
+      root.unmount();
+    }
+  } else {
+    const els = reactOnRailsHtmlElements();
+    for (let i = 0; i < els.length; i += 1) {
+      unmount(els[i]);
+    }
   }
 }
 

node_package/src/reactApis.ts

@@ -0,0 +1,8 @@
+import ReactDOM from 'react-dom';
+
+const reactMajorVersion = ReactDOM.version?.split('.')[0] || 16;
+
+// TODO: once we require React 18, we can remove this and inline everything guarded by it.
+// Not the default export because others may be added for future React versions.
+// eslint-disable-next-line import/prefer-default-export
+export const supportsRootApi = reactMajorVersion >= 18;

node_package/src/reactHydrateOrRender.ts

@@ -1,44 +1,43 @@
-import type { ReactElement } from 'react';
-import ReactDOM from 'react-dom';
-import type { RenderReturnType } from './types';
-
-type HydrateOrRenderType = (domNode: Element, reactElement: ReactElement) => RenderReturnType;
-const supportsReactCreateRoot = ReactDOM.version &&
-  parseInt(ReactDOM.version.split('.')[0], 10) >= 18;
-
-// TODO: once React dependency is updated to >= 18, we can remove this and just
-// import ReactDOM from 'react-dom/client';
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-let reactDomClient: any;
-if (supportsReactCreateRoot) {
-  // This will never throw an exception, but it's the way to tell Webpack the dependency is optional
-  // https://github.com/webpack/webpack/issues/339#issuecomment-47739112
-  // Unfortunately, it only converts the error to a warning.
-  try {
-    // eslint-disable-next-line global-require,import/no-unresolved
-    reactDomClient = require('react-dom/client');
-  } catch (e) {
-    // We should never get here, but if we do, we'll just use the default ReactDOM
-    // and live with the warning.
-    reactDomClient = ReactDOM;
-  }
-}
-
-export const reactHydrate: HydrateOrRenderType = supportsReactCreateRoot ?
-  reactDomClient.hydrateRoot :
-  (domNode, reactElement) => ReactDOM.hydrate(reactElement, domNode);
-
-export function reactRender(domNode: Element, reactElement: ReactElement): RenderReturnType {
-  if (supportsReactCreateRoot) {
-    const root = reactDomClient.createRoot(domNode);
-    root.render(reactElement);
-    return root;
-  }
-
-  // eslint-disable-next-line react/no-render-return-value
-  return ReactDOM.render(reactElement, domNode);
-}
-
-export default function reactHydrateOrRender(domNode: Element, reactElement: ReactElement, hydrate: boolean): RenderReturnType {
-  return hydrate ? reactHydrate(domNode, reactElement) : reactRender(domNode, reactElement);
-}
+import type { ReactElement } from 'react';
+import ReactDOM from 'react-dom';
+import type { RenderReturnType } from './types';
+import { supportsRootApi } from './reactApis';
+
+type HydrateOrRenderType = (domNode: Element, reactElement: ReactElement) => RenderReturnType;
+
+// TODO: once React dependency is updated to >= 18, we can remove this and just
+// import ReactDOM from 'react-dom/client';
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+let reactDomClient: any;
+if (supportsRootApi) {
+  // This will never throw an exception, but it's the way to tell Webpack the dependency is optional
+  // https://github.com/webpack/webpack/issues/339#issuecomment-47739112
+  // Unfortunately, it only converts the error to a warning.
+  try {
+    // eslint-disable-next-line global-require,import/no-unresolved
+    reactDomClient = require('react-dom/client');
+  } catch (e) {
+    // We should never get here, but if we do, we'll just use the default ReactDOM
+    // and live with the warning.
+    reactDomClient = ReactDOM;
+  }
+}
+
+export const reactHydrate: HydrateOrRenderType = supportsRootApi ?
+  reactDomClient.hydrateRoot :
+  (domNode, reactElement) => ReactDOM.hydrate(reactElement, domNode);
+
+export function reactRender(domNode: Element, reactElement: ReactElement): RenderReturnType {
+  if (supportsRootApi) {
+    const root = reactDomClient.createRoot(domNode);
+    root.render(reactElement);
+    return root;
+  }
+
+  // eslint-disable-next-line react/no-render-return-value
+  return ReactDOM.render(reactElement, domNode);
+}
+
+export default function reactHydrateOrRender(domNode: Element, reactElement: ReactElement, hydrate: boolean): RenderReturnType {
+  return hydrate ? reactHydrate(domNode, reactElement) : reactRender(domNode, reactElement);
+}