Fix state leaking when a function component throws on server render (facebook#19212)

pmaccart · trueadm · commit dcb6032dc5fb · 2020-07-08T13:48:01.000+01:00
* add unit test asserting internal hooks state is reset

* Reset internal hooks state before rendering

* reset hooks state on error

* Use expect...toThrow instead of try/catch in test

* reset dev-only hooks state inside resetHooksState

* reset currentlyRenderingComponent to null
diff --git a/packages/react-dom/src/__tests__/ReactDOMServerIntegrationHooks-test.js b/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 () => {
diff --git a/packages/react-dom/src/server/ReactPartialRenderer.js b/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();
     }
   }
 
diff --git a/packages/react-dom/src/server/ReactPartialRendererHooks.js b/packages/react-dom/src/server/ReactPartialRendererHooks.js
@@ -202,24 +202,22 @@ export function finishHooks(
 
     children = Component(props, refOrContext);
   }
+  resetHooksState();
+  return children;
+}
+
+// Reset the internal hooks state if an error occurs while rendering a component
+export function resetHooksState(): void {
+  if (__DEV__) {
+    isInHookUserCodeInDev = false;
+  }
+
   currentlyRenderingComponent = null;
+  didScheduleRenderPhaseUpdate = false;
   firstWorkInProgressHook = null;
   numberOfReRenders = 0;
   renderPhaseUpdates = null;
   workInProgressHook = null;
-  if (__DEV__) {
-    isInHookUserCodeInDev = false;
-  }
-
-  // These were reset above
-  // currentlyRenderingComponent = null;
-  // didScheduleRenderPhaseUpdate = false;
-  // firstWorkInProgressHook = null;
-  // numberOfReRenders = 0;
-  // renderPhaseUpdates = null;
-  // workInProgressHook = null;
-
-  return children;
 }
 
 function readContext<T>(

Original file line number	Diff line number	Diff line change
`@@ -60,6 +60,7 @@ import escapeTextForBrowser from './escapeTextForBrowser';`
`60`	`60`	`import {`
`61`	`61`	`prepareToUseHooks,`
`62`	`62`	`finishHooks,`
	`63`	`+ resetHooksState,`
`63`	`64`	`Dispatcher,`
`64`	`65`	`currentPartialRenderer,`
`65`	`66`	`setCurrentPartialRenderer,`
`@@ -955,6 +956,7 @@ class ReactDOMServerRenderer {`
`955`	`956`	`} finally {`
`956`	`957`	`ReactCurrentDispatcher.current = prevDispatcher;`
`957`	`958`	`setCurrentPartialRenderer(prevPartialRenderer);`
	`959`	`+ resetHooksState();`
`958`	`960`	`}`
`959`	`961`	`}`
`960`	`962`