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.

Author: sebmarkbage

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();
+      });
     });
   });
 });