Fix state leaking when a function component throws on server render

packages/react-dom/src/__tests__/ReactDOMServerIntegrationHooks-test.js

@@ -867,6 +867,44 @@ describe('ReactDOMServerHooks', () => {
     });
   });
 
+  it('renders successfully after a component using hooks throws an error', () => {
+    function ThrowingComponent() {
+      const [value, dispatch] = useReducer((state, action) => {
+        return state + 1;
+      }, 0);
+
+      // throw an error if the count gets too high during the re-render phase
+      if (value >= 3) {
+        throw new Error('Error from ThrowingComponent');
+      } else {
+        // dispatch to trigger a re-render of the component
+        dispatch();
+      }
+
+      return <div>{value}</div>;
+    }
+
+    function NonThrowingComponent() {
+      const [count] = useState(0);
+      return <div>{count}</div>;
+    }
+
+    // First, render a component that will throw an error during a re-render triggered
+    // by a dispatch call.
+    expect(() => ReactDOMServer.renderToString(<ThrowingComponent />)).toThrow(
+      'Error from ThrowingComponent',
+    );
+
+    // Next, assert that we can render a function component using hooks immediately
+    // after an error occurred, which indictates the internal hooks state has been
+    // reset.
+    const container = document.createElement('div');
+    container.innerHTML = ReactDOMServer.renderToString(
+      <NonThrowingComponent />,
+    );
+    expect(container.children[0].textContent).toEqual('0');
+  });
+
   if (__EXPERIMENTAL__) {
     describe('useOpaqueIdentifier', () => {
       it('generates unique ids for server string render', async () => {

packages/react-dom/src/server/ReactPartialRenderer.js

@@ -60,6 +60,7 @@ import escapeTextForBrowser from './escapeTextForBrowser';
 import {
   prepareToUseHooks,
   finishHooks,
+  resetHooksState,
   Dispatcher,
   currentPartialRenderer,
   setCurrentPartialRenderer,
@@ -955,6 +956,7 @@ class ReactDOMServerRenderer {
     } finally {
       ReactCurrentDispatcher.current = prevDispatcher;
       setCurrentPartialRenderer(prevPartialRenderer);
+      resetHooksState();
     }
   }
 

packages/react-dom/src/server/ReactPartialRendererHooks.js

@@ -202,24 +202,21 @@ export function finishHooks(
 
     children = Component(props, refOrContext);
   }
-  currentlyRenderingComponent = null;
-  firstWorkInProgressHook = null;
-  numberOfReRenders = 0;
-  renderPhaseUpdates = null;
-  workInProgressHook = null;
+  resetHooksState();
+  return children;
+}
+
+// Reset the internal hooks state if an error occurs while rendering a component
+export function resetHooksState(): void {
   if (__DEV__) {
     isInHookUserCodeInDev = false;
   }
 
-  // These were reset above
-  // currentlyRenderingComponent = null;
-  // didScheduleRenderPhaseUpdate = false;
-  // firstWorkInProgressHook = null;
-  // numberOfReRenders = 0;
-  // renderPhaseUpdates = null;
-  // workInProgressHook = null;
-
-  return children;
+  didScheduleRenderPhaseUpdate = false;
+  firstWorkInProgressHook = null;
+  numberOfReRenders = 0;
+  renderPhaseUpdates = null;
+  workInProgressHook = null;
 }
 
 function readContext<T>(