Track interactions scheduled during (sync) render phase

Author: Brian Vaughn

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -1162,14 +1162,6 @@ function renderRoot(
 
   const expirationTime = root.nextExpirationTimeToWorkOn;
 
-  let prevInteractions: Set<Interaction> = (null: any);
-  if (enableSchedulerTracking) {
-    // We're about to start new tracked work.
-    // Restore pending interactions so cascading work triggered during the render phase will be accounted for.
-    prevInteractions = __interactionsRef.current;
-    __interactionsRef.current = root.memoizedInteractions;
-  }
-
   // Check if we're starting from a fresh stack, or if we're resuming from
   // previously yielded work.
   if (
@@ -1232,6 +1224,14 @@ function renderRoot(
     }
   }
 
+  let prevInteractions: Set<Interaction> = (null: any);
+  if (enableSchedulerTracking) {
+    // We're about to start new tracked work.
+    // Restore pending interactions so cascading work triggered during the render phase will be accounted for.
+    prevInteractions = __interactionsRef.current;
+    __interactionsRef.current = root.memoizedInteractions;
+  }
+
   let didFatal = false;
 
   startWorkLoopTimer(nextUnitOfWork);

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

@@ -1410,6 +1410,34 @@ describe('Profiler', () => {
       });
     });
 
+    it('should properly track work scheduled during the begin render phase', () => {
+      const callback = jest.fn();
+      let wrapped;
+      const Component = jest.fn(() => {
+        wrapped = SchedulerTracking.unstable_wrap(callback);
+        return null;
+      });
+
+      let interaction;
+      SchedulerTracking.unstable_track('event', mockNow(), () => {
+        const interactions = SchedulerTracking.unstable_getCurrent();
+        expect(interactions.size).toBe(1);
+        interaction = Array.from(interactions)[0];
+        ReactTestRenderer.create(<Component />);
+      });
+
+      expect(Component).toHaveBeenCalledTimes(1);
+      expect(onInteractionScheduledWorkCompleted).not.toHaveBeenCalled();
+      expect(callback).not.toHaveBeenCalled();
+
+      wrapped();
+      expect(callback).toHaveBeenCalledTimes(1);
+      expect(onInteractionScheduledWorkCompleted).toHaveBeenCalledTimes(1);
+      expect(
+        onInteractionScheduledWorkCompleted,
+      ).toHaveBeenLastNotifiedOfInteraction(interaction);
+    });
+
     it('should associate tracked events with their subsequent commits', () => {
       let instance = null;