From 5c43c6f02652b0a3f66678a560e84d3110c573b7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Markb=C3=A5ge?= <sebastian@calyptus.eu>
Date: Tue, 13 Sep 2022 11:18:57 -0400
Subject: [PATCH] Unwind the current workInProgress if it's suspended (#25247)

Usually we complete workInProgress before yielding but if that's the
currently suspended one, we don't yet complete it in case we can
immediately unblock it.

If we get interrupted, however, we must unwind it. Where as we usually
assume that we've already completed it.

This shows up when the current work in progress was a Context that pushed
and then it suspends in its immediate children. If we don't unwind,
it won't pop and so we get an imbalance.
---
 .../src/ReactFiberWorkLoop.new.js             |  4 +-
 .../src/ReactFiberWorkLoop.old.js             |  4 +-
 .../src/__tests__/ReactWakeable-test.js       | 47 +++++++++++++++++++
 3 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js b/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
index fa2e8bd3ecf90..6f62aa4b1c0b1 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
+++ b/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
@@ -1651,7 +1651,9 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
   }
 
   if (workInProgress !== null) {
-    let interruptedWork = workInProgress.return;
+    let interruptedWork = workInProgressIsSuspended
+      ? workInProgress
+      : workInProgress.return;
     while (interruptedWork !== null) {
       const current = interruptedWork.alternate;
       unwindInterruptedWork(
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.old.js b/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
index 57880c15e502f..414e2e3343ed0 100644
--- a/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
+++ b/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
@@ -1651,7 +1651,9 @@ function prepareFreshStack(root: FiberRoot, lanes: Lanes): Fiber {
   }
 
   if (workInProgress !== null) {
-    let interruptedWork = workInProgress.return;
+    let interruptedWork = workInProgressIsSuspended
+      ? workInProgress
+      : workInProgress.return;
     while (interruptedWork !== null) {
       const current = interruptedWork.alternate;
       unwindInterruptedWork(
diff --git a/packages/react-reconciler/src/__tests__/ReactWakeable-test.js b/packages/react-reconciler/src/__tests__/ReactWakeable-test.js
index c0ec87f266413..bb15bc7a06862 100644
--- a/packages/react-reconciler/src/__tests__/ReactWakeable-test.js
+++ b/packages/react-reconciler/src/__tests__/ReactWakeable-test.js
@@ -339,4 +339,51 @@ describe('ReactWakeable', () => {
     expect(Scheduler).toFlushWithoutYielding();
     expect(root).toMatchRenderedOutput('AB');
   });
+
+  // @gate enableUseHook
+  test('interrupting while yielded should reset contexts', async () => {
+    let resolve;
+    const promise = new Promise(r => {
+      resolve = r;
+    });
+
+    const Context = React.createContext();
+
+    const lazy = React.lazy(() => {
+      return promise;
+    });
+
+    function ContextText() {
+      return <Text text={use(Context)} />;
+    }
+
+    function App({text}) {
+      return (
+        <div>
+          <Context.Provider value={text}>
+            {lazy}
+            <ContextText />
+          </Context.Provider>
+        </div>
+      );
+    }
+
+    const root = ReactNoop.createRoot();
+    startTransition(() => {
+      root.render(<App text="world" />);
+    });
+    expect(Scheduler).toFlushUntilNextPaint([]);
+    expect(root).toMatchRenderedOutput(null);
+
+    await resolve({default: <Text key="hi" text="Hello " />});
+
+    // Higher priority update that interrupts the first render
+    ReactNoop.flushSync(() => {
+      root.render(<App text="world!" />);
+    });
+
+    expect(Scheduler).toHaveYielded(['Hello ', 'world!']);
+
+    expect(root).toMatchRenderedOutput(<div>Hello world!</div>);
+  });
 });