When one or more async transitions are on going add all transition types to every root that gets scheduled

Author: sebmarkbage

packages/react-reconciler/src/ReactFiberAsyncAction.js

@@ -25,6 +25,7 @@ import {
   enableComponentPerformanceTrack,
   enableProfilerTimer,
 } from 'shared/ReactFeatureFlags';
+import {clearEntangledAsyncTransitionTypes} from './ReactFiberTransitionTypes';
 
 // If there are multiple, concurrent async actions, they are entangled. All
 // transition updates that occur while the async action is still in progress
@@ -84,6 +85,7 @@ function pingEngtangledActionScope() {
         clearAsyncTransitionTimer();
       }
     }
+    clearEntangledAsyncTransitionTypes();
     if (currentEntangledListeners !== null) {
       // All the actions have finished. Close the entangled async action scope
       // and notify all the listeners.

packages/react-reconciler/src/ReactFiberTransition.js

@@ -12,7 +12,7 @@ import type {
   GestureProvider,
   GestureOptions,
 } from 'shared/ReactTypes';
-import type {Lanes} from './ReactFiberLane';
+import {NoLane, type Lanes} from './ReactFiberLane';
 import type {StackCursor} from './ReactFiberStack';
 import type {Cache, SpawnedCachePool} from './ReactFiberCacheComponent';
 import type {Transition} from 'react/src/ReactStartTransition';
@@ -34,10 +34,17 @@ import {
   retainCache,
   CacheContext,
 } from './ReactFiberCacheComponent';
-import {queueTransitionTypes} from './ReactFiberTransitionTypes';
+import {
+  queueTransitionTypes,
+  entangleAsyncTransitionTypes,
+  entangledTransitionTypes,
+} from './ReactFiberTransitionTypes';
 
 import ReactSharedInternals from 'shared/ReactSharedInternals';
-import {entangleAsyncAction} from './ReactFiberAsyncAction';
+import {
+  entangleAsyncAction,
+  peekEntangledActionLane,
+} from './ReactFiberAsyncAction';
 import {startAsyncTransitionTimer} from './ReactProfilerTimer';
 import {firstScheduledRoot} from './ReactFiberRootScheduler';
 import {
@@ -88,15 +95,31 @@ ReactSharedInternals.S = function onStartTransitionFinishForReconciler(
     const thenable: Thenable<mixed> = (returnValue: any);
     entangleAsyncAction(transition, thenable);
   }
-  if (enableViewTransition && transition.types !== null) {
-    // Within this Transition we should've now scheduled any roots we have updates
-    // to work on. If there are no updates on a root, then the Transition type won't
-    // be applied to that root.
-    // TODO: The exception is if we're to an async action, the updates might come in later.
-    let root = firstScheduledRoot;
-    while (root !== null) {
-      queueTransitionTypes(root, transition.types);
-      root = root.next;
+  if (enableViewTransition) {
+    if (entangledTransitionTypes !== null) {
+      // If we scheduled work on any new roots, we need to add any entangled async
+      // transition types to those roots too.
+      let root = firstScheduledRoot;
+      while (root !== null) {
+        queueTransitionTypes(root, entangledTransitionTypes);
+        root = root.next;
+      }
+    }
+    const transitionTypes = transition.types;
+    if (transitionTypes !== null) {
+      // Within this Transition we should've now scheduled any roots we have updates
+      // to work on. If there are no updates on a root, then the Transition type won't
+      // be applied to that root.
+      let root = firstScheduledRoot;
+      while (root !== null) {
+        queueTransitionTypes(root, transitionTypes);
+        root = root.next;
+      }
+      if (peekEntangledActionLane() !== NoLane) {
+        // If we have entangled, async actions going on, the update associated with
+        // these types might come later. We need to save them for later.
+        entangleAsyncTransitionTypes(transitionTypes);
+      }
     }
   }
   if (prevOnStartTransitionFinish !== null) {

packages/react-reconciler/src/ReactFiberTransitionTypes.js

@@ -14,9 +14,9 @@ import {enableViewTransition} from 'shared/ReactFeatureFlags';
 
 export function queueTransitionTypes(
   root: FiberRoot,
-  transitionTypes: null | TransitionTypes,
+  transitionTypes: TransitionTypes,
 ): void {
-  if (enableViewTransition && transitionTypes !== null) {
+  if (enableViewTransition) {
     // TODO: We should really store transitionTypes per lane in a LaneMap on
     // the root. Then merge it when we commit. We currently assume that all
     // Transitions are entangled.
@@ -33,6 +33,31 @@ export function queueTransitionTypes(
   }
 }
 
+// Store all types while we're entangled with an async Transition.
+export let entangledTransitionTypes: null | TransitionTypes = null;
+
+export function entangleAsyncTransitionTypes(
+  transitionTypes: TransitionTypes,
+): void {
+  if (enableViewTransition) {
+    let queued = entangledTransitionTypes;
+    if (queued === null) {
+      queued = entangledTransitionTypes = [];
+    }
+    for (let i = 0; i < transitionTypes.length; i++) {
+      const transitionType = transitionTypes[i];
+      if (queued.indexOf(transitionType) === -1) {
+        queued.push(transitionType);
+      }
+    }
+  }
+}
+
+export function clearEntangledAsyncTransitionTypes() {
+  // Called when all Async Actions are done.
+  entangledTransitionTypes = null;
+}
+
 export function claimQueuedTransitionTypes(
   root: FiberRoot,
 ): null | TransitionTypes {