Fix ViewTransition onExit cleanup timing on unmount

packages/react-dom/src/__tests__/ReactDOMViewTransition-test.js

@@ -16,8 +16,17 @@ let SuspenseList;
 let ViewTransition;
 let act;
 let assertLog;
+let assertConsoleWarnDev;
 let Scheduler;
 let textCache;
+let finishMockViewTransition;
+let originalStartViewTransition;
+let originalGetBoundingClientRect;
+let originalCSS;
+let originalGetAnimations;
+let originalInnerWidth;
+let originalInnerHeight;
+let originalFonts;
 
 describe('ReactDOMViewTransition', () => {
   let container;
@@ -29,6 +38,7 @@ describe('ReactDOMViewTransition', () => {
     Scheduler = require('scheduler');
     act = require('internal-test-utils').act;
     assertLog = require('internal-test-utils').assertLog;
+    assertConsoleWarnDev = require('internal-test-utils').assertConsoleWarnDev;
     Suspense = React.Suspense;
     ViewTransition = React.ViewTransition;
     if (gate(flags => flags.enableSuspenseList)) {
@@ -38,9 +48,79 @@ describe('ReactDOMViewTransition', () => {
     document.body.appendChild(container);
 
     textCache = new Map();
+
+    finishMockViewTransition = null;
+    originalStartViewTransition = document.startViewTransition;
+    originalGetBoundingClientRect = HTMLElement.prototype.getBoundingClientRect;
+    originalCSS = global.CSS;
+    originalGetAnimations = document.documentElement.getAnimations;
+    originalInnerWidth = window.innerWidth;
+    originalInnerHeight = window.innerHeight;
+    originalFonts = document.fonts;
+    if (originalCSS == null || typeof originalCSS.escape !== 'function') {
+      global.CSS = {
+        escape: value => value,
+      };
+    }
+    document.documentElement.getAnimations = function () {
+      return [];
+    };
+    document.fonts = {
+      status: 'loaded',
+      ready: Promise.resolve(),
+    };
+    Object.defineProperty(window, 'innerWidth', {
+      configurable: true,
+      value: 1024,
+    });
+    Object.defineProperty(window, 'innerHeight', {
+      configurable: true,
+      value: 768,
+    });
+    // Force visibility checks to pass in jsdom so transition callbacks run.
+    HTMLElement.prototype.getBoundingClientRect = function () {
+      return {
+        top: 1,
+        left: 1,
+        right: 11,
+        bottom: 11,
+        width: 10,
+        height: 10,
+        x: 1,
+        y: 1,
+        toJSON() {},
+      };
+    };
+    document.startViewTransition = jest.fn(({update}) => {
+      const ready = Promise.resolve().then(() => update());
+      return {
+        skipTransition() {},
+        ready,
+        finished: new Promise(resolve => {
+          finishMockViewTransition = resolve;
+        }),
+      };
+    });
   });
 
   afterEach(() => {
+    if (finishMockViewTransition !== null) {
+      finishMockViewTransition();
+      finishMockViewTransition = null;
+    }
+    document.startViewTransition = originalStartViewTransition;
+    HTMLElement.prototype.getBoundingClientRect = originalGetBoundingClientRect;
+    global.CSS = originalCSS;
+    document.documentElement.getAnimations = originalGetAnimations;
+    document.fonts = originalFonts;
+    Object.defineProperty(window, 'innerWidth', {
+      configurable: true,
+      value: originalInnerWidth,
+    });
+    Object.defineProperty(window, 'innerHeight', {
+      configurable: true,
+      value: originalInnerHeight,
+    });
     document.body.removeChild(container);
   });
 
@@ -176,4 +256,50 @@ describe('ReactDOMViewTransition', () => {
     expect(container.textContent).toContain('Card 2');
     expect(container.textContent).toContain('Card 3');
   });
+
+  // @gate enableViewTransition
+  it('calls onExit cleanup immediately on unmount', async () => {
+    const log = [];
+
+    function App({show}) {
+      return show ? (
+        <ViewTransition
+          exit="exit-class"
+          onExit={() => {
+            log.push('exit');
+            return () => {
+              log.push('cleanup');
+            };
+          }}>
+          <div>A</div>
+        </ViewTransition>
+      ) : null;
+    }
+
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      React.startTransition(() => {
+        root.render(<App show={true} />);
+      });
+    });
+    document.startViewTransition.mockClear();
+
+    log.length = 0;
+    await act(async () => {
+      React.startTransition(() => {
+        root.render(<App show={false} />);
+      });
+    });
+    if (__DEV__) {
+      assertConsoleWarnDev([
+        'A flushSync update cancelled a View Transition because it was called ' +
+          'while the View Transition was still preparing. To preserve the synchronous ' +
+          "semantics, React had to skip the View Transition. If you can, try to avoid " +
+          "flushSync() in a scenario that's likely to interfere.",
+      ]);
+    }
+    expect(document.startViewTransition).toHaveBeenCalledTimes(1);
+
+    expect(log).toEqual(['exit', 'cleanup']);
+  });
 });

packages/react-reconciler/src/ReactFiberCommitViewTransitions.js

@@ -446,7 +446,7 @@ export function commitExitViewTransitions(deletion: Fiber): void {
         // Therefore it's possible for onShare to be called with only an old snapshot.
         scheduleViewTransitionEvent(deletion, props.onShare);
       } else {
-        scheduleViewTransitionEvent(deletion, props.onExit);
+        scheduleViewTransitionEvent(deletion, props.onExit, true);
       }
     }
     if (appearingViewTransitions !== null) {

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -735,9 +735,13 @@ let pendingEffectsRenderEndTime: number = -0; // Profiling-only
 let pendingPassiveTransitions: Array<Transition> | null = null;
 let pendingRecoverableErrors: null | Array<CapturedValue<mixed>> = null;
 let pendingViewTransition: null | RunningViewTransition = null;
-let pendingViewTransitionEvents: Array<
-  (types: Array<string>) => void | (() => void),
-> | null = null;
+type PendingViewTransitionEvent = {
+  callback: (types: Array<string>) => void | (() => void),
+  // If true, run the cleanup immediately after unmount.
+  cleanupOnUnmount: boolean,
+};
+let pendingViewTransitionEvents: Array<PendingViewTransitionEvent> | null =
+  null;
 let pendingTransitionTypes: null | TransitionTypes = null;
 let pendingDidIncludeRenderPhaseUpdate: boolean = false;
 let pendingSuspendedCommitReason: SuspendedCommitReason = null; // Profiling-only
@@ -906,6 +910,7 @@ export function scheduleViewTransitionEvent(
     instance: ViewTransitionInstance,
     types: Array<string>,
   ) => void | (() => void),
+  cleanupOnUnmount?: boolean,
 ): void {
   if (enableViewTransition) {
     if (callback != null) {
@@ -919,7 +924,10 @@ export function scheduleViewTransitionEvent(
       if (pendingViewTransitionEvents === null) {
         pendingViewTransitionEvents = [];
       }
-      pendingViewTransitionEvents.push(callback.bind(null, instance));
+      pendingViewTransitionEvents.push({
+        callback: callback.bind(null, instance),
+        cleanupOnUnmount: cleanupOnUnmount === true,
+      });
     }
   }
 }
@@ -952,9 +960,10 @@ export function scheduleGestureTransitionEvent(
         if (pendingViewTransitionEvents === null) {
           pendingViewTransitionEvents = [];
         }
-        pendingViewTransitionEvents.push(
-          callback.bind(null, timeline, options, instance),
-        );
+        pendingViewTransitionEvents.push({
+          callback: callback.bind(null, timeline, options, instance),
+          cleanupOnUnmount: false,
+        });
       }
     }
   }
@@ -4276,10 +4285,18 @@ function flushSpawnedWork(): void {
       }
       if (committedViewTransition !== null) {
         for (let i = 0; i < pendingEvents.length; i++) {
-          const viewTransitionEvent = pendingEvents[i];
+          const {callback: viewTransitionEvent, cleanupOnUnmount} =
+            pendingEvents[i];
           const cleanup = viewTransitionEvent(pendingTypes);
           if (cleanup !== undefined) {
-            addViewTransitionFinishedListener(committedViewTransition, cleanup);
+            if (cleanupOnUnmount) {
+              cleanup();
+            } else {
+              addViewTransitionFinishedListener(
+                committedViewTransition,
+                cleanup,
+              );
+            }
           }
         }
       }
@@ -4554,10 +4571,15 @@ function flushGestureAnimations(): void {
         const runningTransition = appliedGesture.running;
         if (runningTransition !== null) {
           for (let i = 0; i < pendingEvents.length; i++) {
-            const viewTransitionEvent = pendingEvents[i];
+            const {callback: viewTransitionEvent, cleanupOnUnmount} =
+              pendingEvents[i];
             const cleanup = viewTransitionEvent(pendingTypes);
             if (cleanup !== undefined) {
-              addViewTransitionFinishedListener(runningTransition, cleanup);
+              if (cleanupOnUnmount) {
+                cleanup();
+              } else {
+                addViewTransitionFinishedListener(runningTransition, cleanup);
+              }
             }
           }
         }