[Hydration Bugfix] Updates to dehydrated content when disableSchedulerTimeoutBasedOnReactExpirationTime is enabled

packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js

@@ -497,6 +497,103 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(span.className).toBe('hi');
   });
 
+  it('regression test: an update to dehydrated content forces a restart to hydrate it first', async () => {
+    // Replicates a bug where React would fail to restart and hydrate when
+    // `disableSchedulerTimeoutBasedOnReactExpirationTime` is true.
+    jest.resetModuleRegistry();
+
+    ReactFeatureFlags = require('shared/ReactFeatureFlags');
+    ReactFeatureFlags.enableSuspenseServerRenderer = true;
+    ReactFeatureFlags.enableSuspenseCallback = true;
+    ReactFeatureFlags.enableFlareAPI = true;
+    ReactFeatureFlags.debugRenderPhaseSideEffects = false;
+    ReactFeatureFlags.debugRenderPhaseSideEffectsForStrictMode = false;
+    // NOTE: Once this feature flag is removed, the preceding test is likely
+    // sufficient to cover this regression case; however, it might be worth
+    // keeping this test around anyway as an extra precaution.
+    ReactFeatureFlags.disableSchedulerTimeoutBasedOnReactExpirationTime = true;
+
+    React = require('react');
+    ReactDOM = require('react-dom');
+    act = require('react-dom/test-utils').act;
+    ReactDOMServer = require('react-dom/server');
+    Scheduler = require('scheduler');
+    Suspense = React.Suspense;
+    SuspenseList = React.unstable_SuspenseList;
+
+    function Text({text}) {
+      Scheduler.unstable_yieldValue(
+        text + ' [' + Scheduler.unstable_getCurrentPriorityLevel() + ']',
+      );
+      return text;
+    }
+
+    function Parent({text, className}) {
+      Scheduler.unstable_yieldValue(
+        'Parent [' + Scheduler.unstable_getCurrentPriorityLevel() + ']',
+      );
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <Text text="Sibling" />
+            <span className={className}>
+              <Text text={text} />
+            </span>
+          </Suspense>
+        </div>
+      );
+    }
+
+    let finalHTML = ReactDOMServer.renderToString(
+      <Parent text="A" className="A" />,
+    );
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+    expect(Scheduler).toHaveYielded(['Parent [3]', 'Sibling [3]', 'A [3]']);
+
+    let span = container.getElementsByTagName('span')[0];
+
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+
+    await act(async () => {
+      // The first render hydrates the shell (everything outside the Suspense
+      // boundary) at normal priority.
+      root.render(<Parent text="A" className="A" />);
+      expect(Scheduler).toFlushUntilNextPaint(['Parent [3]']);
+      expect(span.textContent).toBe('A');
+      expect(span.className).toBe('A');
+
+      // The second render hydrates the child (inside the Suspense boundary)
+      // at idle priority. Let's partially render and stop before it finishes.
+      expect(Scheduler).toFlushAndYieldThrough(['Sibling [5]']);
+
+      // Before the idle hydration finishes, interrupt with an update. This will
+      // start over at normal priority.
+      root.render(<Parent text="B" className="B" />);
+      expect(Scheduler).toFlushUntilNextPaint([
+        'Parent [3]',
+
+        // The Suspense boundary hasn't hydrated yet, but we need to send it new
+        // props. We need to finish hydrating first. So we'll interrupt the
+        // current render, finish hydrating, then start the update again. The
+        // following two entries occur at slightly higher priority. (Parent
+        // doesn't appear as an entry because it already hydrated.)
+        'Sibling [3]',
+        'A [3]',
+      ]);
+      // Hydrating has finished
+      expect(span.textContent).toBe('A');
+      expect(span.className).toBe('A');
+      expect(span).toBe(container.getElementsByTagName('span')[0]);
+
+      // Now that the boundary is hydrated, we can perform the update.
+      expect(Scheduler).toFlushAndYield(['Parent [3]', 'Sibling [3]', 'B [3]']);
+      expect(span.textContent).toBe('B');
+      expect(span.className).toBe('B');
+      expect(span).toBe(container.getElementsByTagName('span')[0]);
+    });
+  });
+
   it('shows the fallback if props have changed before hydration completes and is still suspended', async () => {
     let suspend = false;
     let resolve;

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -2046,8 +2046,16 @@ function updateDehydratedSuspenseComponent(
         // at even higher pri.
         let attemptHydrationAtExpirationTime = renderExpirationTime + 1;
         suspenseState.retryTime = attemptHydrationAtExpirationTime;
+        // TODO: This happens to abort the current render and switch to a
+        // hydration pass, because the priority of the new task is slightly
+        // higher priority, which causes the work loop to cancel the current
+        // Scheduler task via `Scheduler.cancelCallback`. But we should probably
+        // model this entirely within React instead of relying on Scheduler's
+        // semantics for canceling in-progress tasks. I've chosen this approach
+        // for now since it's a fairly non-invasive change and it conceptually
+        // matches how other types of interuptions (e.g. due to input events)
+        // already work.
         scheduleWork(current, attemptHydrationAtExpirationTime);
-        // TODO: Early abort this render.
       } else {
         // We have already tried to ping at a higher priority than we're rendering with
         // so if we got here, we must have failed to hydrate at those levels. We must

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -515,6 +515,9 @@ function scheduleCallbackForRoot(
     // New callback has higher priority than the existing one.
     const existingCallbackNode = root.callbackNode;
     if (existingCallbackNode !== null) {
+      // If this happens during render, this task represents the currently
+      // running task. Canceling has the effect of interrupting the render and
+      // starting over at the higher priority.
       cancelCallback(existingCallbackNode);
     }
     root.callbackExpirationTime = expirationTime;

packages/scheduler/src/Scheduler.js

@@ -75,6 +75,8 @@ var currentPriorityLevel = NormalPriority;
 // This is set while performing work, to prevent re-entrancy.
 var isPerformingWork = false;
 
+var didCancelCurrentTask = false;
+
 var isHostCallbackScheduled = false;
 var isHostTimeoutScheduled = false;
 
@@ -184,16 +186,20 @@ function workLoop(hasTimeRemaining, initialTime) {
       markTaskRun(currentTask, currentTime);
       const continuationCallback = callback(didUserCallbackTimeout);
       currentTime = getCurrentTime();
-      if (typeof continuationCallback === 'function') {
-        currentTask.callback = continuationCallback;
-        markTaskYield(currentTask, currentTime);
+      if (didCancelCurrentTask) {
+        didCancelCurrentTask = false;
       } else {
-        if (enableProfiling) {
-          markTaskCompleted(currentTask, currentTime);
-          currentTask.isQueued = false;
-        }
-        if (currentTask === peek(taskQueue)) {
-          pop(taskQueue);
+        if (typeof continuationCallback === 'function') {
+          currentTask.callback = continuationCallback;
+          markTaskYield(currentTask, currentTime);
+        } else {
+          if (enableProfiling) {
+            markTaskCompleted(currentTask, currentTime);
+            currentTask.isQueued = false;
+          }
+          if (currentTask === peek(taskQueue)) {
+            pop(taskQueue);
+          }
         }
       }
       advanceTimers(currentTime);
@@ -389,6 +395,9 @@ function unstable_cancelCallback(task) {
   // remove from the queue because you can't remove arbitrary nodes from an
   // array based heap, only the first one.)
   task.callback = null;
+  if (task === currentTask) {
+    didCancelCurrentTask = true;
+  }
 }
 
 function unstable_getCurrentPriorityLevel() {
@@ -406,6 +415,7 @@ function unstable_shouldYield() {
       firstTask.callback !== null &&
       firstTask.startTime <= currentTime &&
       firstTask.expirationTime < currentTask.expirationTime) ||
+    didCancelCurrentTask ||
     shouldYieldToHost()
   );
 }

packages/scheduler/src/__tests__/Scheduler-test.js

@@ -288,6 +288,26 @@ describe('Scheduler', () => {
     expect(Scheduler).toFlushWithoutYielding();
   });
 
+  it('cancelling the currently running task', () => {
+    const task = scheduleCallback(NormalPriority, () => {
+      Scheduler.unstable_yieldValue('Start');
+      for (let i = 0; !shouldYield(); i++) {
+        if (i === 5) {
+          // Canceling the current task will cause `shouldYield` to return
+          // `true`. Otherwise this would infinite loop.
+          Scheduler.unstable_yieldValue('Cancel');
+          cancelCallback(task);
+        }
+      }
+      Scheduler.unstable_yieldValue('Finish');
+      // The continuation should be ignored, since the task was
+      // already canceled.
+      return () => Scheduler.unstable_yieldValue('Continuation');
+    });
+
+    expect(Scheduler).toFlushAndYield(['Start', 'Cancel', 'Finish']);
+  });
+
   it('top-level immediate callbacks fire in a subsequent task', () => {
     scheduleCallback(ImmediatePriority, () =>
       Scheduler.unstable_yieldValue('A'),

scripts/jest/matchers/schedulerTestMatchers.js

@@ -57,6 +57,15 @@ function toFlushExpired(Scheduler, expectedYields) {
   });
 }
 
+function toFlushUntilNextPaint(Scheduler, expectedYields) {
+  assertYieldsWereCleared(Scheduler);
+  Scheduler.unstable_flushUntilNextPaint();
+  const actualYields = Scheduler.unstable_clearYields();
+  return captureAssertion(() => {
+    expect(actualYields).toEqual(expectedYields);
+  });
+}
+
 function toHaveYielded(Scheduler, expectedYields) {
   return captureAssertion(() => {
     const actualYields = Scheduler.unstable_clearYields();
@@ -78,6 +87,7 @@ module.exports = {
   toFlushAndYieldThrough,
   toFlushWithoutYielding,
   toFlushExpired,
+  toFlushUntilNextPaint,
   toHaveYielded,
   toFlushAndThrow,
 };