From 049d84c44fe3768d10e798e65972e6b2aab47724 Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@chromatic.com>
Date: Tue, 19 Jul 2022 16:49:01 +1000
Subject: [PATCH] Use `useLayoutEffect` for callback

---
 code/renderers/react/src/render.tsx | 30 ++++++++++++++++++++++++-----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/code/renderers/react/src/render.tsx b/code/renderers/react/src/render.tsx
index 2bcad1628cc5..bf81b32b0dbd 100644
--- a/code/renderers/react/src/render.tsx
+++ b/code/renderers/react/src/render.tsx
@@ -1,7 +1,15 @@
 // @ts-ignore
 import global from 'global';
 
-import React, { Component as ReactComponent, FC, ReactElement, StrictMode, Fragment } from 'react';
+import React, {
+  Component as ReactComponent,
+  FC,
+  ReactElement,
+  StrictMode,
+  Fragment,
+  useLayoutEffect,
+  useRef,
+} from 'react';
 import ReactDOM, { version as reactDomVersion } from 'react-dom';
 import type { Root as ReactRoot } from 'react-dom/client';
 
@@ -26,16 +34,28 @@ export const render: ArgsStoryFn<ReactFramework> = (args, context) => {
   return <Component {...args} />;
 };
 
+const WithCallback: FC<{ callback: () => void; children: ReactElement }> = ({
+  callback,
+  children,
+}) => {
+  // See https://github.com/reactwg/react-18/discussions/5#discussioncomment-2276079
+  const once = useRef(false);
+  useLayoutEffect(() => {
+    if (once.current) return;
+    once.current = true;
+    callback();
+  }, [callback]);
+
+  return children;
+};
+
 const renderElement = async (node: ReactElement, el: Element) => {
   // Create Root Element conditionally for new React 18 Root Api
   const root = await getReactRoot(el);
 
   return new Promise((resolve) => {
     if (root) {
-      root.render(node);
-      setTimeout(() => {
-        resolve(null);
-      }, 0);
+      root.render(<WithCallback callback={() => resolve(null)}>{node}</WithCallback>);
     } else {
       ReactDOM.render(node, el, () => resolve(null));
     }