Fix context propagation for offscreen/fallback trees (#23095)

* Failing test for Context.Consumer in suspended Suspense

See issue #19701.

* Fix context propagation for offscreen trees

* Address nits

* Specify propagation root for Suspense too

* Pass correct propagation root

* Harden test coverage

This test will fail if we remove propagation, or if we propagate with a root node like fiber.return or fiber.return.return. The additional DEV-only error helps detect a different kind of mistake, like if the thing being passed hasn't actually been encountered on the way up. However, we still leave the actual production loop to check against null so that there is no way we loop forever if the propagation root is wrong.

* Remove superfluous warning

Co-authored-by: overlookmotel <[email protected]>

Author: gaearon

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

@@ -174,7 +174,7 @@ import {
   checkIfContextChanged,
   readContext,
   prepareToReadContext,
-  scheduleWorkOnParentPath,
+  scheduleContextWorkOnParentPath,
 } from './ReactFiberNewContext.new';
 import {
   renderWithHooks,
@@ -2754,13 +2754,17 @@ function updateDehydratedSuspenseComponent(
   }
 }
 
-function scheduleWorkOnFiber(fiber: Fiber, renderLanes: Lanes) {
+function scheduleSuspenseWorkOnFiber(
+  fiber: Fiber,
+  renderLanes: Lanes,
+  propagationRoot: Fiber,
+) {
   fiber.lanes = mergeLanes(fiber.lanes, renderLanes);
   const alternate = fiber.alternate;
   if (alternate !== null) {
     alternate.lanes = mergeLanes(alternate.lanes, renderLanes);
   }
-  scheduleWorkOnParentPath(fiber.return, renderLanes);
+  scheduleContextWorkOnParentPath(fiber.return, renderLanes, propagationRoot);
 }
 
 function propagateSuspenseContextChange(
@@ -2776,15 +2780,15 @@ function propagateSuspenseContextChange(
     if (node.tag === SuspenseComponent) {
       const state: SuspenseState | null = node.memoizedState;
       if (state !== null) {
-        scheduleWorkOnFiber(node, renderLanes);
+        scheduleSuspenseWorkOnFiber(node, renderLanes, workInProgress);
       }
     } else if (node.tag === SuspenseListComponent) {
       // If the tail is hidden there might not be an Suspense boundaries
       // to schedule work on. In this case we have to schedule it on the
       // list itself.
       // We don't have to traverse to the children of the list since
       // the list will propagate the change when it rerenders.
-      scheduleWorkOnFiber(node, renderLanes);
+      scheduleSuspenseWorkOnFiber(node, renderLanes, workInProgress);
     } else if (node.child !== null) {
       node.child.return = node;
       node = node.child;

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

@@ -174,7 +174,7 @@ import {
   checkIfContextChanged,
   readContext,
   prepareToReadContext,
-  scheduleWorkOnParentPath,
+  scheduleContextWorkOnParentPath,
 } from './ReactFiberNewContext.old';
 import {
   renderWithHooks,
@@ -2754,13 +2754,17 @@ function updateDehydratedSuspenseComponent(
   }
 }
 
-function scheduleWorkOnFiber(fiber: Fiber, renderLanes: Lanes) {
+function scheduleSuspenseWorkOnFiber(
+  fiber: Fiber,
+  renderLanes: Lanes,
+  propagationRoot: Fiber,
+) {
   fiber.lanes = mergeLanes(fiber.lanes, renderLanes);
   const alternate = fiber.alternate;
   if (alternate !== null) {
     alternate.lanes = mergeLanes(alternate.lanes, renderLanes);
   }
-  scheduleWorkOnParentPath(fiber.return, renderLanes);
+  scheduleContextWorkOnParentPath(fiber.return, renderLanes, propagationRoot);
 }
 
 function propagateSuspenseContextChange(
@@ -2776,15 +2780,15 @@ function propagateSuspenseContextChange(
     if (node.tag === SuspenseComponent) {
       const state: SuspenseState | null = node.memoizedState;
       if (state !== null) {
-        scheduleWorkOnFiber(node, renderLanes);
+        scheduleSuspenseWorkOnFiber(node, renderLanes, workInProgress);
       }
     } else if (node.tag === SuspenseListComponent) {
       // If the tail is hidden there might not be an Suspense boundaries
       // to schedule work on. In this case we have to schedule it on the
       // list itself.
       // We don't have to traverse to the children of the list since
       // the list will propagate the change when it rerenders.
-      scheduleWorkOnFiber(node, renderLanes);
+      scheduleSuspenseWorkOnFiber(node, renderLanes, workInProgress);
     } else if (node.child !== null) {
       node.child.return = node;
       node = node.child;

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

@@ -138,9 +138,10 @@ export function popProvider(
   }
 }
 
-export function scheduleWorkOnParentPath(
+export function scheduleContextWorkOnParentPath(
   parent: Fiber | null,
   renderLanes: Lanes,
+  propagationRoot: Fiber,
 ) {
   // Update the child lanes of all the ancestors, including the alternates.
   let node = parent;
@@ -157,12 +158,26 @@ export function scheduleWorkOnParentPath(
     ) {
       alternate.childLanes = mergeLanes(alternate.childLanes, renderLanes);
     } else {
-      // Neither alternate was updated, which means the rest of the
+      // Neither alternate was updated.
+      // Normally, this would mean that the rest of the
       // ancestor path already has sufficient priority.
+      // However, this is not necessarily true inside offscreen
+      // or fallback trees because childLanes may be inconsistent
+      // with the surroundings. This is why we continue the loop.
+    }
+    if (node === propagationRoot) {
       break;
     }
     node = node.return;
   }
+  if (__DEV__) {
+    if (node !== propagationRoot) {
+      console.error(
+        'Expected to find the propagation root when scheduling context work. ' +
+          'This error is likely caused by a bug in React. Please file an issue.',
+      );
+    }
+  }
 }
 
 export function propagateContextChange<T>(
@@ -246,7 +261,11 @@ function propagateContextChange_eager<T>(
           if (alternate !== null) {
             alternate.lanes = mergeLanes(alternate.lanes, renderLanes);
           }
-          scheduleWorkOnParentPath(fiber.return, renderLanes);
+          scheduleContextWorkOnParentPath(
+            fiber.return,
+            renderLanes,
+            workInProgress,
+          );
 
           // Mark the updated lanes on the list, too.
           list.lanes = mergeLanes(list.lanes, renderLanes);
@@ -284,7 +303,11 @@ function propagateContextChange_eager<T>(
       // because we want to schedule this fiber as having work
       // on its children. We'll use the childLanes on
       // this fiber to indicate that a context has changed.
-      scheduleWorkOnParentPath(parentSuspense, renderLanes);
+      scheduleContextWorkOnParentPath(
+        parentSuspense,
+        renderLanes,
+        workInProgress,
+      );
       nextFiber = fiber.sibling;
     } else {
       // Traverse down.
@@ -365,7 +388,11 @@ function propagateContextChanges<T>(
             if (alternate !== null) {
               alternate.lanes = mergeLanes(alternate.lanes, renderLanes);
             }
-            scheduleWorkOnParentPath(consumer.return, renderLanes);
+            scheduleContextWorkOnParentPath(
+              consumer.return,
+              renderLanes,
+              workInProgress,
+            );
 
             if (!forcePropagateEntireTree) {
               // During lazy propagation, when we find a match, we can defer
@@ -406,7 +433,11 @@ function propagateContextChanges<T>(
       // because we want to schedule this fiber as having work
       // on its children. We'll use the childLanes on
       // this fiber to indicate that a context has changed.
-      scheduleWorkOnParentPath(parentSuspense, renderLanes);
+      scheduleContextWorkOnParentPath(
+        parentSuspense,
+        renderLanes,
+        workInProgress,
+      );
       nextFiber = null;
     } else {
       // Traverse down.

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

@@ -138,9 +138,10 @@ export function popProvider(
   }
 }
 
-export function scheduleWorkOnParentPath(
+export function scheduleContextWorkOnParentPath(
   parent: Fiber | null,
   renderLanes: Lanes,
+  propagationRoot: Fiber,
 ) {
   // Update the child lanes of all the ancestors, including the alternates.
   let node = parent;
@@ -157,12 +158,26 @@ export function scheduleWorkOnParentPath(
     ) {
       alternate.childLanes = mergeLanes(alternate.childLanes, renderLanes);
     } else {
-      // Neither alternate was updated, which means the rest of the
+      // Neither alternate was updated.
+      // Normally, this would mean that the rest of the
       // ancestor path already has sufficient priority.
+      // However, this is not necessarily true inside offscreen
+      // or fallback trees because childLanes may be inconsistent
+      // with the surroundings. This is why we continue the loop.
+    }
+    if (node === propagationRoot) {
       break;
     }
     node = node.return;
   }
+  if (__DEV__) {
+    if (node !== propagationRoot) {
+      console.error(
+        'Expected to find the propagation root when scheduling context work. ' +
+          'This error is likely caused by a bug in React. Please file an issue.',
+      );
+    }
+  }
 }
 
 export function propagateContextChange<T>(
@@ -246,7 +261,11 @@ function propagateContextChange_eager<T>(
           if (alternate !== null) {
             alternate.lanes = mergeLanes(alternate.lanes, renderLanes);
           }
-          scheduleWorkOnParentPath(fiber.return, renderLanes);
+          scheduleContextWorkOnParentPath(
+            fiber.return,
+            renderLanes,
+            workInProgress,
+          );
 
           // Mark the updated lanes on the list, too.
           list.lanes = mergeLanes(list.lanes, renderLanes);
@@ -284,7 +303,11 @@ function propagateContextChange_eager<T>(
       // because we want to schedule this fiber as having work
       // on its children. We'll use the childLanes on
       // this fiber to indicate that a context has changed.
-      scheduleWorkOnParentPath(parentSuspense, renderLanes);
+      scheduleContextWorkOnParentPath(
+        parentSuspense,
+        renderLanes,
+        workInProgress,
+      );
       nextFiber = fiber.sibling;
     } else {
       // Traverse down.
@@ -365,7 +388,11 @@ function propagateContextChanges<T>(
             if (alternate !== null) {
               alternate.lanes = mergeLanes(alternate.lanes, renderLanes);
             }
-            scheduleWorkOnParentPath(consumer.return, renderLanes);
+            scheduleContextWorkOnParentPath(
+              consumer.return,
+              renderLanes,
+              workInProgress,
+            );
 
             if (!forcePropagateEntireTree) {
               // During lazy propagation, when we find a match, we can defer
@@ -406,7 +433,11 @@ function propagateContextChanges<T>(
       // because we want to schedule this fiber as having work
       // on its children. We'll use the childLanes on
       // this fiber to indicate that a context has changed.
-      scheduleWorkOnParentPath(parentSuspense, renderLanes);
+      scheduleContextWorkOnParentPath(
+        parentSuspense,
+        renderLanes,
+        workInProgress,
+      );
       nextFiber = null;
     } else {
       // Traverse down.

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

@@ -1532,5 +1532,68 @@ describe('ReactSuspense', () => {
       expect(Scheduler).toFlushUntilNextPaint(['new value']);
       expect(root).toMatchRenderedOutput('new value');
     });
+
+    it('updates context consumer within child of suspended suspense component when context updates', () => {
+      const {createContext, useState} = React;
+
+      const ValueContext = createContext(null);
+
+      const promiseThatNeverResolves = new Promise(() => {});
+      function Child() {
+        return (
+          <ValueContext.Consumer>
+            {value => {
+              Scheduler.unstable_yieldValue(
+                `Received context value [${value}]`,
+              );
+              if (value === 'default') return <Text text="default" />;
+              throw promiseThatNeverResolves;
+            }}
+          </ValueContext.Consumer>
+        );
+      }
+
+      let setValue;
+      function Wrapper({children}) {
+        const [value, _setValue] = useState('default');
+        setValue = _setValue;
+        return (
+          <ValueContext.Provider value={value}>
+            {children}
+          </ValueContext.Provider>
+        );
+      }
+
+      function App() {
+        return (
+          <Wrapper>
+            <Suspense fallback={<Text text="Loading..." />}>
+              <Child />
+            </Suspense>
+          </Wrapper>
+        );
+      }
+
+      const root = ReactTestRenderer.create(<App />);
+      expect(Scheduler).toHaveYielded([
+        'Received context value [default]',
+        'default',
+      ]);
+      expect(root).toMatchRenderedOutput('default');
+
+      act(() => setValue('new value'));
+      expect(Scheduler).toHaveYielded([
+        'Received context value [new value]',
+        'Loading...',
+      ]);
+      expect(root).toMatchRenderedOutput('Loading...');
+
+      act(() => setValue('default'));
+      expect(Scheduler).toHaveYielded([
+        'Received context value [default]',
+        'default',
+      ]);
+      expect(root).toMatchRenderedOutput('default');
+    });
   });
 });