Add fix for unwinding context in sync update

Author: tyao1

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

@@ -1886,8 +1886,7 @@ describe('ReactDOMServerSelectiveHydration', () => {
     expect(initialSpan).toBe(spanRef);
   });
 
-  // @gate experimental || www
-  it('regression test: can unwind context on selective hydration interruption', async () => {
+  it('regression test: can unwind context on selective hydration interruption for sync updates', async () => {
     const Context = React.createContext('DefaultContext');
 
     function ContextReader(props) {
@@ -1926,9 +1925,6 @@ describe('ReactDOMServerSelectiveHydration', () => {
     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" />);
@@ -1938,14 +1934,12 @@ describe('ReactDOMServerSelectiveHydration', () => {
         'Commit',
       ]);
 
-      TODO_scheduleIdleDOMSchedulerTask(() => {
+      ReactDOM.flushSync(() => {
         root.render(<App a="AA" />);
       });
-      expect(Scheduler).toFlushAndYieldThrough(['App', 'A']);
-
-      dispatchClickEvent(spanA);
-      expect(Scheduler).toHaveYielded(['A']);
-      expect(Scheduler).toFlushAndYield([
+      expect(Scheduler).toHaveYielded([
+        'App',
+        'A',
         'App',
         'AA',
         'DefaultContext',

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -2072,7 +2072,6 @@ function renderRootSync(root: FiberRoot, lanes: Lanes) {
             // Selective hydration. An update flowed into a dehydrated tree.
             // Interrupt the current render so the work loop can switch to the
             // hydration lane.
-            workInProgress = null;
             workInProgressRootExitStatus = RootDidNotComplete;
             break outer;
           }
@@ -2097,8 +2096,12 @@ function renderRootSync(root: FiberRoot, lanes: Lanes) {
   popDispatcher(prevDispatcher);
   popCacheDispatcher(prevCacheDispatcher);
 
-  if (workInProgress !== null) {
-    // This is a sync render, so we should have finished the whole tree.
+  if (
+    workInProgress !== null &&
+    workInProgressRootExitStatus !== RootDidNotComplete
+  ) {
+    // This is a sync render, so we should have finished the whole tree unless
+    // it is interrupted by selective hydration.
     throw new Error(
       'Cannot commit an incomplete root. This error is likely caused by a ' +
         'bug in React. Please file an issue.',