Mark spawned work for client-rendered suspense boundary (#16341)

Currently this is getting marked as Never which is the normal continuation
for a dehydrated boundary, but if it is client-rendered it has a higher
priority. That causes us to drop the interaction tracing for that render.

This colocates the marking where we actually set the expiration time.

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -1903,12 +1903,20 @@ function updateDehydratedSuspenseComponent(
  // they should be.
  let serverDisplayTime = requestCurrentTime();
  // Schedule a normal pri update to render this content.
- workInProgress.expirationTime = computeAsyncExpiration(serverDisplayTime);
+ let newExpirationTime = computeAsyncExpiration(serverDisplayTime);
+ if (enableSchedulerTracing) {
+ markSpawnedWork(newExpirationTime);
+ }
+ workInProgress.expirationTime = newExpirationTime;
  } else {
  // We'll continue hydrating the rest at offscreen priority since we'll already
  // be showing the right content coming from the server, it is no rush.
  workInProgress.expirationTime = Never;
+ if (enableSchedulerTracing) {
+ markSpawnedWork(Never);
+ }
  }
+
  return null;
  }
 

packages/react-reconciler/src/ReactFiberCompleteWork.js

@@ -979,9 +979,6 @@ function completeWork(
  'A dehydrated suspense component was completed without a hydrated node. ' +
  'This is probably a bug in React.',
  );
- if (enableSchedulerTracing) {
- markSpawnedWork(Never);
- }
  skipPastDehydratedSuspenseInstance(workInProgress);
  } else {
  // We should never have been in a hydration state if we didn't have a current.

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

@@ -710,6 +710,78 @@ describe('ReactDOMTracing', () => {
 
  done();
  });
+
+ it('traces interaction across client-rendered hydration', async done => {
+ let suspend = false;
+ let promise = new Promise(() => {});
+ let ref = React.createRef();
+
+ function Child() {
+ if (suspend) {
+ throw promise;
+ } else {
+ return 'Hello';
+ }
+ }
+
+ function App() {
+ return (
+ <div>
+ <React.Suspense fallback="Loading...">
+ <span ref={ref}>
+ <Child />
+ </span>
+ </React.Suspense>
+ </div>
+ );
+ }
+
+ // Render the final HTML.
+ suspend = true;
+ const finalHTML = ReactDOMServer.renderToString(<App />);
+
+ const container = document.createElement('div');
+ container.innerHTML = finalHTML;
+
+ let interaction;
+
+ const root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+
+ // Hydrate without suspending to fill in the client-rendered content.
+ suspend = false;
+ SchedulerTracing.unstable_trace('initialization', 0, () => {
+ interaction = Array.from(SchedulerTracing.unstable_getCurrent())[0];
+
+ root.render(<App />);
+ });
+
+ expect(onWorkStopped).toHaveBeenCalledTimes(1);
+
+ // Advance time a bit so that we get into a new expiration bucket.
+ Scheduler.unstable_advanceTime(300);
+ jest.advanceTimersByTime(300);
+
+ Scheduler.unstable_flushAll();
+ jest.runAllTimers();
+
+ expect(ref.current).not.toBe(null);
+
+ // We should've had two commits that was traced.
+ // First one that hydrates the parent, and then one that hydrates
+ // the boundary at higher than Never priority.
+ expect(onWorkStopped).toHaveBeenCalledTimes(3);
+
+ expect(onInteractionTraced).toHaveBeenCalledTimes(1);
+ expect(onInteractionTraced).toHaveBeenLastNotifiedOfInteraction(
+ interaction,
+ );
+ expect(onInteractionScheduledWorkCompleted).toHaveBeenCalledTimes(1);
+ expect(
+ onInteractionScheduledWorkCompleted,
+ ).toHaveBeenLastNotifiedOfInteraction(interaction);
+
+ done();
+ });
  });
  });
 });