diff --git a/packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js b/packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js
index d31bf6cfc9a41..fcf42c78ecae1 100644
--- a/packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js
+++ b/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 (
+
+
+
+
+
+
+
+
+ );
+ }
+
+ let finalHTML = ReactDOMServer.renderToString(
+ ,
+ );
+ 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();
+ 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();
+ 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;
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index 4ece9c03d25c2..f1c2b6e03f6e0 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.js
+++ b/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
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.js b/packages/react-reconciler/src/ReactFiberWorkLoop.js
index 9d9dad28a41aa..632ba8dd71cd0 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.js
+++ b/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;
diff --git a/packages/scheduler/src/Scheduler.js b/packages/scheduler/src/Scheduler.js
index 4b36cb7dd6a38..11d47198c23f3 100644
--- a/packages/scheduler/src/Scheduler.js
+++ b/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()
);
}
diff --git a/packages/scheduler/src/__tests__/Scheduler-test.js b/packages/scheduler/src/__tests__/Scheduler-test.js
index e0fd9178453cb..eb632fcf6b90a 100644
--- a/packages/scheduler/src/__tests__/Scheduler-test.js
+++ b/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'),
diff --git a/scripts/jest/matchers/schedulerTestMatchers.js b/scripts/jest/matchers/schedulerTestMatchers.js
index ad357df6aab91..99171dc6d5ec3 100644
--- a/scripts/jest/matchers/schedulerTestMatchers.js
+++ b/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,
};