From 9cd52b27fe7be3c320d1e94a016fc416e66569b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Markb=C3=A5ge?= <sema@fb.com>
Date: Fri, 23 Apr 2021 15:24:10 -0400
Subject: [PATCH] Restore context after an error happens (#21341)

Typically we don't need to restore the context here because we assume that
we'll terminate the rest of the subtree so we don't need the correct
context since we're not rendering any siblings.

However, after a nested suspense boundary we need to restore the context.
The boundary could do this but since we're already doing this in the
suspense branch of renderNode, we might as well do it in the error case
which isn't very perf sensitive anyway.
---
 .../src/__tests__/ReactDOMFizzServer-test.js  | 67 +++++++++++++++++++
 packages/react-server/src/ReactFizzServer.js  |  7 ++
 2 files changed, 74 insertions(+)

diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
index 13774ec2de07..a42056ab2492 100644
--- a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
+++ b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
@@ -796,6 +796,73 @@ describe('ReactDOMFizzServer', () => {
     );
   });
 
+  // @gate experimental
+  it('should recover the outer context when an error happens inside a provider', async () => {
+    const ContextA = React.createContext('A0');
+    const ContextB = React.createContext('B0');
+
+    function PrintA() {
+      return (
+        <ContextA.Consumer>{value => <Text text={value} />}</ContextA.Consumer>
+      );
+    }
+
+    class PrintB extends React.Component {
+      static contextType = ContextB;
+      render() {
+        return <Text text={this.context} />;
+      }
+    }
+
+    function Throws() {
+      const value = React.useContext(ContextA);
+      throw new Error(value);
+    }
+
+    const loggedErrors = [];
+    await act(async () => {
+      const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
+        <div>
+          <PrintA />
+          <div>
+            <ContextA.Provider value="A0.1">
+              <Suspense
+                fallback={
+                  <b>
+                    <Text text="Loading..." />
+                  </b>
+                }>
+                <ContextA.Provider value="A0.1.1">
+                  <Throws />
+                </ContextA.Provider>
+              </Suspense>
+              <PrintB />
+            </ContextA.Provider>
+          </div>
+          <PrintA />
+        </div>,
+        writable,
+        {
+          onError(x) {
+            loggedErrors.push(x);
+          },
+        },
+      );
+      startWriting();
+    });
+    expect(loggedErrors.length).toBe(1);
+    expect(loggedErrors[0].message).toEqual('A0.1.1');
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        A0
+        <div>
+          <b>Loading...</b>B0
+        </div>
+        A0
+      </div>,
+    );
+  });
+
   // @gate experimental
   it('client renders a boundary if it errors before finishing the fallback', async () => {
     function App({isClient}) {
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index 9409b57bbbf1..d51ad44ed64e 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -1174,6 +1174,13 @@ function renderNode(request: Request, task: Task, node: ReactNodeList): void {
       // Restore all active ReactContexts to what they were before.
       switchContext(previousContext);
     } else {
+      // Restore the context. We assume that this will be restored by the inner
+      // functions in case nothing throws so we don't use "finally" here.
+      task.blockedSegment.formatContext = previousFormatContext;
+      task.legacyContext = previousLegacyContext;
+      task.context = previousContext;
+      // Restore all active ReactContexts to what they were before.
+      switchContext(previousContext);
       // We assume that we don't need the correct context.
       // Let's terminate the rest of the tree and don't render any siblings.
       throw x;