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.

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;