[Fiber] Trigger default indicator for isomorphic async actions with no root associated (#33190)

Stacked on #33160, #33162, #33186 and #33188.

We have a special case that's awkward for default indicators. When you
start a new async Transition from `React.startTransition` then there's
not yet any associated root with the Transition because you haven't
necessarily `setState` on anything yet until the promise resolves.
That's what `entangleAsyncAction` handles by creating a lane that
everything entangles with until all async actions are done.

If there are no sync updates before the end of the event, we should
trigger a default indicator until either the async action completes
without update or if it gets entangled with some roots we should keep it
going until those roots are done.

fixtures/view-transition/src/components/Page.js

@@ -113,8 +113,8 @@ export default function Page({url, navigate}) {
     <button
       onClick={() =>
         startTransition(async () => {
-          setShowModal(true);
           await sleep(2000);
+          setShowModal(true);
         })
       }>
       Show Modal

packages/react-reconciler/src/ReactFiberAsyncAction.js

@@ -15,7 +15,10 @@ import type {
 import type {Lane} from './ReactFiberLane';
 import type {Transition} from 'react/src/ReactStartTransition';
 
-import {requestTransitionLane} from './ReactFiberRootScheduler';
+import {
+  requestTransitionLane,
+  ensureScheduleIsScheduled,
+} from './ReactFiberRootScheduler';
 import {NoLane} from './ReactFiberLane';
 import {
   hasScheduledTransitionWork,
@@ -24,9 +27,13 @@ import {
 import {
   enableComponentPerformanceTrack,
   enableProfilerTimer,
+  enableDefaultTransitionIndicator,
 } from 'shared/ReactFeatureFlags';
 import {clearEntangledAsyncTransitionTypes} from './ReactFiberTransitionTypes';
 
+import noop from 'shared/noop';
+import reportGlobalError from 'shared/reportGlobalError';
+
 // If there are multiple, concurrent async actions, they are entangled. All
 // transition updates that occur while the async action is still in progress
 // are treated as part of the action.
@@ -46,6 +53,21 @@ let currentEntangledLane: Lane = NoLane;
 // until the async action scope has completed.
 let currentEntangledActionThenable: Thenable<void> | null = null;
 
+// Track the default indicator for every root. undefined means we haven't
+// had any roots registered yet. null means there's more than one callback.
+// If there's more than one callback we bailout to not supporting isomorphic
+// default indicators.
+let isomorphicDefaultTransitionIndicator:
+  | void
+  | null
+  | (() => void | (() => void)) = undefined;
+// The clean up function for the currently running indicator.
+let pendingIsomorphicIndicator: null | (() => void) = null;
+// The number of roots that have pending Transitions that depend on the
+// started isomorphic indicator.
+let pendingEntangledRoots: number = 0;
+let needsIsomorphicIndicator: boolean = false;
+
 export function entangleAsyncAction<S>(
   transition: Transition,
   thenable: Thenable<S>,
@@ -66,6 +88,12 @@ export function entangleAsyncAction<S>(
       },
     };
     currentEntangledActionThenable = entangledThenable;
+    if (enableDefaultTransitionIndicator) {
+      needsIsomorphicIndicator = true;
+      // We'll check if we need a default indicator in a microtask. Ensure
+      // we have this scheduled even if no root is scheduled.
+      ensureScheduleIsScheduled();
+    }
   }
   currentEntangledPendingCount++;
   thenable.then(pingEngtangledActionScope, pingEngtangledActionScope);
@@ -86,6 +114,9 @@ function pingEngtangledActionScope() {
       }
     }
     clearEntangledAsyncTransitionTypes();
+    if (pendingEntangledRoots === 0) {
+      stopIsomorphicDefaultIndicator();
+    }
     if (currentEntangledListeners !== null) {
       // All the actions have finished. Close the entangled async action scope
       // and notify all the listeners.
@@ -98,6 +129,7 @@ function pingEngtangledActionScope() {
       currentEntangledListeners = null;
       currentEntangledLane = NoLane;
       currentEntangledActionThenable = null;
+      needsIsomorphicIndicator = false;
       for (let i = 0; i < listeners.length; i++) {
         const listener = listeners[i];
         listener();
@@ -161,3 +193,71 @@ export function peekEntangledActionLane(): Lane {
 export function peekEntangledActionThenable(): Thenable<void> | null {
   return currentEntangledActionThenable;
 }
+
+export function registerDefaultIndicator(
+  onDefaultTransitionIndicator: () => void | (() => void),
+): void {
+  if (!enableDefaultTransitionIndicator) {
+    return;
+  }
+  if (isomorphicDefaultTransitionIndicator === undefined) {
+    isomorphicDefaultTransitionIndicator = onDefaultTransitionIndicator;
+  } else if (
+    isomorphicDefaultTransitionIndicator !== onDefaultTransitionIndicator
+  ) {
+    isomorphicDefaultTransitionIndicator = null;
+    // Stop any on-going indicator since it's now ambiguous.
+    stopIsomorphicDefaultIndicator();
+  }
+}
+
+export function startIsomorphicDefaultIndicatorIfNeeded() {
+  if (!enableDefaultTransitionIndicator) {
+    return;
+  }
+  if (!needsIsomorphicIndicator) {
+    return;
+  }
+  if (
+    isomorphicDefaultTransitionIndicator != null &&
+    pendingIsomorphicIndicator === null
+  ) {
+    try {
+      pendingIsomorphicIndicator =
+        isomorphicDefaultTransitionIndicator() || noop;
+    } catch (x) {
+      pendingIsomorphicIndicator = noop;
+      reportGlobalError(x);
+    }
+  }
+}
+
+function stopIsomorphicDefaultIndicator() {
+  if (!enableDefaultTransitionIndicator) {
+    return;
+  }
+  if (pendingIsomorphicIndicator !== null) {
+    const cleanup = pendingIsomorphicIndicator;
+    pendingIsomorphicIndicator = null;
+    cleanup();
+  }
+}
+
+function releaseIsomorphicIndicator() {
+  if (--pendingEntangledRoots === 0) {
+    stopIsomorphicDefaultIndicator();
+  }
+}
+
+export function hasOngoingIsomorphicIndicator(): boolean {
+  return pendingIsomorphicIndicator !== null;
+}
+
+export function retainIsomorphicIndicator(): () => void {
+  pendingEntangledRoots++;
+  return releaseIsomorphicIndicator;
+}
+
+export function markIsomorphicIndicatorHandled(): void {
+  needsIsomorphicIndicator = false;
+}

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -125,6 +125,7 @@ export {
   defaultOnRecoverableError,
 } from './ReactFiberErrorLogger';
 import {getLabelForLane, TotalLanes} from 'react-reconciler/src/ReactFiberLane';
+import {registerDefaultIndicator} from './ReactFiberAsyncAction';
 
 type OpaqueRoot = FiberRoot;
 
@@ -259,7 +260,7 @@ export function createContainer(
 ): OpaqueRoot {
   const hydrate = false;
   const initialChildren = null;
-  return createFiberRoot(
+  const root = createFiberRoot(
     containerInfo,
     tag,
     hydrate,
@@ -274,6 +275,8 @@ export function createContainer(
     onDefaultTransitionIndicator,
     transitionCallbacks,
   );
+  registerDefaultIndicator(onDefaultTransitionIndicator);
+  return root;
 }
 
 export function createHydrationContainer(
@@ -323,6 +326,8 @@ export function createHydrationContainer(
     transitionCallbacks,
   );
 
+  registerDefaultIndicator(onDefaultTransitionIndicator);
+
   // TODO: Move this to FiberRoot constructor
   root.context = getContextForSubtree(null);
 

packages/react-reconciler/src/ReactFiberRootScheduler.js

@@ -85,6 +85,13 @@ import {peekEntangledActionLane} from './ReactFiberAsyncAction';
 import noop from 'shared/noop';
 import reportGlobalError from 'shared/reportGlobalError';
 
+import {
+  startIsomorphicDefaultIndicatorIfNeeded,
+  hasOngoingIsomorphicIndicator,
+  retainIsomorphicIndicator,
+  markIsomorphicIndicatorHandled,
+} from './ReactFiberAsyncAction';
+
 // A linked list of all the roots with pending work. In an idiomatic app,
 // there's only a single root, but we do support multi root apps, hence this
 // extra complexity. But this module is optimized for the single root case.
@@ -130,6 +137,20 @@ export function ensureRootIsScheduled(root: FiberRoot): void {
   // without consulting the schedule.
   mightHavePendingSyncWork = true;
 
+  ensureScheduleIsScheduled();
+
+  if (
+    __DEV__ &&
+    !disableLegacyMode &&
+    ReactSharedInternals.isBatchingLegacy &&
+    root.tag === LegacyRoot
+  ) {
+    // Special `act` case: Record whenever a legacy update is scheduled.
+    ReactSharedInternals.didScheduleLegacyUpdate = true;
+  }
+}
+
+export function ensureScheduleIsScheduled(): void {
   // At the end of the current event, go through each of the roots and ensure
   // there's a task scheduled for each one at the correct priority.
   if (__DEV__ && ReactSharedInternals.actQueue !== null) {
@@ -144,16 +165,6 @@ export function ensureRootIsScheduled(root: FiberRoot): void {
       scheduleImmediateRootScheduleTask();
     }
   }
-
-  if (
-    __DEV__ &&
-    !disableLegacyMode &&
-    ReactSharedInternals.isBatchingLegacy &&
-    root.tag === LegacyRoot
-  ) {
-    // Special `act` case: Record whenever a legacy update is scheduled.
-    ReactSharedInternals.didScheduleLegacyUpdate = true;
-  }
 }
 
 export function flushSyncWorkOnAllRoots() {
@@ -339,18 +350,30 @@ function startDefaultTransitionIndicatorIfNeeded() {
   if (!enableDefaultTransitionIndicator) {
     return;
   }
+  // Check if we need to start an isomorphic indicator like if an async action
+  // was started.
+  startIsomorphicDefaultIndicatorIfNeeded();
   // Check all the roots if there are any new indicators needed.
   let root = firstScheduledRoot;
   while (root !== null) {
     if (root.indicatorLanes !== NoLanes && root.pendingIndicator === null) {
       // We have new indicator lanes that requires a loading state. Start the
       // default transition indicator.
-      try {
-        const onDefaultTransitionIndicator = root.onDefaultTransitionIndicator;
-        root.pendingIndicator = onDefaultTransitionIndicator() || noop;
-      } catch (x) {
-        root.pendingIndicator = noop;
-        reportGlobalError(x);
+      if (hasOngoingIsomorphicIndicator()) {
+        // We already have an isomorphic indicator going which means it has to
+        // also apply to this root since it implies all roots have the same one.
+        // We retain this indicator so that it keeps going until we commit this
+        // root.
+        root.pendingIndicator = retainIsomorphicIndicator();
+      } else {
+        try {
+          const onDefaultTransitionIndicator =
+            root.onDefaultTransitionIndicator;
+          root.pendingIndicator = onDefaultTransitionIndicator() || noop;
+        } catch (x) {
+          root.pendingIndicator = noop;
+          reportGlobalError(x);
+        }
       }
     }
     root = root.next;
@@ -708,5 +731,6 @@ export function markIndicatorHandled(root: FiberRoot): void {
     // Clear it from the indicator lanes. We don't need to show a separate
     // loading state for this lane.
     root.indicatorLanes &= ~currentEventTransitionLane;
+    markIsomorphicIndicatorHandled();
   }
 }

packages/react-reconciler/src/__tests__/ReactDefaultTransitionIndicator-test.js

@@ -265,7 +265,6 @@ describe('ReactDefaultTransitionIndicator', () => {
 
     await act(() => {
       // Start an async action but we haven't called setState yet
-      // TODO: This should ideally work with React.startTransition too but we don't know the root.
       start(() => promise);
     });
 
@@ -280,6 +279,132 @@ describe('ReactDefaultTransitionIndicator', () => {
     expect(root).toMatchRenderedOutput('Hi');
   });
 
+  // @gate enableDefaultTransitionIndicator
+  it('triggers the default indicator while an async transition is ongoing (isomorphic)', async () => {
+    let resolve;
+    const promise = new Promise(r => (resolve = r));
+    function App() {
+      return 'Hi';
+    }
+
+    const root = ReactNoop.createRoot({
+      onDefaultTransitionIndicator() {
+        Scheduler.log('start');
+        return () => {
+          Scheduler.log('stop');
+        };
+      },
+    });
+    await act(() => {
+      root.render(<App />);
+    });
+
+    assertLog([]);
+
+    await act(() => {
+      // Start an async action but we haven't called setState yet
+      React.startTransition(() => promise);
+    });
+
+    assertLog(['start']);
+
+    await act(async () => {
+      await resolve('Hello');
+    });
+
+    assertLog(['stop']);
+
+    expect(root).toMatchRenderedOutput('Hi');
+  });
+
+  it('does not triggers isomorphic async action default indicator if there are two different ones', async () => {
+    let resolve;
+    const promise = new Promise(r => (resolve = r));
+    function App() {
+      return 'Hi';
+    }
+
+    const root = ReactNoop.createRoot({
+      onDefaultTransitionIndicator() {
+        Scheduler.log('start');
+        return () => {
+          Scheduler.log('stop');
+        };
+      },
+    });
+    // Initialize second root. This is now ambiguous which indicator to use.
+    ReactNoop.createRoot({
+      onDefaultTransitionIndicator() {
+        Scheduler.log('start2');
+        return () => {
+          Scheduler.log('stop2');
+        };
+      },
+    });
+    await act(() => {
+      root.render(<App />);
+    });
+
+    assertLog([]);
+
+    await act(() => {
+      // Start an async action but we haven't called setState yet
+      React.startTransition(() => promise);
+    });
+
+    assertLog([]);
+
+    await act(async () => {
+      await resolve('Hello');
+    });
+
+    assertLog([]);
+
+    expect(root).toMatchRenderedOutput('Hi');
+  });
+
+  it('does not triggers isomorphic async action default indicator if there is a loading state', async () => {
+    let resolve;
+    const promise = new Promise(r => (resolve = r));
+    let update;
+    function App() {
+      const [state, setState] = useState(false);
+      update = setState;
+      return state ? 'Loading' : 'Hi';
+    }
+
+    const root = ReactNoop.createRoot({
+      onDefaultTransitionIndicator() {
+        Scheduler.log('start');
+        return () => {
+          Scheduler.log('stop');
+        };
+      },
+    });
+    await act(() => {
+      root.render(<App />);
+    });
+
+    assertLog([]);
+
+    await act(() => {
+      update(true);
+      React.startTransition(() => promise.then(() => update(false)));
+    });
+
+    assertLog([]);
+
+    expect(root).toMatchRenderedOutput('Loading');
+
+    await act(async () => {
+      await resolve('Hello');
+    });
+
+    assertLog([]);
+
+    expect(root).toMatchRenderedOutput('Hi');
+  });
+
   it('should not trigger for useDeferredValue (sync)', async () => {
     function Text({text}) {
       Scheduler.log(text);