Schedule work to update Context.Consumer inside Suspense

Author: okmttdhr

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

@@ -20,9 +20,12 @@ import type {SharedQueue} from './ReactUpdateQueue.new';
 import {isPrimaryRenderer} from './ReactFiberHostConfig';
 import {createCursor, push, pop} from './ReactFiberStack.new';
 import {
+  ContextConsumer,
   ContextProvider,
   ClassComponent,
   DehydratedFragment,
+  OffscreenComponent,
+  SuspenseComponent,
 } from './ReactWorkTags';
 import {
   NoLanes,
@@ -157,6 +160,10 @@ export function scheduleWorkOnParentPath(
       !isSubsetOfLanes(alternate.childLanes, renderLanes)
     ) {
       alternate.childLanes = mergeLanes(alternate.childLanes, renderLanes);
+    } else {
+      // Neither alternate was updated, which means the rest of the
+      // ancestor path already has sufficient priority.
+      break;
     }
     node = node.return;
   }
@@ -361,6 +368,69 @@ function propagateContextChanges<T>(
             }
             scheduleWorkOnParentPath(consumer.return, renderLanes);
 
+            // If Context.Consumer is a descendant of Suspense,
+            // `scheduleWorkOnParentPath` may break before it get to Suspense.
+            // Schedule work to avoid bailout before Context.Consumer receives the value.
+            const primaryChildFragment = workInProgress.child;
+            if (
+              workInProgress.tag === SuspenseComponent &&
+              primaryChildFragment !== null &&
+              primaryChildFragment.tag === OffscreenComponent &&
+              consumer.tag === ContextConsumer
+            ) {
+              const workInProgressAlternate = workInProgress.alternate;
+              if (!isSubsetOfLanes(workInProgress.childLanes, renderLanes)) {
+                workInProgress.childLanes = mergeLanes(
+                  workInProgress.childLanes,
+                  renderLanes,
+                );
+                if (workInProgressAlternate !== null) {
+                  workInProgressAlternate.childLanes = mergeLanes(
+                    workInProgressAlternate.childLanes,
+                    renderLanes,
+                  );
+                }
+              } else if (
+                workInProgressAlternate !== null &&
+                !isSubsetOfLanes(
+                  workInProgressAlternate.childLanes,
+                  renderLanes,
+                )
+              ) {
+                workInProgressAlternate.childLanes = mergeLanes(
+                  workInProgressAlternate.childLanes,
+                  renderLanes,
+                );
+              }
+              const primaryChildFragmentAlternate =
+                primaryChildFragment.alternate;
+              if (
+                !isSubsetOfLanes(primaryChildFragment.childLanes, renderLanes)
+              ) {
+                primaryChildFragment.childLanes = mergeLanes(
+                  primaryChildFragment.childLanes,
+                  renderLanes,
+                );
+                if (primaryChildFragmentAlternate !== null) {
+                  primaryChildFragmentAlternate.childLanes = mergeLanes(
+                    primaryChildFragmentAlternate.childLanes,
+                    renderLanes,
+                  );
+                }
+              } else if (
+                primaryChildFragmentAlternate !== null &&
+                !isSubsetOfLanes(
+                  primaryChildFragmentAlternate.childLanes,
+                  renderLanes,
+                )
+              ) {
+                primaryChildFragmentAlternate.childLanes = mergeLanes(
+                  primaryChildFragmentAlternate.childLanes,
+                  renderLanes,
+                );
+              }
+            }
+
             if (!forcePropagateEntireTree) {
               // During lazy propagation, when we find a match, we can defer
               // propagating changes to the children, because we're going to

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

@@ -20,9 +20,12 @@ import type {SharedQueue} from './ReactUpdateQueue.old';
 import {isPrimaryRenderer} from './ReactFiberHostConfig';
 import {createCursor, push, pop} from './ReactFiberStack.old';
 import {
+  ContextConsumer,
   ContextProvider,
   ClassComponent,
   DehydratedFragment,
+  OffscreenComponent,
+  SuspenseComponent,
 } from './ReactWorkTags';
 import {
   NoLanes,
@@ -157,6 +160,10 @@ export function scheduleWorkOnParentPath(
       !isSubsetOfLanes(alternate.childLanes, renderLanes)
     ) {
       alternate.childLanes = mergeLanes(alternate.childLanes, renderLanes);
+    } else {
+      // Neither alternate was updated, which means the rest of the
+      // ancestor path already has sufficient priority.
+      break;
     }
     node = node.return;
   }
@@ -361,6 +368,69 @@ function propagateContextChanges<T>(
             }
             scheduleWorkOnParentPath(consumer.return, renderLanes);
 
+            // If Context.Consumer is a descendant of Suspense,
+            // `scheduleWorkOnParentPath` may break before it get to Suspense.
+            // Schedule work to avoid bailout before Context.Consumer receives the value.
+            const primaryChildFragment = workInProgress.child;
+            if (
+              workInProgress.tag === SuspenseComponent &&
+              primaryChildFragment !== null &&
+              primaryChildFragment.tag === OffscreenComponent &&
+              consumer.tag === ContextConsumer
+            ) {
+              const workInProgressAlternate = workInProgress.alternate;
+              if (!isSubsetOfLanes(workInProgress.childLanes, renderLanes)) {
+                workInProgress.childLanes = mergeLanes(
+                  workInProgress.childLanes,
+                  renderLanes,
+                );
+                if (workInProgressAlternate !== null) {
+                  workInProgressAlternate.childLanes = mergeLanes(
+                    workInProgressAlternate.childLanes,
+                    renderLanes,
+                  );
+                }
+              } else if (
+                workInProgressAlternate !== null &&
+                !isSubsetOfLanes(
+                  workInProgressAlternate.childLanes,
+                  renderLanes,
+                )
+              ) {
+                workInProgressAlternate.childLanes = mergeLanes(
+                  workInProgressAlternate.childLanes,
+                  renderLanes,
+                );
+              }
+              const primaryChildFragmentAlternate =
+                primaryChildFragment.alternate;
+              if (
+                !isSubsetOfLanes(primaryChildFragment.childLanes, renderLanes)
+              ) {
+                primaryChildFragment.childLanes = mergeLanes(
+                  primaryChildFragment.childLanes,
+                  renderLanes,
+                );
+                if (primaryChildFragmentAlternate !== null) {
+                  primaryChildFragmentAlternate.childLanes = mergeLanes(
+                    primaryChildFragmentAlternate.childLanes,
+                    renderLanes,
+                  );
+                }
+              } else if (
+                primaryChildFragmentAlternate !== null &&
+                !isSubsetOfLanes(
+                  primaryChildFragmentAlternate.childLanes,
+                  renderLanes,
+                )
+              ) {
+                primaryChildFragmentAlternate.childLanes = mergeLanes(
+                  primaryChildFragmentAlternate.childLanes,
+                  renderLanes,
+                );
+              }
+            }
+
             if (!forcePropagateEntireTree) {
               // During lazy propagation, when we find a match, we can defer
               // propagating changes to the children, because we're going to