From 8a1e64b39522422b75093aee08dd924b0b4a5e71 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Mon, 12 Sep 2022 23:26:37 -0400
Subject: [PATCH] Unwind the current workInProgress if it's suspended

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 5e9edf7945653..913d5cb96b730 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 d8a0bc883f4d7..0cefbdc893f73 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>);
+  });
 });