wrap SuspenseInstanceRetry callback so scheduler waits for it (#18805)

packages/react-reconciler/src/ReactFiberBeginWork.new.js

@@ -198,6 +198,7 @@ import {
   getWorkInProgressRoot,
   pushRenderLanes,
 } from './ReactFiberWorkLoop.new';
+import {unstable_wrap as Schedule_tracing_wrap} from 'scheduler/tracing';
 
 import {disableLogs, reenableLogs} from 'shared/ConsolePatchingDev';
 
@@ -2336,10 +2337,11 @@ function updateDehydratedSuspenseComponent(
     // Leave the child in place. I.e. the dehydrated fragment.
     workInProgress.child = current.child;
     // Register a callback to retry this boundary once the server has sent the result.
-    registerSuspenseInstanceRetry(
-      suspenseInstance,
-      retryDehydratedSuspenseBoundary.bind(null, current),
-    );
+    let retry = retryDehydratedSuspenseBoundary.bind(null, current);
+    if (enableSchedulerTracing) {
+      retry = Schedule_tracing_wrap(retry);
+    }
+    registerSuspenseInstanceRetry(suspenseInstance, retry);
     return null;
   } else {
     // This is the first attempt.

packages/react-reconciler/src/ReactFiberBeginWork.old.js

@@ -180,6 +180,7 @@ import {
   markUnprocessedUpdateTime,
   getWorkInProgressRoot,
 } from './ReactFiberWorkLoop.old';
+import {unstable_wrap as Schedule_tracing_wrap} from 'scheduler/tracing';
 
 import {disableLogs, reenableLogs} from 'shared/ConsolePatchingDev';
 
@@ -2303,10 +2304,11 @@ function updateDehydratedSuspenseComponent(
     // Leave the child in place. I.e. the dehydrated fragment.
     workInProgress.child = current.child;
     // Register a callback to retry this boundary once the server has sent the result.
-    registerSuspenseInstanceRetry(
-      suspenseInstance,
-      retryDehydratedSuspenseBoundary.bind(null, current),
-    );
+    let retry = retryDehydratedSuspenseBoundary.bind(null, current);
+    if (enableSchedulerTracing) {
+      retry = Schedule_tracing_wrap(retry);
+    }
+    registerSuspenseInstanceRetry(suspenseInstance, retry);
     return null;
   } else {
     // This is the first attempt.

packages/react/src/__tests__/ReactDOMTracing-test.internal.js

@@ -719,6 +719,76 @@ describe('ReactDOMTracing', () => {
         ).toHaveBeenLastNotifiedOfInteraction(interaction);
       });
 
+      // @gate experimental
+      it('traces interaction across suspended hydration from server', async () => {
+        // Copied from ReactDOMHostConfig.js
+        const SUSPENSE_START_DATA = '$';
+        const SUSPENSE_PENDING_START_DATA = '$?';
+
+        const ref = React.createRef();
+
+        function App() {
+          return (
+            <React.Suspense fallback="Loading...">
+              <span ref={ref}>Hello</span>
+            </React.Suspense>
+          );
+        }
+
+        const container = document.createElement('div');
+
+        // Render the final HTML.
+        const finalHTML = ReactDOMServer.renderToString(<App />);
+
+        // Replace the marker with a pending state.
+        // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Guide/Regular_Expressions#Escaping
+        const escapedMarker = SUSPENSE_START_DATA.replace(
+          /[.*+\-?^${}()|[\]\\]/g,
+          '\\$&',
+        );
+        container.innerHTML = finalHTML.replace(
+          new RegExp(escapedMarker, 'g'),
+          SUSPENSE_PENDING_START_DATA,
+        );
+
+        let interaction;
+
+        const root = ReactDOM.createRoot(container, {hydrate: true});
+
+        // Start hydrating but simulate blocking for suspense data from the server.
+        SchedulerTracing.unstable_trace('initialization', 0, () => {
+          interaction = Array.from(SchedulerTracing.unstable_getCurrent())[0];
+
+          root.render(<App />);
+        });
+        Scheduler.unstable_flushAll();
+        jest.runAllTimers();
+
+        expect(ref.current).toBe(null);
+        expect(onInteractionTraced).toHaveBeenCalledTimes(1);
+        expect(onInteractionTraced).toHaveBeenLastNotifiedOfInteraction(
+          interaction,
+        );
+        expect(onInteractionScheduledWorkCompleted).not.toHaveBeenCalled();
+
+        // Unblock rendering, pretend the content is injected by the server.
+        const startNode = container.childNodes[0];
+        expect(startNode).not.toBe(null);
+        expect(startNode.nodeType).toBe(Node.COMMENT_NODE);
+
+        startNode.textContent = SUSPENSE_START_DATA;
+        startNode._reactRetry();
+
+        Scheduler.unstable_flushAll();
+        jest.runAllTimers();
+
+        expect(ref.current).not.toBe(null);
+        expect(onInteractionScheduledWorkCompleted).toHaveBeenCalledTimes(1);
+        expect(
+          onInteractionScheduledWorkCompleted,
+        ).toHaveBeenLastNotifiedOfInteraction(interaction);
+      });
+
       // @gate experimental
       it('traces interaction across client-rendered hydration', () => {
         let suspend = false;