Forbid calling readContext() inside useMemo() for SSR

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

@@ -717,4 +717,36 @@ describe('ReactDOMServerHooks', () => {
       expect(domNode.textContent).toEqual('undefined');
     });
   });
+
+  describe('readContext', () => {
+    function readContext(Context, observedBits) {
+      const dispatcher =
+        React.__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
+          .ReactCurrentDispatcher.current;
+      return dispatcher.readContext(Context, observedBits);
+    }
+
+    itRenders('with a warning inside useMemo and useReducer', async render => {
+      const Context = React.createContext(42);
+
+      function ReadInMemo(props) {
+        let count = React.useMemo(() => readContext(Context), []);
+        return <Text text={count} />;
+      }
+
+      function ReadInReducer(props) {
+        let [count, dispatch] = React.useReducer(() => readContext(Context));
+        if (count !== 42) {
+          dispatch();
+        }
+        return <Text text={count} />;
+      }
+
+      const domNode1 = await render(<ReadInMemo />, 1);
+      expect(domNode1.textContent).toEqual('42');
+
+      const domNode2 = await render(<ReadInReducer />, 1);
+      expect(domNode2.textContent).toEqual('42');
+    });
+  });
 });

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

@@ -49,6 +49,8 @@ let renderPhaseUpdates: Map<UpdateQueue<any>, Update<any>> | null = null;
 let numberOfReRenders: number = 0;
 const RE_RENDER_LIMIT = 25;
 
+let shouldWarnAboutReadingContextInDEV = false;
+
 // In DEV, this is the name of the currently executing primitive hook
 let currentHookNameInDev: ?string;
 
@@ -137,6 +139,9 @@ function createWorkInProgressHook(): Hook {
 
 export function prepareToUseHooks(componentIdentity: Object): void {
   currentlyRenderingComponent = componentIdentity;
+  if (__DEV__) {
+    shouldWarnAboutReadingContextInDEV = false;
+  }
 
   // The following should have already been reset
   // didScheduleRenderPhaseUpdate = false;
@@ -173,6 +178,9 @@ export function finishHooks(
   numberOfReRenders = 0;
   renderPhaseUpdates = null;
   workInProgressHook = null;
+  if (__DEV__) {
+    shouldWarnAboutReadingContextInDEV = false;
+  }
 
   // These were reset above
   // currentlyRenderingComponent = null;
@@ -191,6 +199,15 @@ function readContext<T>(
 ): T {
   let threadID = currentThreadID;
   validateContextBounds(context, threadID);
+  if (__DEV__) {
+    warning(
+      !shouldWarnAboutReadingContextInDEV,
+      '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().',
+    );
+  }
   return context[threadID];
 }
 
@@ -255,7 +272,13 @@ export function useReducer<S, A>(
           const action = update.action;
           // Temporarily clear to forbid calling Hooks.
           currentlyRenderingComponent = null;
+          if (__DEV__) {
+            shouldWarnAboutReadingContextInDEV = true;
+          }
           newState = reducer(newState, action);
+          if (__DEV__) {
+            shouldWarnAboutReadingContextInDEV = false;
+          }
           currentlyRenderingComponent = component;
           update = update.next;
         } while (update !== null);
@@ -311,8 +334,14 @@ function useMemo<T>(nextCreate: () => T, deps: Array<mixed> | void | null): T {
 
   // Temporarily clear to forbid calling Hooks.
   currentlyRenderingComponent = null;
+  if (__DEV__) {
+    shouldWarnAboutReadingContextInDEV = true;
+  }
   const nextValue = nextCreate();
   currentlyRenderingComponent = component;
+  if (__DEV__) {
+    shouldWarnAboutReadingContextInDEV = false;
+  }
   workInProgressHook.memoizedState = [nextValue, nextDeps];
   return nextValue;
 }

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/__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', () => {