Execute afterPrepareSurfacesRunnables after committing the transaction

Some of the runnables will assume the transcation has been committed by
the surface flinger before getting run. For example, the leashes of
insets need to be setup before dispatching to the client. If the order
is not ensured, the transaction sent from the client might be executed
earlier than the one sent from system server which might override the
state of the leash unexpectedly.

This CL executes the runnables when the transaction has been committed.

Fix: 298018626
Test: atest InsetsPolicyTest InsetsStateControllerTest
Change-Id: Ie7bb268f228ae8a8ed15e49b193aa2d0183f31ac
diff --git a/services/core/java/com/android/server/wm/InsetsSourceProvider.java b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
index cd96806..fa76774 100644
--- a/services/core/java/com/android/server/wm/InsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
@@ -520,11 +520,37 @@
         updateVisibility();
         mControl = new InsetsSourceControl(mSource.getId(), mSource.getType(), leash,
                 mClientVisible, surfacePosition, getInsetsHint());
+        mStateController.notifySurfaceTransactionReady(this, getSurfaceTransactionId(leash), true);
 
         ProtoLog.d(WM_DEBUG_WINDOW_INSETS,
                 "InsetsSource Control %s for target %s", mControl, mControlTarget);
     }
 
+    private long getSurfaceTransactionId(SurfaceControl leash) {
+        // Here returns mNativeObject (long) as the ID instead of the leash itself so that
+        // InsetsStateController won't keep referencing the leash unexpectedly.
+        return leash != null ? leash.mNativeObject : 0;
+    }
+
+    /**
+     * This is called when the surface transaction of the leash initialization has been committed.
+     *
+     * @param id Indicates which transaction is committed so that stale callbacks can be dropped.
+     */
+    void onSurfaceTransactionCommitted(long id) {
+        if (mIsLeashReadyForDispatching) {
+            return;
+        }
+        if (mControl == null) {
+            return;
+        }
+        if (id != getSurfaceTransactionId(mControl.getLeash())) {
+            return;
+        }
+        mIsLeashReadyForDispatching = true;
+        mStateController.notifySurfaceTransactionReady(this, 0, false);
+    }
+
     void startSeamlessRotation() {
         if (!mSeamlessRotating) {
             mSeamlessRotating = true;
@@ -545,10 +571,6 @@
         return true;
     }
 
-    void onSurfaceTransactionApplied() {
-        mIsLeashReadyForDispatching = true;
-    }
-
     void setClientVisible(boolean clientVisible) {
         if (mClientVisible == clientVisible) {
             return;
@@ -733,6 +755,7 @@
         public void onAnimationCancelled(SurfaceControl animationLeash) {
             if (mAdapter == this) {
                 mStateController.notifyControlRevoked(mControlTarget, InsetsSourceProvider.this);
+                mStateController.notifySurfaceTransactionReady(InsetsSourceProvider.this, 0, false);
                 mControl = null;
                 mControlTarget = null;
                 mAdapter = null;
diff --git a/services/core/java/com/android/server/wm/InsetsStateController.java b/services/core/java/com/android/server/wm/InsetsStateController.java
index 6b9fcf4..ba578f6 100644
--- a/services/core/java/com/android/server/wm/InsetsStateController.java
+++ b/services/core/java/com/android/server/wm/InsetsStateController.java
@@ -34,6 +34,7 @@
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.SparseArray;
+import android.util.SparseLongArray;
 import android.util.proto.ProtoOutputStream;
 import android.view.InsetsSource;
 import android.view.InsetsSourceControl;
@@ -58,6 +59,7 @@
     private final DisplayContent mDisplayContent;
 
     private final SparseArray<InsetsSourceProvider> mProviders = new SparseArray<>();
+    private final SparseLongArray mSurfaceTransactionIds = new SparseLongArray();
     private final ArrayMap<InsetsControlTarget, ArrayList<InsetsSourceProvider>>
             mControlTargetProvidersMap = new ArrayMap<>();
     private final SparseArray<InsetsControlTarget> mIdControlTargetMap = new SparseArray<>();
@@ -360,14 +362,32 @@
         notifyPendingInsetsControlChanged();
     }
 
+    void notifySurfaceTransactionReady(InsetsSourceProvider provider, long id, boolean ready) {
+        if (ready) {
+            mSurfaceTransactionIds.put(provider.getSource().getId(), id);
+        } else {
+            mSurfaceTransactionIds.delete(provider.getSource().getId());
+        }
+    }
+
     private void notifyPendingInsetsControlChanged() {
         if (mPendingControlChanged.isEmpty()) {
             return;
         }
+        final int size = mSurfaceTransactionIds.size();
+        final SparseLongArray surfaceTransactionIds = new SparseLongArray(size);
+        for (int i = 0; i < size; i++) {
+            surfaceTransactionIds.append(
+                    mSurfaceTransactionIds.keyAt(i), mSurfaceTransactionIds.valueAt(i));
+        }
         mDisplayContent.mWmService.mAnimator.addAfterPrepareSurfacesRunnable(() -> {
-            for (int i = mProviders.size() - 1; i >= 0; i--) {
-                final InsetsSourceProvider provider = mProviders.valueAt(i);
-                provider.onSurfaceTransactionApplied();
+            for (int i = 0; i < size; i++) {
+                final int sourceId = surfaceTransactionIds.keyAt(i);
+                final InsetsSourceProvider provider = mProviders.get(sourceId);
+                if (provider == null) {
+                    continue;
+                }
+                provider.onSurfaceTransactionCommitted(surfaceTransactionIds.valueAt(i));
             }
             final ArraySet<InsetsControlTarget> newControlTargets = new ArraySet<>();
             int displayId = mDisplayContent.getDisplayId();
diff --git a/services/core/java/com/android/server/wm/WindowAnimator.java b/services/core/java/com/android/server/wm/WindowAnimator.java
index b43a454..5e7f1cb 100644
--- a/services/core/java/com/android/server/wm/WindowAnimator.java
+++ b/services/core/java/com/android/server/wm/WindowAnimator.java
@@ -28,6 +28,7 @@
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 
 import android.content.Context;
+import android.os.HandlerExecutor;
 import android.os.Trace;
 import android.util.Slog;
 import android.util.TimeUtils;
@@ -69,6 +70,8 @@
 
     private Choreographer mChoreographer;
 
+    private final HandlerExecutor mExecutor;
+
     /**
      * Indicates whether we have an animation frame callback scheduled, which will happen at
      * vsync-app and then schedule the animation tick at the right time (vsync-sf).
@@ -80,8 +83,7 @@
      * A list of runnable that need to be run after {@link WindowContainer#prepareSurfaces} is
      * executed and the corresponding transaction is closed and applied.
      */
-    private final ArrayList<Runnable> mAfterPrepareSurfacesRunnables = new ArrayList<>();
-    private boolean mInExecuteAfterPrepareSurfacesRunnables;
+    private ArrayList<Runnable> mAfterPrepareSurfacesRunnables = new ArrayList<>();
 
     private final SurfaceControl.Transaction mTransaction;
 
@@ -92,6 +94,7 @@
         mTransaction = service.mTransactionFactory.get();
         service.mAnimationHandler.runWithScissors(
                 () -> mChoreographer = Choreographer.getSfInstance(), 0 /* timeout */);
+        mExecutor = new HandlerExecutor(service.mAnimationHandler);
 
         mAnimationFrameCallback = frameTimeNs -> {
             synchronized (mService.mGlobalLock) {
@@ -197,6 +200,19 @@
             updateRunningExpensiveAnimationsLegacy();
         }
 
+        final ArrayList<Runnable> afterPrepareSurfacesRunnables = mAfterPrepareSurfacesRunnables;
+        if (!afterPrepareSurfacesRunnables.isEmpty()) {
+            mAfterPrepareSurfacesRunnables = new ArrayList<>();
+            mTransaction.addTransactionCommittedListener(mExecutor, () -> {
+                synchronized (mService.mGlobalLock) {
+                    // Traverse in order they were added.
+                    for (int i = 0, size = afterPrepareSurfacesRunnables.size(); i < size; i++) {
+                        afterPrepareSurfacesRunnables.get(i).run();
+                    }
+                    afterPrepareSurfacesRunnables.clear();
+                }
+            });
+        }
         Trace.traceBegin(Trace.TRACE_TAG_WINDOW_MANAGER, "applyTransaction");
         mTransaction.apply();
         Trace.traceEnd(Trace.TRACE_TAG_WINDOW_MANAGER);
@@ -204,7 +220,6 @@
         ProtoLog.i(WM_SHOW_TRANSACTIONS, "<<< CLOSE TRANSACTION animate");
 
         mService.mAtmService.mTaskOrganizerController.dispatchPendingEvents();
-        executeAfterPrepareSurfacesRunnables();
 
         if (DEBUG_WINDOW_TRACE) {
             Slog.i(TAG, "!!! animate: exit"
@@ -286,34 +301,10 @@
 
     /**
      * Adds a runnable to be executed after {@link WindowContainer#prepareSurfaces} is called and
-     * the corresponding transaction is closed and applied.
+     * the corresponding transaction is closed, applied, and committed.
      */
     void addAfterPrepareSurfacesRunnable(Runnable r) {
-        // If runnables are already being handled in executeAfterPrepareSurfacesRunnable, then just
-        // immediately execute the runnable passed in.
-        if (mInExecuteAfterPrepareSurfacesRunnables) {
-            r.run();
-            return;
-        }
-
         mAfterPrepareSurfacesRunnables.add(r);
         scheduleAnimation();
     }
-
-    void executeAfterPrepareSurfacesRunnables() {
-
-        // Don't even think about to start recursing!
-        if (mInExecuteAfterPrepareSurfacesRunnables) {
-            return;
-        }
-        mInExecuteAfterPrepareSurfacesRunnables = true;
-
-        // Traverse in order they were added.
-        final int size = mAfterPrepareSurfacesRunnables.size();
-        for (int i = 0; i < size; i++) {
-            mAfterPrepareSurfacesRunnables.get(i).run();
-        }
-        mAfterPrepareSurfacesRunnables.clear();
-        mInExecuteAfterPrepareSurfacesRunnables = false;
-    }
 }
diff --git a/services/tests/servicestests/utils/com/android/server/testutils/StubTransaction.java b/services/tests/servicestests/utils/com/android/server/testutils/StubTransaction.java
index e27bb4c..b9ece93 100644
--- a/services/tests/servicestests/utils/com/android/server/testutils/StubTransaction.java
+++ b/services/tests/servicestests/utils/com/android/server/testutils/StubTransaction.java
@@ -40,12 +40,19 @@
 public class StubTransaction extends SurfaceControl.Transaction {
 
     private HashSet<Runnable> mWindowInfosReportedListeners = new HashSet<>();
+    private HashSet<SurfaceControl.TransactionCommittedListener> mTransactionCommittedListeners =
+            new HashSet<>();
 
     @Override
     public void apply() {
         for (Runnable listener : mWindowInfosReportedListeners) {
             listener.run();
         }
+        for (SurfaceControl.TransactionCommittedListener listener
+                : mTransactionCommittedListeners) {
+            listener.onTransactionCommitted();
+        }
+        mTransactionCommittedListeners.clear();
     }
 
     @Override
@@ -239,6 +246,9 @@
     @Override
     public SurfaceControl.Transaction addTransactionCommittedListener(Executor executor,
             SurfaceControl.TransactionCommittedListener listener) {
+        SurfaceControl.TransactionCommittedListener listenerInner =
+                () -> executor.execute(listener::onTransactionCommitted);
+        mTransactionCommittedListeners.add(listenerInner);
         return this;
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java
index 0c1fbf3..1a1fe95 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java
@@ -94,6 +94,7 @@
     public void setUp() throws Exception {
         assumeFalse(WindowManagerService.sEnableShellTransitions);
         mAppTransitionController = new AppTransitionController(mWm, mDisplayContent);
+        mWm.mAnimator.ready();
     }
 
     @Test
@@ -855,7 +856,7 @@
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
 
         // Animation run by the remote handler.
         assertTrue(remoteAnimationRunner.isAnimationStarted());
@@ -886,7 +887,7 @@
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity,
                 null /* changingTaskFragment */);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
 
         // Animation is not run by the remote handler because the activity is filling the Task.
         assertFalse(remoteAnimationRunner.isAnimationStarted());
@@ -921,7 +922,7 @@
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity,
                 null /* changingTaskFragment */);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
 
         // Animation run by the remote handler.
         assertTrue(remoteAnimationRunner.isAnimationStarted());
@@ -946,7 +947,7 @@
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
 
         // Animation run by the remote handler.
         assertTrue(remoteAnimationRunner.isAnimationStarted());
@@ -973,7 +974,7 @@
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment1);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
 
         // Animation run by the remote handler.
         assertTrue(remoteAnimationRunner.isAnimationStarted());
@@ -997,7 +998,7 @@
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
 
         // Animation not run by the remote handler.
         assertFalse(remoteAnimationRunner.isAnimationStarted());
@@ -1024,7 +1025,7 @@
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
 
         // Animation should not run by the remote handler when there are non-embedded activities of
         // different UID.
@@ -1051,7 +1052,7 @@
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
 
         // Animation should not run by the remote handler when there is wallpaper in the transition.
         assertFalse(remoteAnimationRunner.isAnimationStarted());
@@ -1085,7 +1086,7 @@
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(activity1, null /* closingActivity */, taskFragment);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
 
         // The animation will be animated remotely by client and all activities are input disabled
         // for untrusted animation.
@@ -1136,7 +1137,7 @@
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
 
         // The animation will be animated remotely by client and all activities are input disabled
         // for untrusted animation.
@@ -1178,7 +1179,7 @@
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
 
         // The animation will be animated remotely by client, but input should not be dropped for
         // fully trusted.
diff --git a/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
index 2085d61..1f15ec3 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
@@ -371,6 +371,7 @@
         mDisplayContent.getInsetsPolicy().updateBarControlTarget(app);
         mDisplayContent.getInsetsPolicy().showTransient(statusBars(),
                 true /* isGestureOnSystemBar */);
+        mWm.mAnimator.ready();
         waitUntilWindowAnimatorIdle();
 
         assertTrue(mDisplayContent.getInsetsPolicy().isTransient(statusBars()));
diff --git a/services/tests/wmtests/src/com/android/server/wm/RemoteAnimationControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/RemoteAnimationControllerTest.java
index 11d9629..a163801 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RemoteAnimationControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RemoteAnimationControllerTest.java
@@ -43,7 +43,6 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyFloat;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
@@ -113,6 +112,7 @@
         runWithScissors(mWm.mH, () -> mHandler = new TestHandler(null, mClock), 0);
         mController = new RemoteAnimationController(mWm, mDisplayContent, mAdapter,
                 mHandler, false /*isActivityEmbedding*/);
+        mWm.mAnimator.ready();
     }
 
     private WindowState createAppOverlayWindow() {
@@ -136,7 +136,7 @@
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_ACTIVITY_OPEN);
-            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+            waitUntilWindowAnimatorIdle();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -168,7 +168,7 @@
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_ACTIVITY_OPEN);
-            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+            waitUntilWindowAnimatorIdle();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -290,7 +290,7 @@
         adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                 mFinishedCallback);
         mController.goodToGo(TRANSIT_OLD_ACTIVITY_OPEN);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
         final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                 ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
         final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -336,7 +336,7 @@
         task.applyAnimationUnchecked(null /* lp */, true /* enter */, TRANSIT_OLD_TASK_OPEN,
                 false /* isVoiceInteraction */, null /* sources */);
         mController.goodToGo(TRANSIT_OLD_TASK_OPEN);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
         final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                 ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
         try {
@@ -363,7 +363,7 @@
             ((AnimationAdapter) record.mThumbnailAdapter).startAnimation(mMockThumbnailLeash,
                     mMockTransaction, ANIMATION_TYPE_WINDOW_ANIMATION, mThumbnailFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_TASK_CHANGE_WINDOWING_MODE);
-            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+            waitUntilWindowAnimatorIdle();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -417,7 +417,7 @@
             ((AnimationAdapter) record.mThumbnailAdapter).startAnimation(mMockThumbnailLeash,
                     mMockTransaction, ANIMATION_TYPE_WINDOW_ANIMATION, mThumbnailFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_TASK_CHANGE_WINDOWING_MODE);
-            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+            waitUntilWindowAnimatorIdle();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -471,7 +471,7 @@
             ((AnimationAdapter) record.mThumbnailAdapter).startAnimation(mMockThumbnailLeash,
                     mMockTransaction, ANIMATION_TYPE_WINDOW_ANIMATION, mThumbnailFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_TASK_CHANGE_WINDOWING_MODE);
-            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+            waitUntilWindowAnimatorIdle();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -526,7 +526,7 @@
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_ACTIVITY_OPEN);
-            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+            waitUntilWindowAnimatorIdle();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -559,7 +559,7 @@
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_ACTIVITY_OPEN);
-            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+            waitUntilWindowAnimatorIdle();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -595,7 +595,7 @@
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_KEYGUARD_GOING_AWAY);
-            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+            waitUntilWindowAnimatorIdle();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -645,7 +645,7 @@
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_KEYGUARD_GOING_AWAY_ON_WALLPAPER);
-            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+            waitUntilWindowAnimatorIdle();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -782,7 +782,7 @@
 
             mDisplayContent.applySurfaceChangesTransaction();
             mController.goodToGo(TRANSIT_OLD_TASK_OPEN);
-            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+            waitUntilWindowAnimatorIdle();
 
             verify(mMockRunner).onAnimationStart(eq(TRANSIT_OLD_TASK_OPEN),
                     any(), any(), any(), any());
@@ -810,7 +810,7 @@
         adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                 mFinishedCallback);
         mController.goodToGo(transit);
-        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        waitUntilWindowAnimatorIdle();
         return adapter;
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
index a8f6fe8..7ab093d 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
@@ -546,7 +546,7 @@
         // This makes sure all previous messages in the handler are fully processed vs. just popping
         // them from the message queue.
         final AtomicBoolean currentMessagesProcessed = new AtomicBoolean(false);
-        wm.mAnimator.getChoreographer().postFrameCallback(time -> {
+        wm.mAnimator.addAfterPrepareSurfacesRunnable(() -> {
             synchronized (currentMessagesProcessed) {
                 currentMessagesProcessed.set(true);
                 currentMessagesProcessed.notifyAll();