[Fizz] Restore context after an error happens

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}) {

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;