Forbid readContext() during render-phase class setState() reducer form

packages/react-reconciler/src/ReactFiberHooks.js

@@ -16,8 +16,8 @@ import {NoWork} from './ReactFiberExpirationTime';
 import {enableHooks} from 'shared/ReactFeatureFlags';
 import {
   readContext,
-  stashContextDependencies,
-  unstashContextDependencies,
+  stashContextDependenciesInDEV,
+  unstashContextDependenciesInDEV,
 } from './ReactFiberNewContext';
 import {
   Update as UpdateEffect,
@@ -604,10 +604,14 @@ export function useReducer<S, A>(
             const action = update.action;
             // Temporarily clear to forbid calling Hooks in a reducer.
             currentlyRenderingFiber = null;
-            stashContextDependencies();
+            if (__DEV__) {
+              stashContextDependenciesInDEV();
+            }
             newState = reducer(newState, action);
             currentlyRenderingFiber = fiber;
-            unstashContextDependencies();
+            if (__DEV__) {
+              unstashContextDependenciesInDEV();
+            }
             update = update.next;
           } while (update !== null);
 
@@ -678,10 +682,14 @@ export function useReducer<S, A>(
             const action = update.action;
             // Temporarily clear to forbid calling Hooks in a reducer.
             currentlyRenderingFiber = null;
-            stashContextDependencies();
+            if (__DEV__) {
+              stashContextDependenciesInDEV();
+            }
             newState = reducer(newState, action);
             currentlyRenderingFiber = fiber;
-            unstashContextDependencies();
+            if (__DEV__) {
+              unstashContextDependenciesInDEV();
+            }
           }
         }
         prevUpdate = update;
@@ -712,7 +720,9 @@ export function useReducer<S, A>(
   }
   // Temporarily clear to forbid calling Hooks in a reducer.
   currentlyRenderingFiber = null;
-  stashContextDependencies();
+  if (__DEV__) {
+    stashContextDependenciesInDEV();
+  }
   // There's no existing queue, so this is the initial render.
   if (reducer === basicStateReducer) {
     // Special case for `useState`.
@@ -723,7 +733,9 @@ export function useReducer<S, A>(
     initialState = reducer(initialState, initialAction);
   }
   currentlyRenderingFiber = fiber;
-  unstashContextDependencies();
+  if (__DEV__) {
+    unstashContextDependenciesInDEV();
+  }
   workInProgressHook.memoizedState = workInProgressHook.baseState = initialState;
   queue = workInProgressHook.queue = {
     last: null,
@@ -957,10 +969,14 @@ export function useMemo<T>(
 
   // Temporarily clear to forbid calling Hooks.
   currentlyRenderingFiber = null;
-  stashContextDependencies();
+  if (__DEV__) {
+    stashContextDependenciesInDEV();
+  }
   const nextValue = nextCreate();
   currentlyRenderingFiber = fiber;
-  unstashContextDependencies();
+  if (__DEV__) {
+    unstashContextDependenciesInDEV();
+  }
   workInProgressHook.memoizedState = [nextValue, nextDeps];
   currentHookNameInDev = null;
   return nextValue;
@@ -1059,10 +1075,14 @@ function dispatchAction<S, A>(
           // Temporarily clear to forbid calling Hooks in a reducer.
           let maybeFiber = currentlyRenderingFiber; // Note: likely null now unlike `fiber`
           currentlyRenderingFiber = null;
-          stashContextDependencies();
+          if (__DEV__) {
+            stashContextDependenciesInDEV();
+          }
           const eagerState = eagerReducer(currentState, action);
           currentlyRenderingFiber = maybeFiber;
-          unstashContextDependencies();
+          if (__DEV__) {
+            unstashContextDependenciesInDEV();
+          }
           // Stash the eagerly computed state, and the reducer used to compute
           // it, on the update object. If the reducer hasn't changed by the
           // time we enter the render phase, then the eager state can be used

packages/react-reconciler/src/ReactFiberNewContext.js

@@ -52,37 +52,30 @@ let currentlyRenderingFiber: Fiber | null = null;
 let lastContextDependency: ContextDependency<mixed> | null = null;
 let lastContextWithAllBitsObserved: ReactContext<any> | null = null;
 
-// We stash the variables above before entering user code in Hooks.
-let stashedCurrentlyRenderingFiber: Fiber | null = null;
-let stashedLastContextDependency: ContextDependency<mixed> | null = null;
-let stashedLastContextWithAllBitsObserved: ReactContext<any> | null = null;
+// Used in DEV to track whether reading context should warn.
+let areContextDependenciesStashedInDEV: boolean = false;
 
 export function resetContextDependences(): void {
   // This is called right before React yields execution, to ensure `readContext`
   // cannot be called outside the render phase.
   currentlyRenderingFiber = null;
   lastContextDependency = null;
   lastContextWithAllBitsObserved = null;
-
-  stashedCurrentlyRenderingFiber = null;
-  stashedLastContextDependency = null;
-  stashedLastContextWithAllBitsObserved = null;
+  if (__DEV__) {
+    areContextDependenciesStashedInDEV = false;
+  }
 }
 
-export function stashContextDependencies(): void {
-  stashedCurrentlyRenderingFiber = currentlyRenderingFiber;
-  stashedLastContextDependency = lastContextDependency;
-  stashedLastContextWithAllBitsObserved = lastContextWithAllBitsObserved;
-
-  currentlyRenderingFiber = null;
-  lastContextDependency = null;
-  lastContextWithAllBitsObserved = null;
+export function stashContextDependenciesInDEV(): void {
+  if (__DEV__) {
+    areContextDependenciesStashedInDEV = true;
+  }
 }
 
-export function unstashContextDependencies(): void {
-  currentlyRenderingFiber = stashedCurrentlyRenderingFiber;
-  lastContextDependency = stashedLastContextDependency;
-  lastContextWithAllBitsObserved = stashedLastContextWithAllBitsObserved;
+export function unstashContextDependenciesInDEV(): void {
+  if (__DEV__) {
+    areContextDependenciesStashedInDEV = false;
+  }
 }
 
 export function pushProvider<T>(providerFiber: Fiber, nextValue: T): void {
@@ -304,6 +297,18 @@ export function readContext<T>(
   context: ReactContext<T>,
   observedBits: void | number | boolean,
 ): T {
+  if (__DEV__) {
+    // This warning would fire if you read context inside a Hook like useMemo.
+    // Unlike the class check below, it's not enforced in production for perf.
+    warning(
+      !areContextDependenciesStashedInDEV,
+      'Context can only be read while React is rendering. ' +
+        'In classes, you can read it in the render method or getDerivedStateFromProps. ' +
+        'In function components, you can read it directly in the function body, but not ' +
+        'inside Hooks like useReducer() or useMemo().',
+    );
+  }
+
   if (lastContextWithAllBitsObserved === context) {
     // Nothing to do. We already observe everything in this context.
   } else if (observedBits === false || observedBits === 0) {

packages/react-reconciler/src/ReactUpdateQueue.js

@@ -88,6 +88,10 @@ import type {Fiber} from './ReactFiber';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 
 import {NoWork} from './ReactFiberExpirationTime';
+import {
+  stashContextDependenciesInDEV,
+  unstashContextDependenciesInDEV,
+} from './ReactFiberNewContext';
 import {Callback, ShouldCapture, DidCapture} from 'shared/ReactSideEffectTags';
 import {ClassComponent} from 'shared/ReactWorkTags';
 
@@ -348,6 +352,7 @@ function getStateFromUpdate<State>(
       if (typeof payload === 'function') {
         // Updater function
         if (__DEV__) {
+          stashContextDependenciesInDEV();
           if (
             debugRenderPhaseSideEffects ||
             (debugRenderPhaseSideEffectsForStrictMode &&
@@ -356,7 +361,11 @@ function getStateFromUpdate<State>(
             payload.call(instance, prevState, nextProps);
           }
         }
-        return payload.call(instance, prevState, nextProps);
+        const nextState = payload.call(instance, prevState, nextProps);
+        if (__DEV__) {
+          unstashContextDependenciesInDEV();
+        }
+        return nextState;
       }
       // State object
       return payload;
@@ -372,6 +381,7 @@ function getStateFromUpdate<State>(
       if (typeof payload === 'function') {
         // Updater function
         if (__DEV__) {
+          stashContextDependenciesInDEV();
           if (
             debugRenderPhaseSideEffects ||
             (debugRenderPhaseSideEffectsForStrictMode &&
@@ -381,6 +391,9 @@ function getStateFromUpdate<State>(
           }
         }
         partialState = payload.call(instance, prevState, nextProps);
+        if (__DEV__) {
+          unstashContextDependenciesInDEV();
+        }
       } else {
         // Partial state object
         partialState = payload;

packages/react-reconciler/src/__tests__/ReactHooks-test.internal.js

@@ -672,7 +672,7 @@ describe('ReactHooks', () => {
     expect(root.toJSON()).toEqual('123');
   });
 
-  it('throws when reading context inside useMemo', () => {
+  it('warns when reading context inside useMemo', () => {
     const {useMemo, createContext} = React;
     const ReactCurrentDispatcher =
       React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
@@ -685,12 +685,12 @@ describe('ReactHooks', () => {
       }, []);
     }
 
-    expect(() => ReactTestRenderer.create(<App />)).toThrow(
+    expect(() => ReactTestRenderer.create(<App />)).toWarnDev(
       'Context can only be read while React is rendering',
     );
   });
 
-  it('throws when reading context inside useMemo after reading outside it', () => {
+  it('warns when reading context inside useMemo after reading outside it', () => {
     const {useMemo, createContext} = React;
     const ReactCurrentDispatcher =
       React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
@@ -707,13 +707,14 @@ describe('ReactHooks', () => {
       }, []);
     }
 
-    expect(() => ReactTestRenderer.create(<App />)).toThrow(
+    expect(() => ReactTestRenderer.create(<App />)).toWarnDev(
       'Context can only be read while React is rendering',
     );
     expect(firstRead).toBe('light');
     expect(secondRead).toBe('light');
   });
 
+  // Throws because there's no runtime cost for being strict here.
   it('throws when reading context inside useEffect', () => {
     const {useEffect, createContext} = React;
     const ReactCurrentDispatcher =
@@ -735,6 +736,7 @@ describe('ReactHooks', () => {
     );
   });
 
+  // Throws because there's no runtime cost for being strict here.
   it('throws when reading context inside useLayoutEffect', () => {
     const {useLayoutEffect, createContext} = React;
     const ReactCurrentDispatcher =
@@ -755,7 +757,7 @@ describe('ReactHooks', () => {
     );
   });
 
-  it('throws when reading context inside useReducer', () => {
+  it('warns when reading context inside useReducer', () => {
     const {useReducer, createContext} = React;
     const ReactCurrentDispatcher =
       React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
@@ -773,13 +775,13 @@ describe('ReactHooks', () => {
       return null;
     }
 
-    expect(() => ReactTestRenderer.create(<App />)).toThrow(
+    expect(() => ReactTestRenderer.create(<App />)).toWarnDev(
       'Context can only be read while React is rendering',
     );
   });
 
   // Edge case.
-  it('throws when reading context inside eager useReducer', () => {
+  it('warns when reading context inside eager useReducer', () => {
     const {useState, createContext} = React;
     const ThemeContext = createContext('light');
 
@@ -810,7 +812,7 @@ describe('ReactHooks', () => {
           <Cls />
         </React.Fragment>,
       ),
-    ).toThrow('Context can only be read while React is rendering');
+    ).toWarnDev('Context can only be read while React is rendering');
   });
 
   it('throws when calling hooks inside useReducer', () => {

packages/react-reconciler/src/__tests__/ReactNewContext-test.internal.js

@@ -1347,6 +1347,31 @@ describe('ReactNewContext', () => {
       expect(ReactNoop.flush()).toEqual(['App', 'App#renderConsumer']);
       expect(ReactNoop.getChildren()).toEqual([span('goodbye')]);
     });
+
+    it('warns when reading context inside render phase class setState updater', () => {
+      const ThemeContext = React.createContext('light');
+
+      class Cls extends React.Component {
+        state = {};
+        render() {
+          this.setState(() => {
+            readContext(ThemeContext);
+          });
+          return null;
+        }
+      }
+
+      ReactNoop.render(<Cls />);
+      expect(ReactNoop.flush).toWarnDev(
+        [
+          'Context can only be read while React is rendering',
+          'Cannot update during an existing state transition',
+        ],
+        {
+          withoutStack: 1,
+        },
+      );
+    });
   });
 
   describe('useContext', () => {