Revert "Re-land: Execute afterPrepareSurfacesRunnables after committing..."

This reverts commit 267c6c21b7edbd8be27c9946a3fc1a2c2f36fba6.

Reason for revert: Caused NullPointerException in DisplayImeController

Change-Id: I6bbb959ff4a5d270bc92780d857783cc813cb988

10 files changed, 64 insertions, 112 deletions

diff --git a/services/core/java/com/android/server/wm/InsetsSourceProvider.java b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
index 9faf1068f27d..b66b8bc2115b 100644
--- a/services/core/java/com/android/server/wm/InsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
@@ -556,37 +556,11 @@ class InsetsSourceProvider {
         }
         mControl = new InsetsSourceControl(mSource.getId(), mSource.getType(), leash,
                 initiallyVisible, 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;
@@ -608,6 +582,10 @@ class InsetsSourceProvider {
         return true;
     }
 
+    void onSurfaceTransactionApplied() {
+        mIsLeashReadyForDispatching = true;
+    }
+
     void setClientVisible(boolean clientVisible) {
         if (mClientVisible == clientVisible) {
             return;
@@ -821,7 +799,6 @@ class InsetsSourceProvider {
         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 7294e3f12560..098a691e0490 100644
--- a/services/core/java/com/android/server/wm/InsetsStateController.java
+++ b/services/core/java/com/android/server/wm/InsetsStateController.java
@@ -34,7 +34,6 @@ import android.os.Trace;
 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;
@@ -60,7 +59,6 @@ class InsetsStateController {
     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<>();
@@ -378,32 +376,14 @@ class InsetsStateController {
         }
     }
 
-    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 = 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));
+            for (int i = mProviders.size() - 1; i >= 0; i--) {
+                final InsetsSourceProvider provider = mProviders.valueAt(i);
+                provider.onSurfaceTransactionApplied();
             }
             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 2342de3676de..13334a5f29b1 100644
--- a/services/core/java/com/android/server/wm/WindowAnimator.java
+++ b/services/core/java/com/android/server/wm/WindowAnimator.java
@@ -27,7 +27,6 @@ import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME
 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,8 +68,6 @@ public class WindowAnimator {
 
     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).
@@ -82,7 +79,8 @@ public class WindowAnimator {
      * A list of runnable that need to be run after {@link WindowContainer#prepareSurfaces} is
      * executed and the corresponding transaction is closed and applied.
      */
-    private ArrayList<Runnable> mAfterPrepareSurfacesRunnables = new ArrayList<>();
+    private final ArrayList<Runnable> mAfterPrepareSurfacesRunnables = new ArrayList<>();
+    private boolean mInExecuteAfterPrepareSurfacesRunnables;
 
     private final SurfaceControl.Transaction mTransaction;
 
@@ -93,7 +91,6 @@ public class WindowAnimator {
         mTransaction = service.mTransactionFactory.get();
         service.mAnimationHandler.runWithScissors(
                 () -> mChoreographer = Choreographer.getSfInstance(), 0 /* timeout */);
-        mExecutor = new HandlerExecutor(service.mAnimationHandler);
 
         mAnimationFrameCallback = frameTimeNs -> {
             synchronized (mService.mGlobalLock) {
@@ -201,19 +198,6 @@ public class WindowAnimator {
             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);
@@ -221,6 +205,7 @@ public class WindowAnimator {
         ProtoLog.i(WM_SHOW_TRANSACTIONS, "<<< CLOSE TRANSACTION animate");
 
         mService.mAtmService.mTaskOrganizerController.dispatchPendingEvents();
+        executeAfterPrepareSurfacesRunnables();
 
         if (DEBUG_WINDOW_TRACE) {
             Slog.i(TAG, "!!! animate: exit"
@@ -302,10 +287,34 @@ public class WindowAnimator {
 
     /**
      * Adds a runnable to be executed after {@link WindowContainer#prepareSurfaces} is called and
-     * the corresponding transaction is closed, applied, and committed.
+     * the corresponding transaction is closed and applied.
      */
     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 b9ece9360980..e27bb4c8c3b6 100644
--- a/services/tests/servicestests/utils/com/android/server/testutils/StubTransaction.java
+++ b/services/tests/servicestests/utils/com/android/server/testutils/StubTransaction.java
@@ -40,19 +40,12 @@ import java.util.concurrent.Executor;
 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
@@ -246,9 +239,6 @@ public class StubTransaction extends SurfaceControl.Transaction {
     @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 c294bc62c7ac..0c1fbf3cb3d7 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java
@@ -94,7 +94,6 @@ public class AppTransitionControllerTest extends WindowTestsBase {
     public void setUp() throws Exception {
         assumeFalse(WindowManagerService.sEnableShellTransitions);
         mAppTransitionController = new AppTransitionController(mWm, mDisplayContent);
-        mWm.mAnimator.ready();
     }
 
     @Test
@@ -856,7 +855,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
 
         // Animation run by the remote handler.
         assertTrue(remoteAnimationRunner.isAnimationStarted());
@@ -887,7 +886,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity,
                 null /* changingTaskFragment */);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
 
         // Animation is not run by the remote handler because the activity is filling the Task.
         assertFalse(remoteAnimationRunner.isAnimationStarted());
@@ -922,7 +921,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity,
                 null /* changingTaskFragment */);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
 
         // Animation run by the remote handler.
         assertTrue(remoteAnimationRunner.isAnimationStarted());
@@ -947,7 +946,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
 
         // Animation run by the remote handler.
         assertTrue(remoteAnimationRunner.isAnimationStarted());
@@ -974,7 +973,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment1);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
 
         // Animation run by the remote handler.
         assertTrue(remoteAnimationRunner.isAnimationStarted());
@@ -998,7 +997,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
 
         // Animation not run by the remote handler.
         assertFalse(remoteAnimationRunner.isAnimationStarted());
@@ -1025,7 +1024,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
 
         // Animation should not run by the remote handler when there are non-embedded activities of
         // different UID.
@@ -1052,7 +1051,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
 
         // Animation should not run by the remote handler when there is wallpaper in the transition.
         assertFalse(remoteAnimationRunner.isAnimationStarted());
@@ -1086,7 +1085,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(activity1, null /* closingActivity */, taskFragment);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
 
         // The animation will be animated remotely by client and all activities are input disabled
         // for untrusted animation.
@@ -1137,7 +1136,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
 
         // The animation will be animated remotely by client and all activities are input disabled
         // for untrusted animation.
@@ -1179,7 +1178,7 @@ public class AppTransitionControllerTest extends WindowTestsBase {
 
         // Prepare and start transition.
         prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
 
         // The animation will be animated remotely by client, but input should not be dropped for
         // fully trusted.
@@ -1303,4 +1302,4 @@ public class AppTransitionControllerTest extends WindowTestsBase {
         doReturn(mock(RemoteAnimationTarget.class)).when(activity).createRemoteAnimationTarget(
                 any());
     }
-}
+}
\ No newline at end of file
diff --git a/services/tests/wmtests/src/com/android/server/wm/ImeInsetsSourceProviderTest.java b/services/tests/wmtests/src/com/android/server/wm/ImeInsetsSourceProviderTest.java
index f70dcebce30d..ea175a5a52b0 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ImeInsetsSourceProviderTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ImeInsetsSourceProviderTest.java
@@ -51,7 +51,6 @@ public class ImeInsetsSourceProviderTest extends WindowTestsBase {
     public void setUp() throws Exception {
         mImeProvider = mDisplayContent.getInsetsStateController().getImeSourceProvider();
         mImeProvider.getSource().setVisible(true);
-        mWm.mAnimator.ready();
     }
 
     @Test
@@ -152,8 +151,8 @@ public class ImeInsetsSourceProviderTest extends WindowTestsBase {
         assertFalse(mImeProvider.isScheduledAndReadyToShowIme());
         assertFalse(mImeProvider.isImeShowing());
 
-        // Waiting for the afterPrepareSurfacesRunnable picks up the show scheduled above.
-        waitUntilWindowAnimatorIdle();
+        // Starting the afterPrepareSurfacesRunnable picks up the show scheduled above.
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
         // No longer scheduled as it was already shown.
         assertFalse(mImeProvider.isScheduledAndReadyToShowIme());
         assertTrue(mImeProvider.isImeShowing());
diff --git a/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java b/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java
index e8d089c61362..619080770f1e 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java
@@ -187,8 +187,8 @@ public class InsetsSourceProviderTest extends WindowTestsBase {
         assertEquals(mProvider.getControlTarget(), target);
         assertNull(mProvider.getLeash(target));
 
-        // Set the leash to be ready for dispatching.
-        mProvider.mIsLeashReadyForDispatching = true;
+        // After surface transactions are applied, the leash is ready for dispatching.
+        mProvider.onSurfaceTransactionApplied();
         assertNotNull(mProvider.getLeash(target));
 
         // We do have fake control for the fake control target, but that has no leash.
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 d0d7c06bd706..964264d82b65 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
@@ -388,7 +388,6 @@ public class InsetsStateControllerTest extends WindowTestsBase {
         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 c5cbedb9193c..6f7d0dced484 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RemoteAnimationControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RemoteAnimationControllerTest.java
@@ -110,7 +110,6 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
         runWithScissors(mWm.mH, () -> mHandler = new TestHandler(null, mClock), 0);
         mController = new RemoteAnimationController(mWm, mDisplayContent, mAdapter,
                 mHandler, false /*isActivityEmbedding*/);
-        mWm.mAnimator.ready();
     }
 
     private WindowState createAppOverlayWindow() {
@@ -134,7 +133,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_ACTIVITY_OPEN);
-            waitUntilWindowAnimatorIdle();
+            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -166,7 +165,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_ACTIVITY_OPEN);
-            waitUntilWindowAnimatorIdle();
+            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -288,7 +287,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
         adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                 mFinishedCallback);
         mController.goodToGo(TRANSIT_OLD_ACTIVITY_OPEN);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
         final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                 ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
         final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -334,7 +333,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
         task.applyAnimationUnchecked(null /* lp */, true /* enter */, TRANSIT_OLD_TASK_OPEN,
                 false /* isVoiceInteraction */, null /* sources */);
         mController.goodToGo(TRANSIT_OLD_TASK_OPEN);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
         final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                 ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
         try {
@@ -361,7 +360,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
             ((AnimationAdapter) record.mThumbnailAdapter).startAnimation(mMockThumbnailLeash,
                     mMockTransaction, ANIMATION_TYPE_WINDOW_ANIMATION, mThumbnailFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_TASK_CHANGE_WINDOWING_MODE);
-            waitUntilWindowAnimatorIdle();
+            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -415,7 +414,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
             ((AnimationAdapter) record.mThumbnailAdapter).startAnimation(mMockThumbnailLeash,
                     mMockTransaction, ANIMATION_TYPE_WINDOW_ANIMATION, mThumbnailFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_TASK_CHANGE_WINDOWING_MODE);
-            waitUntilWindowAnimatorIdle();
+            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -469,7 +468,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
             ((AnimationAdapter) record.mThumbnailAdapter).startAnimation(mMockThumbnailLeash,
                     mMockTransaction, ANIMATION_TYPE_WINDOW_ANIMATION, mThumbnailFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_TASK_CHANGE_WINDOWING_MODE);
-            waitUntilWindowAnimatorIdle();
+            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -524,7 +523,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_ACTIVITY_OPEN);
-            waitUntilWindowAnimatorIdle();
+            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -557,7 +556,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_ACTIVITY_OPEN);
-            waitUntilWindowAnimatorIdle();
+            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -593,7 +592,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_KEYGUARD_GOING_AWAY);
-            waitUntilWindowAnimatorIdle();
+            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -643,7 +642,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
             adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                     mFinishedCallback);
             mController.goodToGo(TRANSIT_OLD_KEYGUARD_GOING_AWAY_ON_WALLPAPER);
-            waitUntilWindowAnimatorIdle();
+            mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
             final ArgumentCaptor<RemoteAnimationTarget[]> appsCaptor =
                     ArgumentCaptor.forClass(RemoteAnimationTarget[].class);
             final ArgumentCaptor<RemoteAnimationTarget[]> wallpapersCaptor =
@@ -743,7 +742,7 @@ public class RemoteAnimationControllerTest extends WindowTestsBase {
         adapter.startAnimation(mMockLeash, mMockTransaction, ANIMATION_TYPE_APP_TRANSITION,
                 mFinishedCallback);
         mController.goodToGo(transit);
-        waitUntilWindowAnimatorIdle();
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
         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 de3c1313b0ee..1e39f0b963b7 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
@@ -571,7 +571,7 @@ public class SystemServicesTestRule implements TestRule {
         // 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.addAfterPrepareSurfacesRunnable(() -> {
+        wm.mAnimator.getChoreographer().postFrameCallback(time -> {
             synchronized (currentMessagesProcessed) {
                 currentMessagesProcessed.set(true);
                 currentMessagesProcessed.notifyAll();