fix: restore execution context after RetryAfterError completed (#21766)

* test: Add failing test due to executionContext not being restored

* fix: restore execution context after RetryAfterError completed

* Poke codesandbox/ci

* Completely restore executionContext

* expect a specific error

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

@@ -246,6 +246,33 @@ function runActTests(label, render, unmount, rerender) {
         ]);
       });
 
+      // @gate __DEV__
+      it('warns if a setState is called outside of act(...) after a component threw', () => {
+        let setValue = null;
+        function App({defaultValue}) {
+          if (defaultValue === undefined) {
+            throw new Error('some error');
+          }
+          const [value, _setValue] = React.useState(defaultValue);
+          setValue = _setValue;
+          return value;
+        }
+
+        expect(() => {
+          act(() => {
+            render(<App defaultValue={undefined} />, container);
+          });
+        }).toThrow('some error');
+
+        act(() => {
+          rerender(<App defaultValue={0} />, container);
+        });
+
+        expect(() => setValue(1)).toErrorDev([
+          'An update to App inside a test was not wrapped in act(...).',
+        ]);
+      });
+
       describe('fake timers', () => {
         beforeEach(() => {
           jest.useFakeTimers();

packages/react-reconciler/src/ReactFiberWorkLoop.new.js

@@ -779,6 +779,7 @@ function performConcurrentWorkOnRoot(root, didTimeout) {
       : renderRootSync(root, lanes);
   if (exitStatus !== RootIncomplete) {
     if (exitStatus === RootErrored) {
+      const prevExecutionContext = executionContext;
       executionContext |= RetryAfterError;
 
       // If an error occurred during hydration,
@@ -800,6 +801,8 @@ function performConcurrentWorkOnRoot(root, didTimeout) {
         lanes = errorRetryLanes;
         exitStatus = renderRootSync(root, errorRetryLanes);
       }
+
+      executionContext = prevExecutionContext;
     }
 
     if (exitStatus === RootFatalErrored) {
@@ -972,6 +975,7 @@ function performSyncWorkOnRoot(root) {
 
   let exitStatus = renderRootSync(root, lanes);
   if (root.tag !== LegacyRoot && exitStatus === RootErrored) {
+    const prevExecutionContext = executionContext;
     executionContext |= RetryAfterError;
 
     // If an error occurred during hydration,
@@ -993,6 +997,8 @@ function performSyncWorkOnRoot(root) {
       lanes = errorRetryLanes;
       exitStatus = renderRootSync(root, lanes);
     }
+
+    executionContext = prevExecutionContext;
   }
 
   if (exitStatus === RootFatalErrored) {

packages/react-reconciler/src/ReactFiberWorkLoop.old.js

@@ -779,6 +779,7 @@ function performConcurrentWorkOnRoot(root, didTimeout) {
       : renderRootSync(root, lanes);
   if (exitStatus !== RootIncomplete) {
     if (exitStatus === RootErrored) {
+      const prevExecutionContext = executionContext;
       executionContext |= RetryAfterError;
 
       // If an error occurred during hydration,
@@ -800,6 +801,8 @@ function performConcurrentWorkOnRoot(root, didTimeout) {
         lanes = errorRetryLanes;
         exitStatus = renderRootSync(root, errorRetryLanes);
       }
+
+      executionContext = prevExecutionContext;
     }
 
     if (exitStatus === RootFatalErrored) {
@@ -972,6 +975,7 @@ function performSyncWorkOnRoot(root) {
 
   let exitStatus = renderRootSync(root, lanes);
   if (root.tag !== LegacyRoot && exitStatus === RootErrored) {
+    const prevExecutionContext = executionContext;
     executionContext |= RetryAfterError;
 
     // If an error occurred during hydration,
@@ -993,6 +997,8 @@ function performSyncWorkOnRoot(root) {
       lanes = errorRetryLanes;
       exitStatus = renderRootSync(root, lanes);
     }
+
+    executionContext = prevExecutionContext;
   }
 
   if (exitStatus === RootFatalErrored) {