Fix unwinding context during selective hydration

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 <span>{text}</span>;
+    }
+    const ChildWithBoundary = React.memo(function({text}) {
+      return (
+        <Suspense fallback="Loading...">
+          <Child text={text} />
+        </Suspense>
+      );
+    });
+
+    function App({a}) {
+      Scheduler.unstable_yieldValue('App');
+      React.useEffect(() => {
+        Scheduler.unstable_yieldValue('Commit');
+      });
+      return (
+        <>
+          <Context.Provider value="SiblingContext">
+            <ChildWithBoundary text={a} />
+          </Context.Provider>
+          <ContextReader />
+        </>
+      );
+    }
+    const finalHTML = ReactDOMServer.renderToString(<App a="A" />);
+    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, <App a="A" />);
+      expect(Scheduler).toFlushAndYieldThrough([
+        'App',
+        'DefaultContext',
+        'Commit',
+      ]);
+
+      TODO_scheduleIdleDOMSchedulerTask(() => {
+        root.render(<App a="AA" />);
+      });
+      expect(Scheduler).toFlushAndYieldThrough(['App', 'A']);
+
+      dispatchClickEvent(spanA);
+      expect(Scheduler).toHaveYielded(['A']);
+      expect(Scheduler).toFlushAndYield([
+        'App',
+        'AA',
+        'DefaultContext',
+        'Commit',
+      ]);
+    });
+  });
 });

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,

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')