Merge "Optimize activity snapshot for shell transition." into udc-qpr-dev am: 5f49480288 am: 5e1b56982c

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/23898517

Change-Id: Icc5f57699de722f2f42543a9231175f5268d371c
Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
diff --git a/services/core/java/com/android/server/wm/ActivitySnapshotController.java b/services/core/java/com/android/server/wm/ActivitySnapshotController.java
index 105b2bb..148bf9b 100644
--- a/services/core/java/com/android/server/wm/ActivitySnapshotController.java
+++ b/services/core/java/com/android/server/wm/ActivitySnapshotController.java
@@ -16,16 +16,14 @@
 
 package com.android.server.wm;
 
-import static com.android.server.wm.SnapshotController.ACTIVITY_CLOSE;
-import static com.android.server.wm.SnapshotController.ACTIVITY_OPEN;
-import static com.android.server.wm.SnapshotController.TASK_CLOSE;
-import static com.android.server.wm.SnapshotController.TASK_OPEN;
+import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.ActivityManager;
 import android.os.Environment;
 import android.os.SystemProperties;
+import android.os.Trace;
 import android.util.ArraySet;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -35,7 +33,6 @@
 import com.android.server.LocalServices;
 import com.android.server.pm.UserManagerInternal;
 import com.android.server.wm.BaseAppSnapshotPersister.PersistInfoProvider;
-import com.android.server.wm.SnapshotController.TransitionState;
 
 import java.io.File;
 import java.util.ArrayList;
@@ -62,12 +59,6 @@
     static final String SNAPSHOTS_DIRNAME = "activity_snapshots";
 
     /**
-     * The pending activities which should capture snapshot when process transition finish.
-     */
-    @VisibleForTesting
-    final ArraySet<ActivityRecord> mPendingCaptureActivity = new ArraySet<>();
-
-    /**
      * The pending activities which should remove snapshot from memory when process transition
      * finish.
      */
@@ -86,6 +77,10 @@
     @VisibleForTesting
     final ArraySet<ActivityRecord> mPendingLoadActivity = new ArraySet<>();
 
+    private final ArraySet<ActivityRecord> mOnBackPressedActivities = new ArraySet<>();
+
+    private final ArrayList<ActivityRecord> mTmpBelowActivities = new ArrayList<>();
+    private final ArrayList<WindowContainer> mTmpTransitionParticipants = new ArrayList<>();
     private final SnapshotPersistQueue mSnapshotPersistQueue;
     private final PersistInfoProvider mPersistInfoProvider;
     private final AppSnapshotLoader mSnapshotLoader;
@@ -117,20 +112,6 @@
         setSnapshotEnabled(snapshotEnabled);
     }
 
-    void systemReady() {
-        if (shouldDisableSnapshots()) {
-            return;
-        }
-        mService.mSnapshotController.registerTransitionStateConsumer(
-                ACTIVITY_OPEN, this::handleOpenActivityTransition);
-        mService.mSnapshotController.registerTransitionStateConsumer(
-                ACTIVITY_CLOSE, this::handleCloseActivityTransition);
-        mService.mSnapshotController.registerTransitionStateConsumer(
-                TASK_OPEN, this::handleOpenTaskTransition);
-        mService.mSnapshotController.registerTransitionStateConsumer(
-                TASK_CLOSE, this::handleCloseTaskTransition);
-    }
-
     @Override
     protected float initSnapshotScale() {
         final float config = mService.mContext.getResources().getFloat(
@@ -173,6 +154,7 @@
 
                         @Override
                         void write() {
+                            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "cleanUpUserFiles");
                             final File file = mPersistInfoProvider.getDirectory(userId);
                             if (file.exists()) {
                                 final File[] contents = file.listFiles();
@@ -182,15 +164,30 @@
                                     }
                                 }
                             }
+                            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
                         }
                     });
         }
     }
 
+    void addOnBackPressedActivity(ActivityRecord ar) {
+        if (shouldDisableSnapshots()) {
+            return;
+        }
+        mOnBackPressedActivities.add(ar);
+    }
+
+    void clearOnBackPressedActivities() {
+        if (shouldDisableSnapshots()) {
+            return;
+        }
+        mOnBackPressedActivities.clear();
+    }
+
     /**
-     * Prepare to handle on transition start. Clear all temporary fields.
+     * Prepare to collect any change for snapshots processing. Clear all temporary fields.
      */
-    void preTransitionStart() {
+    void beginSnapshotProcess() {
         if (shouldDisableSnapshots()) {
             return;
         }
@@ -198,18 +195,22 @@
     }
 
     /**
-     * on transition start has notified, start process data.
+     * End collect any change for snapshots processing, start process data.
      */
-    void postTransitionStart() {
+    void endSnapshotProcess() {
         if (shouldDisableSnapshots()) {
             return;
         }
-        onCommitTransition();
+        for (int i = mOnBackPressedActivities.size() - 1; i >= 0; --i) {
+            handleActivityTransition(mOnBackPressedActivities.valueAt(i));
+        }
+        mOnBackPressedActivities.clear();
+        mTmpTransitionParticipants.clear();
+        postProcess();
     }
 
     @VisibleForTesting
     void resetTmpFields() {
-        mPendingCaptureActivity.clear();
         mPendingRemoveActivity.clear();
         mPendingDeleteActivity.clear();
         mPendingLoadActivity.clear();
@@ -218,31 +219,13 @@
     /**
      * Start process all pending activities for a transition.
      */
-    private void onCommitTransition() {
+    private void postProcess() {
         if (DEBUG) {
-            Slog.d(TAG, "ActivitySnapshotController#onCommitTransition result:"
-                    + " capture " + mPendingCaptureActivity
+            Slog.d(TAG, "ActivitySnapshotController#postProcess result:"
                     + " remove " + mPendingRemoveActivity
                     + " delete " + mPendingDeleteActivity
                     + " load " + mPendingLoadActivity);
         }
-        // task snapshots
-        for (int i = mPendingCaptureActivity.size() - 1; i >= 0; i--) {
-            recordSnapshot(mPendingCaptureActivity.valueAt(i));
-        }
-        // clear mTmpRemoveActivity from cache
-        for (int i = mPendingRemoveActivity.size() - 1; i >= 0; i--) {
-            final ActivityRecord ar = mPendingRemoveActivity.valueAt(i);
-            final int code = getSystemHashCode(ar);
-            mCache.onIdRemoved(code);
-        }
-        // clear snapshot on cache and delete files
-        for (int i = mPendingDeleteActivity.size() - 1; i >= 0; i--) {
-            final ActivityRecord ar = mPendingDeleteActivity.valueAt(i);
-            final int code = getSystemHashCode(ar);
-            mCache.onIdRemoved(code);
-            removeIfUserSavedFileExist(code, ar.mUserId);
-        }
         // load snapshot to cache
         for (int i = mPendingLoadActivity.size() - 1; i >= 0; i--) {
             final ActivityRecord ar = mPendingLoadActivity.valueAt(i);
@@ -258,6 +241,8 @@
                             new SnapshotPersistQueue.WriteQueueItem(mPersistInfoProvider) {
                                 @Override
                                 void write() {
+                                    Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER,
+                                            "load_activity_snapshot");
                                     final TaskSnapshot snapshot = mSnapshotLoader.loadTask(code,
                                             userId, false /* loadLowResolutionBitmap */);
                                     synchronized (mService.getWindowManagerLock()) {
@@ -265,16 +250,36 @@
                                             mCache.putSnapshot(ar, snapshot);
                                         }
                                     }
+                                    Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
                                 }
                             });
                 }
             }
         }
+        // clear mTmpRemoveActivity from cache
+        for (int i = mPendingRemoveActivity.size() - 1; i >= 0; i--) {
+            final ActivityRecord ar = mPendingRemoveActivity.valueAt(i);
+            final int code = getSystemHashCode(ar);
+            mCache.onIdRemoved(code);
+        }
+        // clear snapshot on cache and delete files
+        for (int i = mPendingDeleteActivity.size() - 1; i >= 0; i--) {
+            final ActivityRecord ar = mPendingDeleteActivity.valueAt(i);
+            final int code = getSystemHashCode(ar);
+            mCache.onIdRemoved(code);
+            removeIfUserSavedFileExist(code, ar.mUserId);
+        }
         // don't keep any reference
         resetTmpFields();
     }
 
-    private void recordSnapshot(ActivityRecord activity) {
+    void recordSnapshot(ActivityRecord activity) {
+        if (shouldDisableSnapshots()) {
+            return;
+        }
+        if (DEBUG) {
+            Slog.d(TAG, "ActivitySnapshotController#recordSnapshot " + activity);
+        }
         final TaskSnapshot snapshot = recordSnapshotInner(activity, false /* allowSnapshotHome */);
         if (snapshot != null) {
             final int code = getSystemHashCode(activity);
@@ -285,15 +290,20 @@
     /**
      * Called when the visibility of an app changes outside the regular app transition flow.
      */
-    void notifyAppVisibilityChanged(ActivityRecord appWindowToken, boolean visible) {
+    void notifyAppVisibilityChanged(ActivityRecord ar, boolean visible) {
         if (shouldDisableSnapshots()) {
             return;
         }
+        final Task task = ar.getTask();
+        if (task == null) {
+            return;
+        }
+        // Doesn't need to capture activity snapshot when it converts from translucent.
         if (!visible) {
             resetTmpFields();
-            addBelowTopActivityIfExist(appWindowToken.getTask(), mPendingRemoveActivity,
+            addBelowActivityIfExist(ar, mPendingRemoveActivity, false,
                     "remove-snapshot");
-            onCommitTransition();
+            postProcess();
         }
     }
 
@@ -301,65 +311,146 @@
         return System.identityHashCode(activity);
     }
 
-    void handleOpenActivityTransition(TransitionState<ActivityRecord> transitionState) {
-        ArraySet<ActivityRecord> participant = transitionState.getParticipant(false /* open */);
-        for (ActivityRecord ar : participant) {
-            mPendingCaptureActivity.add(ar);
-            // remove the snapshot for the one below close
-            final ActivityRecord below = ar.getTask().getActivityBelow(ar);
-            if (below != null) {
-                mPendingRemoveActivity.add(below);
+    @VisibleForTesting
+    void handleTransitionFinish(@NonNull ArrayList<WindowContainer> windows) {
+        mTmpTransitionParticipants.clear();
+        mTmpTransitionParticipants.addAll(windows);
+        for (int i = mTmpTransitionParticipants.size() - 1; i >= 0; --i) {
+            final WindowContainer next = mTmpTransitionParticipants.get(i);
+            if (next.asTask() != null) {
+                handleTaskTransition(next.asTask());
+            } else if (next.asTaskFragment() != null) {
+                final TaskFragment tf = next.asTaskFragment();
+                final ActivityRecord ar = tf.getTopMostActivity();
+                if (ar != null) {
+                    handleActivityTransition(ar);
+                }
+            } else if (next.asActivityRecord() != null) {
+                handleActivityTransition(next.asActivityRecord());
             }
         }
     }
 
-    void handleCloseActivityTransition(TransitionState<ActivityRecord> transitionState) {
-        ArraySet<ActivityRecord> participant = transitionState.getParticipant(true /* open */);
-        for (ActivityRecord ar : participant) {
+    private void handleActivityTransition(@NonNull ActivityRecord ar) {
+        if (shouldDisableSnapshots()) {
+            return;
+        }
+        if (ar.isVisibleRequested()) {
             mPendingDeleteActivity.add(ar);
             // load next one if exists.
-            final ActivityRecord below = ar.getTask().getActivityBelow(ar);
-            if (below != null) {
-                mPendingLoadActivity.add(below);
-            }
+            addBelowActivityIfExist(ar, mPendingLoadActivity, true, "load-snapshot");
+        } else {
+            // remove the snapshot for the one below close
+            addBelowActivityIfExist(ar, mPendingRemoveActivity, true, "remove-snapshot");
         }
     }
 
-    void handleCloseTaskTransition(TransitionState<Task> closeTaskTransitionRecord) {
-        ArraySet<Task> participant = closeTaskTransitionRecord.getParticipant(false /* open */);
-        for (Task close : participant) {
-            // this is close task transition
-            // remove the N - 1 from cache
-            addBelowTopActivityIfExist(close, mPendingRemoveActivity, "remove-snapshot");
+    private void handleTaskTransition(Task task) {
+        if (shouldDisableSnapshots()) {
+            return;
         }
-    }
-
-    void handleOpenTaskTransition(TransitionState<Task> openTaskTransitionRecord) {
-        ArraySet<Task> participant = openTaskTransitionRecord.getParticipant(true /* open */);
-        for (Task open : participant) {
-            // this is close task transition
-            // remove the N - 1 from cache
-            addBelowTopActivityIfExist(open, mPendingLoadActivity, "load-snapshot");
+        final ActivityRecord topActivity = task.getTopMostActivity();
+        if (topActivity == null) {
+            return;
+        }
+        if (task.isVisibleRequested()) {
+            // this is open task transition
+            // load the N - 1 to cache
+            addBelowActivityIfExist(topActivity, mPendingLoadActivity, true, "load-snapshot");
             // Move the activities to top of mSavedFilesInOrder, so when purge happen, there
             // will trim the persisted files from the most non-accessed.
-            adjustSavedFileOrder(open);
+            adjustSavedFileOrder(task);
+        } else {
+            // this is close task transition
+            // remove the N - 1 from cache
+            addBelowActivityIfExist(topActivity, mPendingRemoveActivity, true, "remove-snapshot");
         }
     }
 
-    // Add the top -1 activity to a set if it exists.
-    private void addBelowTopActivityIfExist(Task task, ArraySet<ActivityRecord> set,
-            String debugMessage) {
-        final ActivityRecord topActivity = task.getTopMostActivity();
-        if (topActivity != null) {
-            final ActivityRecord below = task.getActivityBelow(topActivity);
-            if (below != null) {
-                set.add(below);
-                if (DEBUG) {
-                    Slog.d(TAG, "ActivitySnapshotController#addBelowTopActivityIfExist "
-                            + below + " from " + debugMessage);
-                }
+    /**
+     * Add the top -1 activity to a set if it exists.
+     * @param inTransition true if the activity must participant in transition.
+     */
+    private void addBelowActivityIfExist(ActivityRecord currentActivity,
+            ArraySet<ActivityRecord> set, boolean inTransition, String debugMessage) {
+        getActivityBelow(currentActivity, inTransition, mTmpBelowActivities);
+        for (int i = mTmpBelowActivities.size() - 1; i >= 0; --i) {
+            set.add(mTmpBelowActivities.get(i));
+            if (DEBUG) {
+                Slog.d(TAG, "ActivitySnapshotController#addBelowTopActivityIfExist "
+                        + mTmpBelowActivities.get(i) + " from " + debugMessage);
             }
         }
+        mTmpBelowActivities.clear();
+    }
+
+    private void getActivityBelow(ActivityRecord currentActivity, boolean inTransition,
+            ArrayList<ActivityRecord> result) {
+        final Task currentTask = currentActivity.getTask();
+        if (currentTask == null) {
+            return;
+        }
+        final ActivityRecord initPrev = currentTask.getActivityBelow(currentActivity);
+        if (initPrev == null) {
+            return;
+        }
+        final TaskFragment currTF = currentActivity.getTaskFragment();
+        final TaskFragment prevTF = initPrev.getTaskFragment();
+        final TaskFragment prevAdjacentTF = prevTF != null
+                ? prevTF.getAdjacentTaskFragment() : null;
+        if (currTF == prevTF && currTF != null || prevAdjacentTF == null) {
+            // Current activity and previous one is in the same task fragment, or
+            // previous activity is not in a task fragment, or
+            // previous activity's task fragment doesn't adjacent to any others.
+            if (!inTransition || isInParticipant(initPrev, mTmpTransitionParticipants)) {
+                result.add(initPrev);
+            }
+            return;
+        }
+
+        if (prevAdjacentTF == currTF) {
+            // previous activity A is adjacent to current activity B.
+            // Try to find anyone below previous activityA, which are C and D if exists.
+            // A | B
+            // C (| D)
+            getActivityBelow(initPrev, inTransition, result);
+        } else {
+            // previous activity C isn't adjacent to current activity A.
+            // A
+            // B | C
+            final Task prevAdjacentTask = prevAdjacentTF.getTask();
+            if (prevAdjacentTask == currentTask) {
+                final int currentIndex = currTF != null
+                        ? currentTask.mChildren.indexOf(currTF)
+                        : currentTask.mChildren.indexOf(currentActivity);
+                final int prevAdjacentIndex =
+                        prevAdjacentTask.mChildren.indexOf(prevAdjacentTF);
+                // prevAdjacentTF already above currentActivity
+                if (prevAdjacentIndex > currentIndex) {
+                    return;
+                }
+            }
+            if (!inTransition || isInParticipant(initPrev, mTmpTransitionParticipants)) {
+                result.add(initPrev);
+            }
+            // prevAdjacentTF is adjacent to another one
+            final ActivityRecord prevAdjacentActivity = prevAdjacentTF.getTopMostActivity();
+            if (prevAdjacentActivity != null && (!inTransition
+                    || isInParticipant(prevAdjacentActivity, mTmpTransitionParticipants))) {
+                result.add(prevAdjacentActivity);
+            }
+        }
+    }
+
+    static boolean isInParticipant(ActivityRecord ar,
+            ArrayList<WindowContainer> transitionParticipants) {
+        for (int i = transitionParticipants.size() - 1; i >= 0; --i) {
+            final WindowContainer wc = transitionParticipants.get(i);
+            if (ar == wc || ar.isDescendantOf(wc)) {
+                return true;
+            }
+        }
+        return false;
     }
 
     private void adjustSavedFileOrder(Task nextTopTask) {
@@ -376,6 +467,9 @@
 
     @Override
     void onAppRemoved(ActivityRecord activity) {
+        if (shouldDisableSnapshots()) {
+            return;
+        }
         super.onAppRemoved(activity);
         final int code = getSystemHashCode(activity);
         removeIfUserSavedFileExist(code, activity.mUserId);
@@ -386,6 +480,9 @@
 
     @Override
     void onAppDied(ActivityRecord activity) {
+        if (shouldDisableSnapshots()) {
+            return;
+        }
         super.onAppDied(activity);
         final int code = getSystemHashCode(activity);
         removeIfUserSavedFileExist(code, activity.mUserId);
@@ -440,7 +537,7 @@
     private void removeIfUserSavedFileExist(int code, int userId) {
         final UserSavedFile usf = getUserFiles(userId).get(code);
         if (usf != null) {
-            mUserSavedFiles.remove(code);
+            mUserSavedFiles.get(userId).remove(code);
             mSavedFilesInOrder.remove(usf);
             mPersister.removeSnap(code, userId);
         }
@@ -490,11 +587,13 @@
                     new SnapshotPersistQueue.WriteQueueItem(mPersistInfoProvider) {
                         @Override
                         void write() {
+                            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "activity_remove_files");
                             for (int i = files.size() - 1; i >= 0; --i) {
                                 final UserSavedFile usf = files.get(i);
                                 mSnapshotPersistQueue.deleteSnapshot(
                                         usf.mFileId, usf.mUserId, mPersistInfoProvider);
                             }
+                            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
                         }
                     });
         }
diff --git a/services/core/java/com/android/server/wm/BackNavigationController.java b/services/core/java/com/android/server/wm/BackNavigationController.java
index 976641b..993c016 100644
--- a/services/core/java/com/android/server/wm/BackNavigationController.java
+++ b/services/core/java/com/android/server/wm/BackNavigationController.java
@@ -1177,6 +1177,8 @@
                 if (!composeAnimations(mCloseTarget, mOpenTarget, openActivity)) {
                     return null;
                 }
+                mCloseTarget.mTransitionController.mSnapshotController
+                        .mActivitySnapshotController.clearOnBackPressedActivities();
                 applyPreviewStrategy(mOpenAdaptor, openActivity);
 
                 final IBackAnimationFinishedCallback callback = makeAnimationFinishedCallback();
@@ -1222,6 +1224,8 @@
             // Call it again to make sure the activity could be visible while handling the pending
             // animation.
             activity.commitVisibility(true, true);
+            activity.mTransitionController.mSnapshotController
+                    .mActivitySnapshotController.addOnBackPressedActivity(activity);
         }
         activity.mLaunchTaskBehind = true;
 
@@ -1248,6 +1252,9 @@
         // Restore the launch-behind state.
         activity.mTaskSupervisor.scheduleLaunchTaskBehindComplete(activity.token);
         activity.mLaunchTaskBehind = false;
+        // Ignore all change
+        activity.mTransitionController.mSnapshotController
+                .mActivitySnapshotController.clearOnBackPressedActivities();
         ProtoLog.d(WM_DEBUG_BACK_PREVIEW,
                 "Setting Activity.mLauncherTaskBehind to false. Activity=%s",
                 activity);
diff --git a/services/core/java/com/android/server/wm/SnapshotController.java b/services/core/java/com/android/server/wm/SnapshotController.java
index badcfa9..37f9730 100644
--- a/services/core/java/com/android/server/wm/SnapshotController.java
+++ b/services/core/java/com/android/server/wm/SnapshotController.java
@@ -16,185 +16,36 @@
 
 package com.android.server.wm;
 
+import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
+import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 import static android.view.WindowManager.TRANSIT_CLOSE;
 import static android.view.WindowManager.TRANSIT_FIRST_CUSTOM;
 import static android.view.WindowManager.TRANSIT_OPEN;
 import static android.view.WindowManager.TRANSIT_TO_BACK;
 import static android.view.WindowManager.TRANSIT_TO_FRONT;
 
-import android.annotation.IntDef;
-import android.util.ArraySet;
-import android.util.Slog;
-import android.util.SparseArray;
+import android.os.Trace;
 import android.view.WindowManager;
 
-import com.android.internal.annotations.VisibleForTesting;
-
 import java.io.PrintWriter;
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
-import java.util.function.Consumer;
 
 /**
  * Integrates common functionality from TaskSnapshotController and ActivitySnapshotController.
  */
 class SnapshotController {
-    private static final boolean DEBUG = false;
-    private static final String TAG = AbsAppSnapshotController.TAG;
-
-    static final int ACTIVITY_OPEN = 1;
-    static final int ACTIVITY_CLOSE = 2;
-    static final int TASK_OPEN = 4;
-    static final int TASK_CLOSE = 8;
-    /** @hide */
-    @Retention(RetentionPolicy.SOURCE)
-    @IntDef(
-            value = {ACTIVITY_OPEN,
-                    ACTIVITY_CLOSE,
-                    TASK_OPEN,
-                    TASK_CLOSE})
-    @interface TransitionStateType {}
-
     private final SnapshotPersistQueue mSnapshotPersistQueue;
     final TaskSnapshotController mTaskSnapshotController;
     final ActivitySnapshotController mActivitySnapshotController;
 
-    private final ArraySet<Task> mTmpCloseTasks = new ArraySet<>();
-    private final ArraySet<Task> mTmpOpenTasks = new ArraySet<>();
-
-    private final SparseArray<TransitionState> mTmpOpenCloseRecord = new SparseArray<>();
-    private final ArraySet<Integer> mTmpAnalysisRecord = new ArraySet<>();
-    private final SparseArray<ArrayList<Consumer<TransitionState>>> mTransitionStateConsumer =
-            new SparseArray<>();
-    private int mActivatedType;
-
-    private final ActivityOrderCheck mActivityOrderCheck = new ActivityOrderCheck();
-    private final ActivityOrderCheck.AnalysisResult mResultHandler = (type, close, open) -> {
-        addTransitionRecord(type, true/* open */, open);
-        addTransitionRecord(type, false/* open */, close);
-    };
-
-    private static class ActivityOrderCheck {
-        private ActivityRecord mOpenActivity;
-        private ActivityRecord mCloseActivity;
-        private int mOpenIndex = -1;
-        private int mCloseIndex = -1;
-
-        private void reset() {
-            mOpenActivity = null;
-            mCloseActivity = null;
-            mOpenIndex = -1;
-            mCloseIndex = -1;
-        }
-
-        private void setTarget(boolean open, ActivityRecord ar, int index) {
-            if (open) {
-                mOpenActivity = ar;
-                mOpenIndex = index;
-            } else {
-                mCloseActivity = ar;
-                mCloseIndex = index;
-            }
-        }
-
-        void analysisOrder(ArraySet<ActivityRecord> closeApps,
-                ArraySet<ActivityRecord> openApps, Task task, AnalysisResult result) {
-            for (int j = closeApps.size() - 1; j >= 0; j--) {
-                final ActivityRecord ar = closeApps.valueAt(j);
-                if (ar.getTask() == task) {
-                    setTarget(false, ar, task.mChildren.indexOf(ar));
-                    break;
-                }
-            }
-            for (int j = openApps.size() - 1; j >= 0; j--) {
-                final ActivityRecord ar = openApps.valueAt(j);
-                if (ar.getTask() == task) {
-                    setTarget(true, ar, task.mChildren.indexOf(ar));
-                    break;
-                }
-            }
-            if (mOpenIndex > mCloseIndex && mCloseIndex != -1) {
-                result.onCheckResult(ACTIVITY_OPEN, mCloseActivity, mOpenActivity);
-            } else if (mOpenIndex < mCloseIndex && mOpenIndex != -1) {
-                result.onCheckResult(ACTIVITY_CLOSE, mCloseActivity, mOpenActivity);
-            }
-            reset();
-        }
-        private interface AnalysisResult {
-            void onCheckResult(@TransitionStateType int type,
-                    ActivityRecord close, ActivityRecord open);
-        }
-    }
-
-    private void addTransitionRecord(int type, boolean open, WindowContainer target) {
-        TransitionState record = mTmpOpenCloseRecord.get(type);
-        if (record == null) {
-            record =  new TransitionState();
-            mTmpOpenCloseRecord.set(type, record);
-        }
-        record.addParticipant(target, open);
-        mTmpAnalysisRecord.add(type);
-    }
-
-    private void clearRecord() {
-        mTmpOpenCloseRecord.clear();
-        mTmpAnalysisRecord.clear();
-    }
-
-    static class TransitionState<TYPE extends WindowContainer> {
-        private final ArraySet<TYPE> mOpenParticipant = new ArraySet<>();
-        private final ArraySet<TYPE> mCloseParticipant = new ArraySet<>();
-
-        void addParticipant(TYPE target, boolean open) {
-            final ArraySet<TYPE> participant = open
-                    ? mOpenParticipant : mCloseParticipant;
-            participant.add(target);
-        }
-
-        ArraySet<TYPE> getParticipant(boolean open) {
-            return open ? mOpenParticipant : mCloseParticipant;
-        }
-    }
-
     SnapshotController(WindowManagerService wms) {
         mSnapshotPersistQueue = new SnapshotPersistQueue();
         mTaskSnapshotController = new TaskSnapshotController(wms, mSnapshotPersistQueue);
         mActivitySnapshotController = new ActivitySnapshotController(wms, mSnapshotPersistQueue);
     }
 
-    void registerTransitionStateConsumer(@TransitionStateType int type,
-            Consumer<TransitionState> consumer) {
-        ArrayList<Consumer<TransitionState>> consumers = mTransitionStateConsumer.get(type);
-        if (consumers == null) {
-            consumers = new ArrayList<>();
-            mTransitionStateConsumer.set(type, consumers);
-        }
-        if (!consumers.contains(consumer)) {
-            consumers.add(consumer);
-        }
-        mActivatedType |= type;
-    }
-
-    void unregisterTransitionStateConsumer(int type, Consumer<TransitionState> consumer) {
-        final ArrayList<Consumer<TransitionState>> consumers = mTransitionStateConsumer.get(type);
-        if (consumers == null) {
-            return;
-        }
-        consumers.remove(consumer);
-        if (consumers.size() == 0) {
-            mActivatedType &= ~type;
-        }
-    }
-
-    private boolean hasTransitionStateConsumer(@TransitionStateType int type) {
-        return (mActivatedType & type) != 0;
-    }
-
     void systemReady() {
         mSnapshotPersistQueue.systemReady();
-        mTaskSnapshotController.systemReady();
-        mActivitySnapshotController.systemReady();
     }
 
     void setPause(boolean paused) {
@@ -212,47 +63,69 @@
     }
 
     void notifyAppVisibilityChanged(ActivityRecord appWindowToken, boolean visible) {
-        if (!visible && hasTransitionStateConsumer(TASK_CLOSE)) {
-            final Task task = appWindowToken.getTask();
-            if (task == null || task.isVisibleRequested()) {
-                return;
+        mActivitySnapshotController.notifyAppVisibilityChanged(appWindowToken, visible);
+    }
+
+    // For legacy transition, which won't support activity snapshot
+    void onTransitionStarting(DisplayContent displayContent) {
+        mTaskSnapshotController.handleClosingApps(displayContent.mClosingApps);
+    }
+
+    // For shell transition, record snapshots before transaction start.
+    void onTransactionReady(@WindowManager.TransitionType int type,
+            ArrayList<Transition.ChangeInfo> changeInfos) {
+        final boolean isTransitionOpen = isTransitionOpen(type);
+        final boolean isTransitionClose = isTransitionClose(type);
+        if (!isTransitionOpen && !isTransitionClose && type < TRANSIT_FIRST_CUSTOM) {
+            return;
+        }
+        for (int i = changeInfos.size() - 1; i >= 0; --i) {
+            Transition.ChangeInfo info = changeInfos.get(i);
+            // Intentionally skip record snapshot for changes originated from PiP.
+            if (info.mWindowingMode == WINDOWING_MODE_PINNED) continue;
+            if (info.mContainer.asTask() != null && !info.mContainer.isVisibleRequested()) {
+                mTaskSnapshotController.recordSnapshot(info.mContainer.asTask(),
+                        false /* allowSnapshotHome */);
             }
-            // close task transition
-            addTransitionRecord(TASK_CLOSE, false /*open*/, task);
-            mActivitySnapshotController.preTransitionStart();
-            notifyTransition(TASK_CLOSE);
-            mActivitySnapshotController.postTransitionStart();
-            clearRecord();
+            // Won't need to capture activity snapshot in close transition.
+            if (isTransitionClose) {
+                continue;
+            }
+            if (info.mContainer.asActivityRecord() != null
+                    || info.mContainer.asTaskFragment() != null) {
+                final TaskFragment tf = info.mContainer.asTaskFragment();
+                final ActivityRecord ar = tf != null ? tf.getTopMostActivity()
+                        : info.mContainer.asActivityRecord();
+                final boolean taskVis = ar != null && ar.getTask().isVisibleRequested();
+                if (ar != null && !ar.isVisibleRequested() && taskVis) {
+                    mActivitySnapshotController.recordSnapshot(ar);
+                }
+            }
         }
     }
 
-    // For legacy transition
-    void onTransitionStarting(DisplayContent displayContent) {
-        handleAppTransition(displayContent.mClosingApps, displayContent.mOpeningApps);
-    }
-
-    // For shell transition, adapt to legacy transition.
-    void onTransitionReady(@WindowManager.TransitionType int type,
-            ArraySet<WindowContainer> participants) {
+    void onTransitionFinish(@WindowManager.TransitionType int type,
+            ArrayList<Transition.ChangeInfo> changeInfos) {
         final boolean isTransitionOpen = isTransitionOpen(type);
         final boolean isTransitionClose = isTransitionClose(type);
         if (!isTransitionOpen && !isTransitionClose && type < TRANSIT_FIRST_CUSTOM
-                || (mActivatedType == 0)) {
+                || (changeInfos.isEmpty())) {
             return;
         }
-        final ArraySet<ActivityRecord> openingApps = new ArraySet<>();
-        final ArraySet<ActivityRecord> closingApps = new ArraySet<>();
-
-        for (int i = participants.size() - 1; i >= 0; --i) {
-            final ActivityRecord ar = participants.valueAt(i).asActivityRecord();
-            if (ar == null || ar.getTask() == null) continue;
-            if (ar.isVisibleRequested()) {
-                openingApps.add(ar);
-            } else {
-                closingApps.add(ar);
+        Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "SnapshotController_analysis");
+        mActivitySnapshotController.beginSnapshotProcess();
+        final ArrayList<WindowContainer> windows = new ArrayList<>();
+        for (int i = changeInfos.size() - 1; i >= 0; --i) {
+            final WindowContainer wc = changeInfos.get(i).mContainer;
+            if (wc.asTask() == null && wc.asTaskFragment() == null
+                    && wc.asActivityRecord() == null) {
+                continue;
             }
+            windows.add(wc);
         }
-        handleAppTransition(closingApps, openingApps);
+        mActivitySnapshotController.handleTransitionFinish(windows);
+        mActivitySnapshotController.endSnapshotProcess();
+        Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
     }
 
     private static boolean isTransitionOpen(int type) {
@@ -262,78 +135,6 @@
         return type == TRANSIT_CLOSE || type == TRANSIT_TO_BACK;
     }
 
-    @VisibleForTesting
-    void handleAppTransition(ArraySet<ActivityRecord> closingApps,
-            ArraySet<ActivityRecord> openApps) {
-        if (mActivatedType == 0) {
-            return;
-        }
-        analysisTransition(closingApps, openApps);
-        mActivitySnapshotController.preTransitionStart();
-        for (Integer transitionType : mTmpAnalysisRecord) {
-            notifyTransition(transitionType);
-        }
-        mActivitySnapshotController.postTransitionStart();
-        clearRecord();
-    }
-
-    private void notifyTransition(int transitionType) {
-        final TransitionState record = mTmpOpenCloseRecord.get(transitionType);
-        final ArrayList<Consumer<TransitionState>> consumers =
-                mTransitionStateConsumer.get(transitionType);
-        for (Consumer<TransitionState> consumer : consumers) {
-            consumer.accept(record);
-        }
-    }
-
-    private void analysisTransition(ArraySet<ActivityRecord> closingApps,
-            ArraySet<ActivityRecord> openingApps) {
-        getParticipantTasks(closingApps, mTmpCloseTasks, false /* isOpen */);
-        getParticipantTasks(openingApps, mTmpOpenTasks, true /* isOpen */);
-        if (DEBUG) {
-            Slog.d(TAG, "AppSnapshotController#analysisTransition participants"
-                    + " mTmpCloseTasks " + mTmpCloseTasks
-                    + " mTmpOpenTasks " + mTmpOpenTasks);
-        }
-        for (int i = mTmpCloseTasks.size() - 1; i >= 0; i--) {
-            final Task closeTask = mTmpCloseTasks.valueAt(i);
-            if (mTmpOpenTasks.contains(closeTask)) {
-                if (hasTransitionStateConsumer(ACTIVITY_OPEN)
-                        || hasTransitionStateConsumer(ACTIVITY_CLOSE)) {
-                    mActivityOrderCheck.analysisOrder(closingApps, openingApps, closeTask,
-                            mResultHandler);
-                }
-            } else if (hasTransitionStateConsumer(TASK_CLOSE)) {
-                // close task transition
-                addTransitionRecord(TASK_CLOSE, false /*open*/, closeTask);
-            }
-        }
-        if (hasTransitionStateConsumer(TASK_OPEN)) {
-            for (int i = mTmpOpenTasks.size() - 1; i >= 0; i--) {
-                final Task openTask = mTmpOpenTasks.valueAt(i);
-                if (!mTmpCloseTasks.contains(openTask)) {
-                    // this is open task transition
-                    addTransitionRecord(TASK_OPEN, true /*open*/, openTask);
-                }
-            }
-        }
-        mTmpCloseTasks.clear();
-        mTmpOpenTasks.clear();
-    }
-
-    private void getParticipantTasks(ArraySet<ActivityRecord> activityRecords, ArraySet<Task> tasks,
-            boolean isOpen) {
-        for (int i = activityRecords.size() - 1; i >= 0; i--) {
-            final ActivityRecord activity = activityRecords.valueAt(i);
-            final Task task = activity.getTask();
-            if (task == null) continue;
-
-            if (isOpen == activity.isVisibleRequested()) {
-                tasks.add(task);
-            }
-        }
-    }
-
     void dump(PrintWriter pw, String prefix) {
         mTaskSnapshotController.dump(pw, prefix);
         mActivitySnapshotController.dump(pw, prefix);
diff --git a/services/core/java/com/android/server/wm/SnapshotPersistQueue.java b/services/core/java/com/android/server/wm/SnapshotPersistQueue.java
index 58e1c54..f4f641f 100644
--- a/services/core/java/com/android/server/wm/SnapshotPersistQueue.java
+++ b/services/core/java/com/android/server/wm/SnapshotPersistQueue.java
@@ -17,6 +17,7 @@
 package com.android.server.wm;
 
 import static android.graphics.Bitmap.CompressFormat.JPEG;
+import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
@@ -26,6 +27,7 @@
 import android.graphics.Bitmap;
 import android.os.Process;
 import android.os.SystemClock;
+import android.os.Trace;
 import android.util.AtomicFile;
 import android.util.Slog;
 import android.window.TaskSnapshot;
@@ -249,6 +251,7 @@
 
         @Override
         void write() {
+            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "StoreWriteQueueItem");
             if (!mPersistInfoProvider.createDirectory(mUserId)) {
                 Slog.e(TAG, "Unable to create snapshot directory for user dir="
                         + mPersistInfoProvider.getDirectory(mUserId));
@@ -263,6 +266,7 @@
             if (failed) {
                 deleteSnapshot(mId, mUserId, mPersistInfoProvider);
             }
+            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
         }
 
         boolean writeProto() {
@@ -373,7 +377,9 @@
 
         @Override
         void write() {
+            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "DeleteWriteQueueItem");
             deleteSnapshot(mId, mUserId, mPersistInfoProvider);
+            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
         }
     }
 }
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index c747c09..4eb4290 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -16,7 +16,6 @@
 
 package com.android.server.wm;
 
-import static com.android.server.wm.SnapshotController.TASK_CLOSE;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_SCREENSHOT;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 
@@ -77,13 +76,6 @@
         setSnapshotEnabled(snapshotEnabled);
     }
 
-    void systemReady() {
-        if (!shouldDisableSnapshots()) {
-            mService.mSnapshotController.registerTransitionStateConsumer(TASK_CLOSE,
-                    this::handleTaskClose);
-        }
-    }
-
     static PersistInfoProvider createPersistInfoProvider(WindowManagerService service,
             BaseAppSnapshotPersister.DirectoryResolver resolver) {
         final float highResTaskSnapshotScale = service.mContext.getResources().getFloat(
@@ -116,20 +108,23 @@
                 enableLowResSnapshots, lowResScaleFactor, use16BitFormat);
     }
 
-    void handleTaskClose(SnapshotController.TransitionState<Task> closeTaskTransitionRecord) {
+    // Still needed for legacy transition.(AppTransitionControllerTest)
+    void handleClosingApps(ArraySet<ActivityRecord> closingApps) {
         if (shouldDisableSnapshots()) {
             return;
         }
+        // We need to take a snapshot of the task if and only if all activities of the task are
+        // either closing or hidden.
         mTmpTasks.clear();
-        final ArraySet<Task> tasks = closeTaskTransitionRecord.getParticipant(false /* open */);
-        if (mService.mAtmService.getTransitionController().isShellTransitionsEnabled()) {
-            mTmpTasks.addAll(tasks);
-        } else {
-            for (Task task : tasks) {
-                getClosingTasksInner(task, mTmpTasks);
-            }
+        for (int i = closingApps.size() - 1; i >= 0; i--) {
+            final ActivityRecord activity = closingApps.valueAt(i);
+            final Task task = activity.getTask();
+            if (task == null) continue;
+
+            getClosingTasksInner(task, mTmpTasks);
         }
         snapshotTasks(mTmpTasks);
+        mTmpTasks.clear();
         mSkipClosingAppSnapshotTasks.clear();
     }
 
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotPersister.java b/services/core/java/com/android/server/wm/TaskSnapshotPersister.java
index cd15119..3e8c017 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotPersister.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotPersister.java
@@ -16,6 +16,9 @@
 
 package com.android.server.wm;
 
+import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
+
+import android.os.Trace;
 import android.util.ArraySet;
 import android.window.TaskSnapshot;
 
@@ -102,6 +105,7 @@
 
         @Override
         void write() {
+            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "RemoveObsoleteFilesQueueItem");
             final ArraySet<Integer> newPersistedTaskIds;
             synchronized (mLock) {
                 newPersistedTaskIds = new ArraySet<>(mPersistedTaskIdsSinceLastRemoveObsolete);
@@ -120,6 +124,7 @@
                     }
                 }
             }
+            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
         }
 
         @VisibleForTesting
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index aad17aa..1566bb2c 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -1191,8 +1191,6 @@
                                         "  Skipping post-transition snapshot for task %d",
                                         task.mTaskId);
                             }
-                            snapController.mActivitySnapshotController
-                                    .notifyAppVisibilityChanged(ar, false /* visible */);
                         }
                         ar.commitVisibility(false /* visible */, false /* performLayout */,
                                 true /* fromTransition */);
@@ -1389,6 +1387,7 @@
         // Handle back animation if it's already started.
         mController.mAtm.mBackNavigationController.onTransitionFinish(mTargets, this);
         mController.mFinishingTransition = null;
+        mController.mSnapshotController.onTransitionFinish(mType, mTargets);
     }
 
     void abort() {
@@ -1593,16 +1592,7 @@
         // transferred. If transition is transient, IME won't be moved during the transition and
         // the tasks are still live, so we take the snapshot at the end of the transition instead.
         if (mTransientLaunches == null) {
-            for (int i = mParticipants.size() - 1; i >= 0; --i) {
-                final ActivityRecord ar = mParticipants.valueAt(i).asActivityRecord();
-                if (ar == null || ar.getTask() == null
-                        || ar.getTask().isVisibleRequested()) continue;
-                final ChangeInfo change = mChanges.get(ar);
-                // Intentionally skip record snapshot for changes originated from PiP.
-                if (change != null && change.mWindowingMode == WINDOWING_MODE_PINNED) continue;
-                mController.mSnapshotController.mTaskSnapshotController.recordSnapshot(
-                        ar.getTask(), false /* allowSnapshotHome */);
-            }
+            mController.mSnapshotController.onTransactionReady(mType, mTargets);
         }
 
         // This is non-null only if display has changes. It handles the visible windows that don't
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivitySnapshotControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivitySnapshotControllerTests.java
index 0eca8c9..98f1843 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivitySnapshotControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivitySnapshotControllerTests.java
@@ -18,6 +18,9 @@
 
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
 
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+
 import static org.junit.Assert.assertEquals;
 
 import android.platform.test.annotations.Presubmit;
@@ -28,6 +31,8 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.ArrayList;
+
 /**
  * Test class for {@link ActivitySnapshotController}.
  *
@@ -42,13 +47,13 @@
     private ActivitySnapshotController mActivitySnapshotController;
     @Before
     public void setUp() throws Exception {
+        spyOn(mWm.mSnapshotController.mActivitySnapshotController);
         mActivitySnapshotController = mWm.mSnapshotController.mActivitySnapshotController;
+        doReturn(false).when(mActivitySnapshotController).shouldDisableSnapshots();
         mActivitySnapshotController.resetTmpFields();
     }
     @Test
     public void testOpenActivityTransition() {
-        final SnapshotController.TransitionState transitionState =
-                new SnapshotController.TransitionState();
         final Task task = createTask(mDisplayContent);
         // note for createAppWindow: the new child is added at index 0
         final WindowState openingWindow = createAppWindow(task,
@@ -59,14 +64,12 @@
                 "closingWindow");
         closingWindow.mActivityRecord.commitVisibility(
                 false /* visible */, true /* performLayout */);
-        transitionState.addParticipant(closingWindow.mActivityRecord, false);
-        transitionState.addParticipant(openingWindow.mActivityRecord, true);
-        mActivitySnapshotController.handleOpenActivityTransition(transitionState);
+        final ArrayList<WindowContainer> windows = new ArrayList<>();
+        windows.add(openingWindow.mActivityRecord);
+        windows.add(closingWindow.mActivityRecord);
+        mActivitySnapshotController.handleTransitionFinish(windows);
 
-        assertEquals(1, mActivitySnapshotController.mPendingCaptureActivity.size());
         assertEquals(0, mActivitySnapshotController.mPendingRemoveActivity.size());
-        assertEquals(closingWindow.mActivityRecord,
-                mActivitySnapshotController.mPendingCaptureActivity.valueAt(0));
         mActivitySnapshotController.resetTmpFields();
 
         // simulate three activity
@@ -74,19 +77,15 @@
                 "belowClose");
         belowClose.mActivityRecord.commitVisibility(
                 false /* visible */, true /* performLayout */);
-        mActivitySnapshotController.handleOpenActivityTransition(transitionState);
-        assertEquals(1, mActivitySnapshotController.mPendingCaptureActivity.size());
+        windows.add(belowClose.mActivityRecord);
+        mActivitySnapshotController.handleTransitionFinish(windows);
         assertEquals(1, mActivitySnapshotController.mPendingRemoveActivity.size());
-        assertEquals(closingWindow.mActivityRecord,
-                mActivitySnapshotController.mPendingCaptureActivity.valueAt(0));
         assertEquals(belowClose.mActivityRecord,
                 mActivitySnapshotController.mPendingRemoveActivity.valueAt(0));
     }
 
     @Test
     public void testCloseActivityTransition() {
-        final SnapshotController.TransitionState transitionState =
-                new SnapshotController.TransitionState();
         final Task task = createTask(mDisplayContent);
         // note for createAppWindow: the new child is added at index 0
         final WindowState closingWindow = createAppWindow(task, ACTIVITY_TYPE_STANDARD,
@@ -97,10 +96,10 @@
                 ACTIVITY_TYPE_STANDARD, "openingWindow");
         openingWindow.mActivityRecord.commitVisibility(
                 true /* visible */, true /* performLayout */);
-        transitionState.addParticipant(closingWindow.mActivityRecord, false);
-        transitionState.addParticipant(openingWindow.mActivityRecord, true);
-        mActivitySnapshotController.handleCloseActivityTransition(transitionState);
-        assertEquals(0, mActivitySnapshotController.mPendingCaptureActivity.size());
+        final ArrayList<WindowContainer> windows = new ArrayList<>();
+        windows.add(openingWindow.mActivityRecord);
+        windows.add(closingWindow.mActivityRecord);
+        mActivitySnapshotController.handleTransitionFinish(windows);
         assertEquals(1, mActivitySnapshotController.mPendingDeleteActivity.size());
         assertEquals(openingWindow.mActivityRecord,
                 mActivitySnapshotController.mPendingDeleteActivity.valueAt(0));
@@ -111,8 +110,8 @@
                 "belowOpen");
         belowOpen.mActivityRecord.commitVisibility(
                 false /* visible */, true /* performLayout */);
-        mActivitySnapshotController.handleCloseActivityTransition(transitionState);
-        assertEquals(0, mActivitySnapshotController.mPendingCaptureActivity.size());
+        windows.add(belowOpen.mActivityRecord);
+        mActivitySnapshotController.handleTransitionFinish(windows);
         assertEquals(1, mActivitySnapshotController.mPendingDeleteActivity.size());
         assertEquals(1, mActivitySnapshotController.mPendingLoadActivity.size());
         assertEquals(openingWindow.mActivityRecord,
@@ -123,10 +122,6 @@
 
     @Test
     public void testTaskTransition() {
-        final SnapshotController.TransitionState taskCloseTransition =
-                new SnapshotController.TransitionState();
-        final SnapshotController.TransitionState taskOpenTransition =
-                new SnapshotController.TransitionState();
         final Task closeTask = createTask(mDisplayContent);
         // note for createAppWindow: the new child is added at index 0
         final WindowState closingWindow = createAppWindow(closeTask, ACTIVITY_TYPE_STANDARD,
@@ -147,10 +142,10 @@
                 "openingWindowBelow");
         openingWindowBelow.mActivityRecord.commitVisibility(
                 false /* visible */, true /* performLayout */);
-        taskCloseTransition.addParticipant(closeTask, false);
-        taskOpenTransition.addParticipant(openTask, true);
-        mActivitySnapshotController.handleCloseTaskTransition(taskCloseTransition);
-        mActivitySnapshotController.handleOpenTaskTransition(taskOpenTransition);
+        final ArrayList<WindowContainer> windows = new ArrayList<>();
+        windows.add(closeTask);
+        windows.add(openTask);
+        mActivitySnapshotController.handleTransitionFinish(windows);
 
         assertEquals(1, mActivitySnapshotController.mPendingRemoveActivity.size());
         assertEquals(closingWindowBelow.mActivityRecord,
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppSnapshotControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/AppSnapshotControllerTests.java
deleted file mode 100644
index 83af1814..0000000
--- a/services/tests/wmtests/src/com/android/server/wm/AppSnapshotControllerTests.java
+++ /dev/null
@@ -1,171 +0,0 @@
-/*
- * Copyright (C) 2022 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.server.wm;
-
-import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
-
-import static com.android.server.wm.SnapshotController.ACTIVITY_CLOSE;
-import static com.android.server.wm.SnapshotController.ACTIVITY_OPEN;
-import static com.android.server.wm.SnapshotController.TASK_CLOSE;
-import static com.android.server.wm.SnapshotController.TASK_OPEN;
-
-import static org.junit.Assert.assertTrue;
-
-import android.platform.test.annotations.Presubmit;
-import android.util.ArraySet;
-
-import androidx.test.filters.SmallTest;
-
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-
-
-/**
- * Test class for {@link SnapshotController}.
- *
- * Build/Install/Run:
- *  *  atest WmTests:AppSnapshotControllerTests
- */
-@SmallTest
-@Presubmit
-@RunWith(WindowTestRunner.class)
-public class AppSnapshotControllerTests extends WindowTestsBase {
-    final ArraySet<ActivityRecord> mClosingApps = new ArraySet<>();
-    final ArraySet<ActivityRecord> mOpeningApps = new ArraySet<>();
-
-    final TransitionMonitor mOpenActivityMonitor = new TransitionMonitor();
-    final TransitionMonitor mCloseActivityMonitor = new TransitionMonitor();
-    final TransitionMonitor mOpenTaskMonitor = new TransitionMonitor();
-    final TransitionMonitor mCloseTaskMonitor = new TransitionMonitor();
-
-    @Before
-    public void setUp() throws Exception {
-        resetStatus();
-        mWm.mSnapshotController.registerTransitionStateConsumer(
-                ACTIVITY_CLOSE, mCloseActivityMonitor::handleTransition);
-        mWm.mSnapshotController.registerTransitionStateConsumer(
-                ACTIVITY_OPEN, mOpenActivityMonitor::handleTransition);
-        mWm.mSnapshotController.registerTransitionStateConsumer(
-                TASK_CLOSE, mCloseTaskMonitor::handleTransition);
-        mWm.mSnapshotController.registerTransitionStateConsumer(
-                TASK_OPEN, mOpenTaskMonitor::handleTransition);
-    }
-
-    @After
-    public void tearDown() throws Exception {
-        mWm.mSnapshotController.unregisterTransitionStateConsumer(
-                ACTIVITY_CLOSE, mCloseActivityMonitor::handleTransition);
-        mWm.mSnapshotController.unregisterTransitionStateConsumer(
-                ACTIVITY_OPEN, mOpenActivityMonitor::handleTransition);
-        mWm.mSnapshotController.unregisterTransitionStateConsumer(
-                TASK_CLOSE, mCloseTaskMonitor::handleTransition);
-        mWm.mSnapshotController.unregisterTransitionStateConsumer(
-                TASK_OPEN, mOpenTaskMonitor::handleTransition);
-    }
-
-    private static class TransitionMonitor {
-        private final ArraySet<WindowContainer> mOpenParticipant = new ArraySet<>();
-        private final ArraySet<WindowContainer> mCloseParticipant = new ArraySet<>();
-        void handleTransition(SnapshotController.TransitionState<ActivityRecord> state) {
-            mOpenParticipant.addAll(state.getParticipant(true /* open */));
-            mCloseParticipant.addAll(state.getParticipant(false /* close */));
-        }
-        void reset() {
-            mOpenParticipant.clear();
-            mCloseParticipant.clear();
-        }
-    }
-
-    private void resetStatus() {
-        mClosingApps.clear();
-        mOpeningApps.clear();
-        mOpenActivityMonitor.reset();
-        mCloseActivityMonitor.reset();
-        mOpenTaskMonitor.reset();
-        mCloseTaskMonitor.reset();
-    }
-
-    @Test
-    public void testHandleAppTransition_openActivityTransition() {
-        final Task task = createTask(mDisplayContent);
-        // note for createAppWindow: the new child is added at index 0
-        final WindowState openingWindow = createAppWindow(task,
-                ACTIVITY_TYPE_STANDARD, "openingWindow");
-        openingWindow.mActivityRecord.commitVisibility(
-                true /* visible */, true /* performLayout */);
-        final WindowState closingWindow = createAppWindow(task, ACTIVITY_TYPE_STANDARD,
-                "closingWindow");
-        closingWindow.mActivityRecord.commitVisibility(
-                false /* visible */, true /* performLayout */);
-        mClosingApps.add(closingWindow.mActivityRecord);
-        mOpeningApps.add(openingWindow.mActivityRecord);
-        mWm.mSnapshotController.handleAppTransition(mClosingApps, mOpeningApps);
-        assertTrue(mOpenActivityMonitor.mCloseParticipant.contains(closingWindow.mActivityRecord));
-        assertTrue(mOpenActivityMonitor.mOpenParticipant.contains(openingWindow.mActivityRecord));
-    }
-
-    @Test
-    public void testHandleAppTransition_closeActivityTransition() {
-        final Task task = createTask(mDisplayContent);
-        // note for createAppWindow: the new child is added at index 0
-        final WindowState closingWindow = createAppWindow(task, ACTIVITY_TYPE_STANDARD,
-                "closingWindow");
-        closingWindow.mActivityRecord.commitVisibility(
-                false /* visible */, true /* performLayout */);
-        final WindowState openingWindow = createAppWindow(task,
-                ACTIVITY_TYPE_STANDARD, "openingWindow");
-        openingWindow.mActivityRecord.commitVisibility(
-                true /* visible */, true /* performLayout */);
-        mClosingApps.add(closingWindow.mActivityRecord);
-        mOpeningApps.add(openingWindow.mActivityRecord);
-        mWm.mSnapshotController.handleAppTransition(mClosingApps, mOpeningApps);
-        assertTrue(mCloseActivityMonitor.mCloseParticipant.contains(closingWindow.mActivityRecord));
-        assertTrue(mCloseActivityMonitor.mOpenParticipant.contains(openingWindow.mActivityRecord));
-    }
-
-    @Test
-    public void testHandleAppTransition_TaskTransition() {
-        final Task closeTask = createTask(mDisplayContent);
-        // note for createAppWindow: the new child is added at index 0
-        final WindowState closingWindow = createAppWindow(closeTask, ACTIVITY_TYPE_STANDARD,
-                "closingWindow");
-        closingWindow.mActivityRecord.commitVisibility(
-                false /* visible */, true /* performLayout */);
-        final WindowState closingWindowBelow = createAppWindow(closeTask, ACTIVITY_TYPE_STANDARD,
-                "closingWindowBelow");
-        closingWindowBelow.mActivityRecord.commitVisibility(
-                false /* visible */, true /* performLayout */);
-
-        final Task openTask = createTask(mDisplayContent);
-        final WindowState openingWindow = createAppWindow(openTask, ACTIVITY_TYPE_STANDARD,
-                "openingWindow");
-        openingWindow.mActivityRecord.commitVisibility(
-                true /* visible */, true /* performLayout */);
-        final WindowState openingWindowBelow = createAppWindow(openTask, ACTIVITY_TYPE_STANDARD,
-                "openingWindowBelow");
-        openingWindowBelow.mActivityRecord.commitVisibility(
-                false /* visible */, true /* performLayout */);
-
-        mClosingApps.add(closingWindow.mActivityRecord);
-        mOpeningApps.add(openingWindow.mActivityRecord);
-        mWm.mSnapshotController.handleAppTransition(mClosingApps, mOpeningApps);
-        assertTrue(mCloseTaskMonitor.mCloseParticipant.contains(closeTask));
-        assertTrue(mOpenTaskMonitor.mOpenParticipant.contains(openTask));
-    }
-}
diff --git a/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java b/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
index 40b1521..07cdfaf 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
@@ -46,7 +46,6 @@
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
-import static com.android.server.wm.SnapshotController.TASK_CLOSE;
 import static com.android.server.wm.WindowContainer.POSITION_TOP;
 
 import static org.junit.Assert.assertEquals;
@@ -1387,8 +1386,6 @@
     @Test
     public void testTransientLaunch() {
         spyOn(mWm.mSnapshotController.mTaskSnapshotController);
-        mWm.mSnapshotController.registerTransitionStateConsumer(TASK_CLOSE,
-                mWm.mSnapshotController.mTaskSnapshotController::handleTaskClose);
         final ArrayList<ActivityRecord> enteringAnimReports = new ArrayList<>();
         final TransitionController controller = new TestTransitionController(mAtm) {
             @Override