Schedule passive effects using the regular ensureRootIsScheduled flow

This treats workInProgressRoot work and rootWithPendingPassiveEffects the same way.

Basically as long as there's some work on the root, yield the current task.
Including passive effects.

packages/react-reconciler/src/ReactFiberRootScheduler.js

@@ -20,6 +20,7 @@ import {
   enableProfilerNestedUpdatePhase,
   enableComponentPerformanceTrack,
   enableSiblingPrerendering,
+  enableYieldingBeforePassive,
 } from 'shared/ReactFeatureFlags';
 import {
   NoLane,
@@ -41,6 +42,8 @@ import {
   getExecutionContext,
   getWorkInProgressRoot,
   getWorkInProgressRootRenderLanes,
+  getRootWithPendingPassiveEffects,
+  getPendingPassiveEffectsLanes,
   isWorkLoopSuspendedOnData,
   performWorkOnRoot,
 } from './ReactFiberWorkLoop';
@@ -324,12 +327,21 @@ function scheduleTaskForRootDuringMicrotask(
   markStarvedLanesAsExpired(root, currentTime);
 
   // Determine the next lanes to work on, and their priority.
+  const rootWithPendingPassiveEffects = getRootWithPendingPassiveEffects();
+  const pendingPassiveEffectsLanes = getPendingPassiveEffectsLanes();
   const workInProgressRoot = getWorkInProgressRoot();
   const workInProgressRootRenderLanes = getWorkInProgressRootRenderLanes();
-  const nextLanes = getNextLanes(
-    root,
-    root === workInProgressRoot ? workInProgressRootRenderLanes : NoLanes,
-  );
+  const nextLanes =
+    enableYieldingBeforePassive && root === rootWithPendingPassiveEffects
+      ? // This will schedule the callback at the priority of the lane but we used to
+        // always schedule it at NormalPriority. Discrete will flush it sync anyway.
+        // So the only difference is Idle and it doesn't seem necessarily right for that
+        // to get upgraded beyond something important just because we're past commit.
+        pendingPassiveEffectsLanes
+      : getNextLanes(
+          root,
+          root === workInProgressRoot ? workInProgressRootRenderLanes : NoLanes,
+        );
 
   const existingCallbackNode = root.callbackNode;
   if (

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -40,6 +40,7 @@ import {
   disableDefaultPropsExceptForClasses,
   enableSiblingPrerendering,
   enableComponentPerformanceTrack,
+  enableYieldingBeforePassive,
 } from 'shared/ReactFeatureFlags';
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 import is from 'shared/objectIs';
@@ -610,7 +611,6 @@ export function getRenderTargetTime(): number {
 
 let legacyErrorBoundariesThatAlreadyFailed: Set<mixed> | null = null;
 
-let rootDoesHavePassiveEffects: boolean = false;
 let rootWithPendingPassiveEffects: FiberRoot | null = null;
 let pendingPassiveEffectsLanes: Lanes = NoLanes;
 let pendingPassiveEffectsRemainingLanes: Lanes = NoLanes;
@@ -638,6 +638,14 @@ export function getWorkInProgressRootRenderLanes(): Lanes {
   return workInProgressRootRenderLanes;
 }
 
+export function getRootWithPendingPassiveEffects(): FiberRoot | null {
+  return rootWithPendingPassiveEffects;
+}
+
+export function getPendingPassiveEffectsLanes(): Lanes {
+  return pendingPassiveEffectsLanes;
+}
+
 export function isWorkLoopSuspendedOnData(): boolean {
   return (
     workInProgressSuspendedReason === SuspendedOnData ||
@@ -3210,12 +3218,6 @@ function commitRootImpl(
     );
   }
 
-  // commitRoot never returns a continuation; it always finishes synchronously.
-  // So we can clear these now to allow a new callback to be scheduled.
-  root.callbackNode = null;
-  root.callbackPriority = NoLane;
-  root.cancelPendingCommit = null;
-
   // Check which lanes no longer have any work scheduled on them, and mark
   // those as finished.
   let remainingLanes = mergeLanes(finishedWork.lanes, finishedWork.childLanes);
@@ -3253,6 +3255,7 @@ function commitRootImpl(
   // might get scheduled in the commit phase. (See #16714.)
   // TODO: Delete all other places that schedule the passive effect callback
   // They're redundant.
+  let rootDoesHavePassiveEffects: boolean = false;
   if (
     // If this subtree rendered with profiling this commit, we need to visit it to log it.
     (enableProfilerTimer &&
@@ -3261,17 +3264,25 @@ function commitRootImpl(
     (finishedWork.subtreeFlags & PassiveMask) !== NoFlags ||
     (finishedWork.flags & PassiveMask) !== NoFlags
   ) {
-    if (!rootDoesHavePassiveEffects) {
-      rootDoesHavePassiveEffects = true;
-      pendingPassiveEffectsRemainingLanes = remainingLanes;
-      pendingPassiveEffectsRenderEndTime = completedRenderEndTime;
-      // workInProgressTransitions might be overwritten, so we want
-      // to store it in pendingPassiveTransitions until they get processed
-      // We need to pass this through as an argument to commitRoot
-      // because workInProgressTransitions might have changed between
-      // the previous render and commit if we throttle the commit
-      // with setTimeout
-      pendingPassiveTransitions = transitions;
+    rootDoesHavePassiveEffects = true;
+    pendingPassiveEffectsRemainingLanes = remainingLanes;
+    pendingPassiveEffectsRenderEndTime = completedRenderEndTime;
+    // workInProgressTransitions might be overwritten, so we want
+    // to store it in pendingPassiveTransitions until they get processed
+    // We need to pass this through as an argument to commitRoot
+    // because workInProgressTransitions might have changed between
+    // the previous render and commit if we throttle the commit
+    // with setTimeout
+    pendingPassiveTransitions = transitions;
+    if (enableYieldingBeforePassive) {
+      // We don't schedule a separate task for flushing passive effects.
+      // Instead, we just rely on ensureRootIsScheduled below to schedule
+      // a callback for us to flush the passive effects.
+    } else {
+      // So we can clear these now to allow a new callback to be scheduled.
+      root.callbackNode = null;
+      root.callbackPriority = NoLane;
+      root.cancelPendingCommit = null;
       scheduleCallback(NormalSchedulerPriority, () => {
         if (enableProfilerTimer && enableComponentPerformanceTrack) {
           // Track the currently executing event if there is one so we can ignore this
@@ -3285,6 +3296,12 @@ function commitRootImpl(
         return null;
       });
     }
+  } else {
+    // If we don't have passive effects, we're not going to need to perform more work
+    // so we can clear the callback now.
+    root.callbackNode = null;
+    root.callbackPriority = NoLane;
+    root.cancelPendingCommit = null;
   }
 
   if (enableProfilerTimer) {
@@ -3441,10 +3458,6 @@ function commitRootImpl(
     onCommitRootTestSelector();
   }
 
-  // Always call this before exiting `commitRoot`, to ensure that any
-  // additional work on this root is scheduled.
-  ensureRootIsScheduled(root);
-
   if (recoverableErrors !== null) {
     // There were errors during this render, but recovered from them without
     // needing to surface it to the UI. We log them here.
@@ -3480,6 +3493,10 @@ function commitRootImpl(
     flushPassiveEffects();
   }
 
+  // Always call this before exiting `commitRoot`, to ensure that any
+  // additional work on this root is scheduled.
+  ensureRootIsScheduled(root);
+
   // Read this again, since a passive effect might have updated it
   remainingLanes = root.pendingLanes;
 
@@ -3648,6 +3665,13 @@ function flushPassiveEffectsImpl(wasDelayedCommit: void | boolean) {
   // because it's only used for profiling), but it's a refactor hazard.
   pendingPassiveEffectsLanes = NoLanes;
 
+  if (enableYieldingBeforePassive) {
+    // We've finished our work for this render pass.
+    root.callbackNode = null;
+    root.callbackPriority = NoLane;
+    root.cancelPendingCommit = null;
+  }
+
   if ((executionContext & (RenderContext | CommitContext)) !== NoContext) {
     throw new Error('Cannot flush passive effects while already rendering.');
   }
@@ -3745,6 +3769,14 @@ function flushPassiveEffectsImpl(wasDelayedCommit: void | boolean) {
     didScheduleUpdateDuringPassiveEffects = false;
   }
 
+  if (enableYieldingBeforePassive) {
+    // Next, we reschedule any remaining work in a new task since it's a new
+    // sequence of work. We wait until the end to do this in case the passive
+    // effect schedules higher priority work than we had remaining. That way
+    // we don't schedule an early callback that gets cancelled anyway.
+    ensureRootIsScheduled(root);
+  }
+
   // TODO: Move to commitPassiveMountEffects
   onPostCommitRootDevTools(root);
   if (enableProfilerTimer && enableProfilerCommitHooks) {