diff --git a/packages/react-dom/src/__tests__/ReactDOMServerSelectiveHydration-test.internal.js b/packages/react-dom/src/__tests__/ReactDOMServerSelectiveHydration-test.internal.js
index 3469c4faf0683..f4483e9bca8f5 100644
--- a/packages/react-dom/src/__tests__/ReactDOMServerSelectiveHydration-test.internal.js
+++ b/packages/react-dom/src/__tests__/ReactDOMServerSelectiveHydration-test.internal.js
@@ -1496,12 +1496,10 @@ describe('ReactDOMServerSelectiveHydration', () => {
// Start rendering. This will force the first boundary to hydrate
// by scheduling it at one higher pri than Idle.
expect(Scheduler).toFlushAndYieldThrough([
- // An update was scheduled to force hydrate the boundary, but React will
- // continue rendering at Idle until the next time React yields. This is
- // fine though because it will switch to the hydration level when it
- // re-enters the work loop.
'App',
- 'AA',
+
+ // Start hydrating A
+ 'A',
]);
// Hover over A which (could) schedule at one higher pri than Idle.
@@ -1772,4 +1770,72 @@ describe('ReactDOMServerSelectiveHydration', () => {
document.body.removeChild(container);
});
+
+ // @gate experimental || www
+ it('regression test: can unwind context on selective hydration interruption', async () => {
+ const Context = React.createContext('DefaultContext');
+
+ function ContextReader(props) {
+ const value = React.useContext(Context);
+ Scheduler.unstable_yieldValue(value);
+ return null;
+ }
+
+ function Child({text}) {
+ Scheduler.unstable_yieldValue(text);
+ return {text};
+ }
+ const ChildWithBoundary = React.memo(function({text}) {
+ return (
+
+
+
+ );
+ });
+
+ function App({a}) {
+ Scheduler.unstable_yieldValue('App');
+ React.useEffect(() => {
+ Scheduler.unstable_yieldValue('Commit');
+ });
+ return (
+ <>
+
+
+
+
+ >
+ );
+ }
+ const finalHTML = ReactDOMServer.renderToString();
+ expect(Scheduler).toHaveYielded(['App', 'A', 'DefaultContext']);
+ const container = document.createElement('div');
+ container.innerHTML = finalHTML;
+ document.body.appendChild(container);
+
+ const spanA = container.getElementsByTagName('span')[0];
+
+ await act(async () => {
+ const root = ReactDOMClient.hydrateRoot(container, );
+ expect(Scheduler).toFlushAndYieldThrough([
+ 'App',
+ 'DefaultContext',
+ 'Commit',
+ ]);
+
+ TODO_scheduleIdleDOMSchedulerTask(() => {
+ root.render();
+ });
+ expect(Scheduler).toFlushAndYieldThrough(['App', 'A']);
+
+ dispatchClickEvent(spanA);
+ expect(Scheduler).toHaveYielded(['A']);
+ expect(Scheduler).toFlushAndYield([
+ 'App',
+ 'AA',
+ 'DefaultContext',
+ 'Commit',
+ ]);
+ });
+ });
});
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index 78372d1762916..bde65936f1c01 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.js
@@ -282,6 +282,14 @@ import {
const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
+// A special exception that's used to unwind the stack when an update flows
+// into a dehydrated boundary.
+export const SelectiveHydrationException: mixed = new Error(
+ "This is not a real error. It's an implementation detail of React's " +
+ "selective hydration feature. If this leaks into userspace, it's a bug in " +
+ 'React. Please file an issue.',
+);
+
let didReceiveUpdate: boolean = false;
let didWarnAboutBadClass;
@@ -2860,6 +2868,16 @@ function updateDehydratedSuspenseComponent(
attemptHydrationAtLane,
eventTime,
);
+
+ // Throw a special object that signals to the work loop that it should
+ // interrupt the current render.
+ //
+ // Because we're inside a React-only execution stack, we don't
+ // strictly need to throw here — we could instead modify some internal
+ // work loop state. But using an exception means we don't need to
+ // check for this case on every iteration of the work loop. So doing
+ // it this way moves the check out of the fast path.
+ throw SelectiveHydrationException;
} 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
@@ -2870,15 +2888,17 @@ function updateDehydratedSuspenseComponent(
}
}
- // If we have scheduled higher pri work above, this will just abort the render
- // since we now have higher priority work. We'll try to infinitely suspend until
- // we yield. TODO: We could probably just force yielding earlier instead.
- renderDidSuspendDelayIfPossible();
- // If we rendered synchronously, we won't yield so have to render something.
- // This will cause us to delete any existing content.
+ // If we did not selectively hydrate, we'll continue rendering without
+ // hydrating. Mark this tree as suspended to prevent it from committing
+ // outside a transition.
+ //
+ // This path should only happen if the hydration lane already suspended.
+ // Currently, it also happens during sync updates because there is no
+ // hydration lane for sync updates.
// TODO: We should ideally have a sync hydration lane that we can apply to do
// a pass where we hydrate this subtree in place using the previous Context and then
// reapply the update afterwards.
+ renderDidSuspendDelayIfPossible();
return retrySuspenseComponentWithoutHydrating(
current,
workInProgress,
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.js b/packages/react-reconciler/src/ReactFiberWorkLoop.js
index 0c86cbb324020..2a1166cc43c1a 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.js
+++ b/packages/react-reconciler/src/ReactFiberWorkLoop.js
@@ -175,6 +175,7 @@ import {
} from './ReactEventPriorities';
import {requestCurrentTransition, NoTransition} from './ReactFiberTransition';
import {
+ SelectiveHydrationException,
beginWork as originalBeginWork,
replayFunctionComponent,
} from './ReactFiberBeginWork';
@@ -316,13 +317,14 @@ let workInProgress: Fiber | null = null;
// The lanes we're rendering
let workInProgressRootRenderLanes: Lanes = NoLanes;
-opaque type SuspendedReason = 0 | 1 | 2 | 3 | 4 | 5;
+opaque type SuspendedReason = 0 | 1 | 2 | 3 | 4 | 5 | 6;
const NotSuspended: SuspendedReason = 0;
const SuspendedOnError: SuspendedReason = 1;
const SuspendedOnData: SuspendedReason = 2;
const SuspendedOnImmediate: SuspendedReason = 3;
const SuspendedOnDeprecatedThrowPromise: SuspendedReason = 4;
const SuspendedAndReadyToUnwind: SuspendedReason = 5;
+const SuspendedOnHydration: SuspendedReason = 6;
// When this is true, the work-in-progress fiber just suspended (or errored) and
// we've yet to unwind the stack. In some cases, we may yield to the main thread
@@ -1701,6 +1703,31 @@ export function getRenderLanes(): Lanes {
return renderLanes;
}
+function resetWorkInProgressStack() {
+ if (workInProgress === null) return;
+ let interruptedWork;
+ if (workInProgressSuspendedReason === NotSuspended) {
+ // Normal case. Work-in-progress hasn't started yet. Unwind all
+ // its parents.
+ interruptedWork = workInProgress.return;
+ } else {
+ // Work-in-progress is in suspended state. Reset the work loop and unwind
+ // both the suspended fiber and all its parents.
+ resetSuspendedWorkLoopOnUnwind();
+ interruptedWork = workInProgress;
+ }
+ while (interruptedWork !== null) {
+ const current = interruptedWork.alternate;
+ unwindInterruptedWork(
+ current,
+ interruptedWork,
+ workInProgressRootRenderLanes,
+ );
+ interruptedWork = interruptedWork.return;
+ }
+ workInProgress = null;
+}
+
function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
root.finishedWork = null;
root.finishedLanes = NoLanes;
@@ -1714,28 +1741,7 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
cancelTimeout(timeoutHandle);
}
- if (workInProgress !== null) {
- let interruptedWork;
- if (workInProgressSuspendedReason === NotSuspended) {
- // Normal case. Work-in-progress hasn't started yet. Unwind all
- // its parents.
- interruptedWork = workInProgress.return;
- } else {
- // Work-in-progress is in suspended state. Reset the work loop and unwind
- // both the suspended fiber and all its parents.
- resetSuspendedWorkLoopOnUnwind();
- interruptedWork = workInProgress;
- }
- while (interruptedWork !== null) {
- const current = interruptedWork.alternate;
- unwindInterruptedWork(
- current,
- interruptedWork,
- workInProgressRootRenderLanes,
- );
- interruptedWork = interruptedWork.return;
- }
- }
+ resetWorkInProgressStack();
workInProgressRoot = root;
const rootWorkInProgress = createWorkInProgress(root.current, null);
workInProgress = rootWorkInProgress;
@@ -1797,6 +1803,17 @@ function handleThrow(root, thrownValue): void {
workInProgressSuspendedReason = shouldAttemptToSuspendUntilDataResolves()
? SuspendedOnData
: SuspendedOnImmediate;
+ } else if (thrownValue === SelectiveHydrationException) {
+ // An update flowed into a dehydrated boundary. Before we can apply the
+ // update, we need to finish hydrating. Interrupt the work-in-progress
+ // render so we can restart at the hydration lane.
+ //
+ // The ideal implementation would be able to switch contexts without
+ // unwinding the current stack.
+ //
+ // We could name this something more general but as of now it's the only
+ // case where we think this should happen.
+ workInProgressSuspendedReason = SuspendedOnHydration;
} else {
// This is a regular error.
const isWakeable =
@@ -2038,7 +2055,7 @@ function renderRootSync(root: FiberRoot, lanes: Lanes) {
markRenderStarted(lanes);
}
- do {
+ outer: do {
try {
if (
workInProgressSuspendedReason !== NotSuspended &&
@@ -2054,11 +2071,23 @@ function renderRootSync(root: FiberRoot, lanes: Lanes) {
// function and fork the behavior some other way.
const unitOfWork = workInProgress;
const thrownValue = workInProgressThrownValue;
- workInProgressSuspendedReason = NotSuspended;
- workInProgressThrownValue = null;
- unwindSuspendedUnitOfWork(unitOfWork, thrownValue);
-
- // Continue with the normal work loop.
+ switch (workInProgressSuspendedReason) {
+ case SuspendedOnHydration: {
+ // Selective hydration. An update flowed into a dehydrated tree.
+ // Interrupt the current render so the work loop can switch to the
+ // hydration lane.
+ resetWorkInProgressStack();
+ workInProgressRootExitStatus = RootDidNotComplete;
+ break outer;
+ }
+ default: {
+ // Continue with the normal work loop.
+ workInProgressSuspendedReason = NotSuspended;
+ workInProgressThrownValue = null;
+ unwindSuspendedUnitOfWork(unitOfWork, thrownValue);
+ break;
+ }
+ }
}
workLoopSync();
break;
@@ -2216,6 +2245,14 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
unwindSuspendedUnitOfWork(unitOfWork, thrownValue);
break;
}
+ case SuspendedOnHydration: {
+ // Selective hydration. An update flowed into a dehydrated tree.
+ // Interrupt the current render so the work loop can switch to the
+ // hydration lane.
+ resetWorkInProgressStack();
+ workInProgressRootExitStatus = RootDidNotComplete;
+ break outer;
+ }
default: {
throw new Error(
'Unexpected SuspendedReason. This is a bug in React.',
@@ -3741,6 +3778,7 @@ if (__DEV__ && replayFailedUnitOfWorkWithInvokeGuardedCallback) {
if (
didSuspendOrErrorWhileHydratingDEV() ||
originalError === SuspenseException ||
+ originalError === SelectiveHydrationException ||
(originalError !== null &&
typeof originalError === 'object' &&
typeof originalError.then === 'function')