Support to record multiple activities in the same task as one snapshot.

So the snapshot can be easier to be used with current SnapshotDrawer,
also the overall structure would be much similar with associated task
starting window.

- Record a single activity snapshot while multiple activities in the
  same leaf task are closing together, this can be extention if mulitple
  adjacent will happen.
- Share the snapshot from activities: because the key of cache must be
  an activity, so each activity can be the key to find this snapshot
  from cache, but the activities must be exactly match when calling
  getSnapshot.
- Remove the snapshot if one of activity removed, usually the
  configuration of rest activities will be changed, so doesn't need to
  keep the snapshot.

Bug: 259497289
Bug: 274997067
Test: atest ActivitySnapshotControllerTests
Test: test on sample app, monitor the snapshot of adjacent activities
should work like single activity snapshot.

Change-Id: I357010ab79eda165564e36fe7203fc508dbc357b

7 files changed, 413 insertions, 108 deletions

diff --git a/services/core/java/com/android/server/wm/AbsAppSnapshotController.java b/services/core/java/com/android/server/wm/AbsAppSnapshotController.java
index 05da9dfe7921..e5c743cc69e4 100644
--- a/services/core/java/com/android/server/wm/AbsAppSnapshotController.java
+++ b/services/core/java/com/android/server/wm/AbsAppSnapshotController.java
@@ -18,6 +18,7 @@ package com.android.server.wm;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_DREAM;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_RECENTS;
+
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_SCREENSHOT;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
@@ -180,16 +181,8 @@ abstract class AbsAppSnapshotController<TYPE extends WindowContainer,
         if (snapshot == null) {
             return null;
         }
-        final HardwareBuffer buffer = snapshot.getHardwareBuffer();
-        if (buffer.getWidth() == 0 || buffer.getHeight() == 0) {
-            buffer.close();
-            Slog.e(TAG, "Invalid snapshot dimensions " + buffer.getWidth() + "x"
-                    + buffer.getHeight());
-            return null;
-        } else {
-            mCache.putSnapshot(source, snapshot);
-            return snapshot;
-        }
+        mCache.putSnapshot(source, snapshot);
+        return snapshot;
     }
 
     @VisibleForTesting
@@ -210,6 +203,11 @@ abstract class AbsAppSnapshotController<TYPE extends WindowContainer,
 
     @Nullable
     TaskSnapshot snapshot(TYPE source) {
+        return snapshot(source, mHighResSnapshotScale);
+    }
+
+    @Nullable
+    TaskSnapshot snapshot(TYPE source, float scale) {
         TaskSnapshot.Builder builder = new TaskSnapshot.Builder();
         final Rect crop = prepareTaskSnapshot(source, builder);
         if (crop == null) {
@@ -218,7 +216,7 @@ abstract class AbsAppSnapshotController<TYPE extends WindowContainer,
         }
         Trace.traceBegin(Trace.TRACE_TAG_WINDOW_MANAGER, "createSnapshot");
         final ScreenCapture.ScreenshotHardwareBuffer screenshotBuffer = createSnapshot(source,
-                mHighResSnapshotScale, crop, builder);
+                scale, crop, builder);
         Trace.traceEnd(Trace.TRACE_TAG_WINDOW_MANAGER);
         if (screenshotBuffer == null) {
             // Failed to acquire image. Has been logged.
@@ -227,7 +225,19 @@ abstract class AbsAppSnapshotController<TYPE extends WindowContainer,
         builder.setCaptureTime(SystemClock.elapsedRealtimeNanos());
         builder.setSnapshot(screenshotBuffer.getHardwareBuffer());
         builder.setColorSpace(screenshotBuffer.getColorSpace());
-        return builder.build();
+        final TaskSnapshot snapshot = builder.build();
+        return validateSnapshot(snapshot);
+    }
+
+    private static TaskSnapshot validateSnapshot(@NonNull TaskSnapshot snapshot) {
+        final HardwareBuffer buffer = snapshot.getHardwareBuffer();
+        if (buffer.getWidth() == 0 || buffer.getHeight() == 0) {
+            buffer.close();
+            Slog.e(TAG, "Invalid snapshot dimensions " + buffer.getWidth() + "x"
+                    + buffer.getHeight());
+            return null;
+        }
+        return snapshot;
     }
 
     @Nullable
@@ -432,7 +442,7 @@ abstract class AbsAppSnapshotController<TYPE extends WindowContainer,
         InsetUtils.addInsets(contentInsets, letterboxInsets);
         // Note, the app theme snapshot is never translucent because we enforce a non-translucent
         // color above
-        return new TaskSnapshot(
+        final TaskSnapshot taskSnapshot = new TaskSnapshot(
                 System.currentTimeMillis() /* id */,
                 SystemClock.elapsedRealtimeNanos() /* captureTime */,
                 topActivity.mActivityComponent, hwBitmap.getHardwareBuffer(),
@@ -441,6 +451,7 @@ abstract class AbsAppSnapshotController<TYPE extends WindowContainer,
                 contentInsets, letterboxInsets, false /* isLowResolution */,
                 false /* isRealSnapshot */, source.getWindowingMode(),
                 getAppearance(source), false /* isTranslucent */, false /* hasImeSurface */);
+        return validateSnapshot(taskSnapshot);
     }
 
     static Rect getSystemBarInsets(Rect frame, InsetsState state) {
diff --git a/services/core/java/com/android/server/wm/ActivitySnapshotController.java b/services/core/java/com/android/server/wm/ActivitySnapshotController.java
index 7af494c296de..a692167bbbf9 100644
--- a/services/core/java/com/android/server/wm/ActivitySnapshotController.java
+++ b/services/core/java/com/android/server/wm/ActivitySnapshotController.java
@@ -25,6 +25,7 @@ import android.os.Environment;
 import android.os.SystemProperties;
 import android.os.Trace;
 import android.util.ArraySet;
+import android.util.IntArray;
 import android.util.Slog;
 import android.util.SparseArray;
 import android.window.TaskSnapshot;
@@ -36,6 +37,7 @@ import com.android.server.wm.BaseAppSnapshotPersister.PersistInfoProvider;
 import com.android.window.flags.Flags;
 
 import java.io.File;
+import java.io.PrintWriter;
 import java.util.ArrayList;
 
 /**
@@ -136,11 +138,26 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
                 false /* enableLowResSnapshots */, 0 /* lowResScaleFactor */, use16BitFormat);
     }
 
-    /** Retrieves a snapshot for an activity from cache. */
+    /**
+     * Retrieves a snapshot for a set of activities from cache.
+     * This will only return the snapshot IFF input activities exist entirely in the snapshot.
+     * Sample: If the snapshot was captured with activity A and B, here will return null if the
+     * input activity is only [A] or [B], it must be [A, B]
+     */
     @Nullable
-    TaskSnapshot getSnapshot(ActivityRecord ar) {
-        final int code = getSystemHashCode(ar);
-        return mCache.getSnapshot(code);
+    TaskSnapshot getSnapshot(@NonNull ActivityRecord[] activities) {
+        if (activities.length == 0) {
+            return null;
+        }
+        final UserSavedFile tmpUsf = findSavedFile(activities[0]);
+        if (tmpUsf == null || tmpUsf.mActivityIds.size() != activities.length) {
+            return null;
+        }
+        int fileId = 0;
+        for (int i = activities.length - 1; i >= 0; --i) {
+            fileId ^= getSystemHashCode(activities[i]);
+        }
+        return tmpUsf.mFileId == fileId ? mCache.getSnapshot(tmpUsf.mActivityIds.get(0)) : null;
     }
 
     private void cleanUpUserFiles(int userId) {
@@ -229,33 +246,16 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
                     + " load " + mPendingLoadActivity);
         }
         // load snapshot to cache
-        for (int i = mPendingLoadActivity.size() - 1; i >= 0; i--) {
-            final ActivityRecord ar = mPendingLoadActivity.valueAt(i);
-            final int code = getSystemHashCode(ar);
-            final int userId = ar.mUserId;
-            if (mCache.getSnapshot(code) != null) {
-                // already in cache, skip
-                continue;
-            }
-            if (containsFile(code, userId)) {
-                synchronized (mSnapshotPersistQueue.getLock()) {
-                    mSnapshotPersistQueue.insertQueueAtFirstLocked(
-                            new LoadActivitySnapshotItem(ar, code, userId, mPersistInfoProvider));
-                }
-            }
-        }
+        loadActivitySnapshot();
         // 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);
+            removeCachedFiles(ar);
         }
         // 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);
+            removeIfUserSavedFileExist(ar);
         }
         // don't keep any reference
         resetTmpFields();
@@ -264,28 +264,38 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
     class LoadActivitySnapshotItem extends SnapshotPersistQueue.WriteQueueItem {
         private final int mCode;
         private final int mUserId;
-        private final ActivityRecord mActivityRecord;
+        private final ActivityRecord[] mActivities;
 
-        LoadActivitySnapshotItem(@NonNull ActivityRecord ar, int code, int userId,
+        LoadActivitySnapshotItem(@NonNull ActivityRecord[] activities, int code, int userId,
                 @NonNull PersistInfoProvider persistInfoProvider) {
             super(persistInfoProvider);
-            mActivityRecord = ar;
+            mActivities = activities;
             mCode = code;
             mUserId = userId;
         }
 
         @Override
         void write() {
-            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER,
-                    "load_activity_snapshot");
-            final TaskSnapshot snapshot = mSnapshotLoader.loadTask(mCode,
-                    mUserId, false /* loadLowResolutionBitmap */);
-            synchronized (mService.getWindowManagerLock()) {
-                if (snapshot != null && !mActivityRecord.finishing) {
-                    mCache.putSnapshot(mActivityRecord, snapshot);
+            try {
+                Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER,
+                        "load_activity_snapshot");
+                final TaskSnapshot snapshot = mSnapshotLoader.loadTask(mCode,
+                        mUserId, false /* loadLowResolutionBitmap */);
+                if (snapshot == null) {
+                    return;
+                }
+                synchronized (mService.getWindowManagerLock()) {
+                    // Verify the snapshot is still needed, and the activity is not finishing
+                    if (!hasRecord(mActivities[0])) {
+                        return;
+                    }
+                    for (ActivityRecord ar : mActivities) {
+                        mCache.putSnapshot(ar, snapshot);
+                    }
                 }
+            } finally {
+                Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
             }
-            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
         }
 
         @Override
@@ -297,18 +307,81 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
         }
     }
 
-    void recordSnapshot(ActivityRecord activity) {
-        if (shouldDisableSnapshots()) {
+    void loadActivitySnapshot() {
+        if (mPendingLoadActivity.isEmpty()) {
+            return;
+        }
+        // Only load if saved file exists.
+        final ArraySet<UserSavedFile> loadingFiles = new ArraySet<>();
+        for (int i = mPendingLoadActivity.size() - 1; i >= 0; i--) {
+            final ActivityRecord ar = mPendingLoadActivity.valueAt(i);
+            final UserSavedFile usf = findSavedFile(ar);
+            if (usf != null) {
+                loadingFiles.add(usf);
+            }
+        }
+        // Filter out the activity if the snapshot was removed.
+        for (int i = loadingFiles.size() - 1; i >= 0; i--) {
+            final UserSavedFile usf = loadingFiles.valueAt(i);
+            final ActivityRecord[] activities = usf.filterExistActivities(mPendingLoadActivity);
+            if (activities == null) {
+                continue;
+            }
+            if (getSnapshot(activities) != null) {
+                // Found the cache in memory, so skip loading from file.
+                continue;
+            }
+            loadSnapshotInner(activities, usf);
+        }
+    }
+
+    @VisibleForTesting
+    void loadSnapshotInner(ActivityRecord[] activities, UserSavedFile usf) {
+        synchronized (mSnapshotPersistQueue.getLock()) {
+            mSnapshotPersistQueue.insertQueueAtFirstLocked(new LoadActivitySnapshotItem(
+                    activities, usf.mFileId, usf.mUserId, mPersistInfoProvider));
+        }
+    }
+
+    /**
+     * Record one or multiple activities within a snapshot where those activities must belong to
+     * the same task.
+     * @param activity If the request activity is more than one, try to record those activities
+     *                 as a single snapshot, so those activities should belong to the same task.
+     */
+    void recordSnapshot(@NonNull ArrayList<ActivityRecord> activity) {
+        if (shouldDisableSnapshots() || activity.isEmpty()) {
             return;
         }
         if (DEBUG) {
             Slog.d(TAG, "ActivitySnapshotController#recordSnapshot " + activity);
         }
-        final TaskSnapshot snapshot = recordSnapshotInner(activity);
-        if (snapshot != null) {
-            final int code = getSystemHashCode(activity);
-            addUserSavedFile(code, activity.mUserId, snapshot);
+        final int size = activity.size();
+        final int[] mixedCode = new int[size];
+        if (size == 1) {
+            final ActivityRecord singleActivity = activity.get(0);
+            final TaskSnapshot snapshot = recordSnapshotInner(singleActivity);
+            if (snapshot != null) {
+                mixedCode[0] = getSystemHashCode(singleActivity);
+                addUserSavedFile(singleActivity.mUserId, snapshot, mixedCode);
+            }
+            return;
         }
+
+        final Task mainTask = activity.get(0).getTask();
+        // Snapshot by task controller with activity's scale.
+        final TaskSnapshot snapshot = mService.mTaskSnapshotController
+                .snapshot(mainTask, mHighResSnapshotScale);
+        if (snapshot == null) {
+            return;
+        }
+
+        for (int i = 0; i < activity.size(); ++i) {
+            final ActivityRecord next = activity.get(i);
+            mCache.putSnapshot(next, snapshot);
+            mixedCode[i] = getSystemHashCode(next);
+        }
+        addUserSavedFile(mainTask.mUserId, snapshot, mixedCode);
     }
 
     /**
@@ -331,7 +404,8 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
         }
     }
 
-    private static int getSystemHashCode(ActivityRecord activity) {
+    @VisibleForTesting
+    static int getSystemHashCode(ActivityRecord activity) {
         return System.identityHashCode(activity);
     }
 
@@ -362,7 +436,13 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
         if (ar.isVisibleRequested()) {
             mPendingDeleteActivity.add(ar);
             // load next one if exists.
-            addBelowActivityIfExist(ar, mPendingLoadActivity, true, "load-snapshot");
+            // Note if this transition is happen between two TaskFragment, the next N - 1 activity
+            // may not participant in this transition.
+            // Sample:
+            //   [TF1] close
+            //   [TF2] open
+            //   Bottom Activity <- Able to load this even it didn't participant the transition.
+            addBelowActivityIfExist(ar, mPendingLoadActivity, false, "load-snapshot");
         } else {
             // remove the snapshot for the one below close
             addBelowActivityIfExist(ar, mPendingRemoveActivity, true, "remove-snapshot");
@@ -478,10 +558,8 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
     }
 
     private void adjustSavedFileOrder(Task nextTopTask) {
-        final int userId = nextTopTask.mUserId;
         nextTopTask.forAllActivities(ar -> {
-            final int code = getSystemHashCode(ar);
-            final UserSavedFile usf = getUserFiles(userId).get(code);
+            final UserSavedFile usf = findSavedFile(ar);
             if (usf != null) {
                 mSavedFilesInOrder.remove(usf);
                 mSavedFilesInOrder.add(usf);
@@ -494,9 +572,7 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
         if (shouldDisableSnapshots()) {
             return;
         }
-        super.onAppRemoved(activity);
-        final int code = getSystemHashCode(activity);
-        removeIfUserSavedFileExist(code, activity.mUserId);
+        removeIfUserSavedFileExist(activity);
         if (DEBUG) {
             Slog.d(TAG, "ActivitySnapshotController#onAppRemoved delete snapshot " + activity);
         }
@@ -507,9 +583,7 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
         if (shouldDisableSnapshots()) {
             return;
         }
-        super.onAppDied(activity);
-        final int code = getSystemHashCode(activity);
-        removeIfUserSavedFileExist(code, activity.mUserId);
+        removeIfUserSavedFileExist(activity);
         if (DEBUG) {
             Slog.d(TAG, "ActivitySnapshotController#onAppDied delete snapshot " + activity);
         }
@@ -558,55 +632,92 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
         return mUserSavedFiles.get(userId);
     }
 
-    private void removeIfUserSavedFileExist(int code, int userId) {
-        final UserSavedFile usf = getUserFiles(userId).get(code);
+    UserSavedFile findSavedFile(@NonNull ActivityRecord ar) {
+        final int code = getSystemHashCode(ar);
+        return findSavedFile(ar.mUserId, code);
+    }
+
+    UserSavedFile findSavedFile(int userId, int code) {
+        final SparseArray<UserSavedFile> usfs = getUserFiles(userId);
+        return usfs.get(code);
+    }
+
+    private void removeCachedFiles(ActivityRecord ar) {
+        final UserSavedFile usf = findSavedFile(ar);
+        if (usf != null) {
+            for (int i = usf.mActivityIds.size() - 1; i >= 0; --i) {
+                final int activityId = usf.mActivityIds.get(i);
+                mCache.onIdRemoved(activityId);
+            }
+        }
+    }
+
+    private void removeIfUserSavedFileExist(ActivityRecord ar) {
+        final UserSavedFile usf = findSavedFile(ar);
         if (usf != null) {
-            mUserSavedFiles.get(userId).remove(code);
+            final SparseArray<UserSavedFile> usfs = getUserFiles(ar.mUserId);
+            for (int i = usf.mActivityIds.size() - 1; i >= 0; --i) {
+                final int activityId = usf.mActivityIds.get(i);
+                usf.remove(activityId);
+                mCache.onIdRemoved(activityId);
+                usfs.remove(activityId);
+            }
             mSavedFilesInOrder.remove(usf);
-            mPersister.removeSnapshot(code, userId);
+            mPersister.removeSnapshot(usf.mFileId, ar.mUserId);
         }
     }
 
-    private boolean containsFile(int code, int userId) {
-        return getUserFiles(userId).get(code) != null;
+    @VisibleForTesting
+    boolean hasRecord(@NonNull ActivityRecord ar) {
+        return findSavedFile(ar) != null;
     }
 
-    private void addUserSavedFile(int code, int userId, TaskSnapshot snapshot) {
-        final SparseArray<UserSavedFile> savedFiles = getUserFiles(userId);
-        final UserSavedFile savedFile = savedFiles.get(code);
-        if (savedFile == null) {
-            final UserSavedFile usf = new UserSavedFile(code, userId);
-            savedFiles.put(code, usf);
-            mSavedFilesInOrder.add(usf);
-            mPersister.persistSnapshot(code, userId, snapshot);
+    @VisibleForTesting
+    void addUserSavedFile(int userId, TaskSnapshot snapshot, @NonNull int[] code) {
+        final UserSavedFile savedFile = findSavedFile(userId, code[0]);
+        if (savedFile != null) {
+            Slog.w(TAG, "Duplicate request for recording activity snapshot " + savedFile);
+            return;
+        }
+        int fileId = 0;
+        for (int i = code.length - 1; i >= 0; --i) {
+            fileId ^= code[i];
+        }
+        final UserSavedFile usf = new UserSavedFile(fileId, userId);
+        SparseArray<UserSavedFile> usfs = getUserFiles(userId);
+        for (int i = code.length - 1; i >= 0; --i) {
+            usfs.put(code[i], usf);
+        }
+        usf.mActivityIds.addAll(code);
+        mSavedFilesInOrder.add(usf);
+        mPersister.persistSnapshot(fileId, userId, snapshot);
 
-            if (mSavedFilesInOrder.size() > MAX_PERSIST_SNAPSHOT_COUNT * 2) {
-                purgeSavedFile();
-            }
+        if (mSavedFilesInOrder.size() > MAX_PERSIST_SNAPSHOT_COUNT * 2) {
+            purgeSavedFile();
         }
     }
 
     private void purgeSavedFile() {
         final int savedFileCount = mSavedFilesInOrder.size();
         final int removeCount = savedFileCount - MAX_PERSIST_SNAPSHOT_COUNT;
-        final ArrayList<UserSavedFile> usfs = new ArrayList<>();
-        if (removeCount > 0) {
-            final int removeTillIndex = savedFileCount - removeCount;
-            for (int i = savedFileCount - 1; i > removeTillIndex; --i) {
-                final UserSavedFile usf = mSavedFilesInOrder.remove(i);
-                if (usf != null) {
-                    final SparseArray<UserSavedFile> records = getUserFiles(usf.mUserId);
-                    records.remove(usf.mFileId);
-                    usfs.add(usf);
-                }
-            }
+        if (removeCount < 1) {
+            return;
         }
-        if (usfs.size() > 0) {
-            removeSnapshotFiles(usfs);
+
+        final ArrayList<UserSavedFile> removeTargets = new ArrayList<>();
+        for (int i = removeCount - 1; i >= 0; --i) {
+            final UserSavedFile usf = mSavedFilesInOrder.remove(i);
+            final SparseArray<UserSavedFile> files = mUserSavedFiles.get(usf.mUserId);
+            for (int j = usf.mActivityIds.size() - 1; j >= 0; --j) {
+                mCache.removeRunningEntry(usf.mActivityIds.get(j));
+                files.remove(usf.mActivityIds.get(j));
+            }
+            removeTargets.add(usf);
         }
+        removeSnapshotFiles(removeTargets);
     }
 
-    private void removeSnapshotFiles(ArrayList<UserSavedFile> files) {
+    private void removeSnapshotFiles(@NonNull ArrayList<UserSavedFile> files) {
         synchronized (mSnapshotPersistQueue.getLock()) {
             mSnapshotPersistQueue.sendToQueueLocked(
                     new SnapshotPersistQueue.WriteQueueItem(mPersistInfoProvider) {
@@ -624,12 +735,85 @@ class ActivitySnapshotController extends AbsAppSnapshotController<ActivityRecord
         }
     }
 
+    @Override
+    void dump(PrintWriter pw, String prefix) {
+        super.dump(pw, prefix);
+        final String doublePrefix = prefix + "  ";
+        final String triplePrefix = doublePrefix + "  ";
+        for (int i = mUserSavedFiles.size() - 1; i >= 0; --i) {
+            final SparseArray<UserSavedFile> usfs = mUserSavedFiles.valueAt(i);
+            pw.println(doublePrefix + "UserSavedFile userId=" + mUserSavedFiles.keyAt(i));
+            final ArraySet<UserSavedFile> sets = new ArraySet<>();
+            for (int j = usfs.size() - 1; j >= 0; --j) {
+                sets.add(usfs.valueAt(j));
+            }
+            for (int j = sets.size() - 1; j >= 0; --j) {
+                pw.println(triplePrefix + "SavedFile=" + sets.valueAt(j));
+            }
+        }
+    }
+
     static class UserSavedFile {
-        int mFileId;
-        int mUserId;
+        // The unique id as filename.
+        final int mFileId;
+        final int mUserId;
+
+        /**
+         * The Id of all activities which are includes in the snapshot.
+         */
+        final IntArray mActivityIds = new IntArray();
+
         UserSavedFile(int fileId, int userId) {
             mFileId = fileId;
             mUserId = userId;
         }
+
+        boolean contains(int code) {
+            return mActivityIds.contains(code);
+        }
+
+        void remove(int code) {
+            final int index = mActivityIds.indexOf(code);
+            if (index >= 0) {
+                mActivityIds.remove(index);
+            }
+        }
+
+        ActivityRecord[] filterExistActivities(
+                @NonNull ArraySet<ActivityRecord> pendingLoadActivity) {
+            ArrayList<ActivityRecord> matchedActivities = null;
+            for (int i = pendingLoadActivity.size() - 1; i >= 0; --i) {
+                final ActivityRecord ar = pendingLoadActivity.valueAt(i);
+                if (contains(getSystemHashCode(ar))) {
+                    if (matchedActivities == null) {
+                        matchedActivities = new ArrayList<>();
+                    }
+                    matchedActivities.add(ar);
+                }
+            }
+            if (matchedActivities == null || matchedActivities.size() != mActivityIds.size()) {
+                return null;
+            }
+            return matchedActivities.toArray(new ActivityRecord[0]);
+        }
+
+        @Override
+        public String toString() {
+            StringBuilder sb = new StringBuilder(128);
+            sb.append("UserSavedFile{");
+            sb.append(Integer.toHexString(System.identityHashCode(this)));
+            sb.append(" fileId=");
+            sb.append(Integer.toHexString(mFileId));
+            sb.append(", activityIds=[");
+            for (int i = mActivityIds.size() - 1; i >= 0; --i) {
+                sb.append(Integer.toHexString(mActivityIds.get(i)));
+                if (i > 0) {
+                    sb.append(',');
+                }
+            }
+            sb.append("]");
+            sb.append("}");
+            return sb.toString();
+        }
     }
 }
diff --git a/services/core/java/com/android/server/wm/BackNavigationController.java b/services/core/java/com/android/server/wm/BackNavigationController.java
index 22d17b596c4c..b94206dd700a 100644
--- a/services/core/java/com/android/server/wm/BackNavigationController.java
+++ b/services/core/java/com/android/server/wm/BackNavigationController.java
@@ -1327,7 +1327,8 @@ class BackNavigationController {
                 return mAnimationTarget;
             }
 
-            void createStartingSurface(@NonNull WindowContainer closeWindow) {
+            void createStartingSurface(@NonNull WindowContainer closeWindow,
+                    @NonNull ActivityRecord[] visibleOpenActivities) {
                 if (!mIsOpen) {
                     return;
                 }
@@ -1346,7 +1347,7 @@ class BackNavigationController {
                 if (mainActivity == null) {
                     return;
                 }
-                final TaskSnapshot snapshot = getSnapshot(mTarget);
+                final TaskSnapshot snapshot = getSnapshot(mTarget, visibleOpenActivities);
                 mRequestedStartingSurfaceId = openTask.mAtmService.mTaskOrganizerController
                         .addWindowlessStartingSurface(openTask, mainActivity,
                                 // Choose configuration from closeWindow, because the configuration
@@ -1489,7 +1490,8 @@ class BackNavigationController {
                         // Try to draw two snapshot within a WindowlessStartingWindow, or find
                         // another key for StartingWindowRecordManager.
                         && openAnimationAdaptor.length == 1) {
-                    openAnimationAdaptor[0].createStartingSurface(closeWindow);
+                    openAnimationAdaptor[0].createStartingSurface(closeWindow,
+                            visibleOpenActivities);
                 } else {
                     for (int i = visibleOpenActivities.length - 1; i >= 0; --i) {
                         setLaunchBehind(visibleOpenActivities[i]);
@@ -1671,7 +1673,8 @@ class BackNavigationController {
         mPendingAnimationBuilder = null;
     }
 
-    static TaskSnapshot getSnapshot(@NonNull WindowContainer w) {
+    static TaskSnapshot getSnapshot(@NonNull WindowContainer w,
+            ActivityRecord[] visibleOpenActivities) {
         if (w.asTask() != null) {
             final Task task = w.asTask();
             return task.mRootWindowContainer.mWindowManager.mTaskSnapshotController.getSnapshot(
@@ -1681,7 +1684,8 @@ class BackNavigationController {
 
         if (w.asActivityRecord() != null) {
             final ActivityRecord ar = w.asActivityRecord();
-            return ar.mWmService.mSnapshotController.mActivitySnapshotController.getSnapshot(ar);
+            return ar.mWmService.mSnapshotController.mActivitySnapshotController
+                    .getSnapshot(visibleOpenActivities);
         }
         return null;
     }
diff --git a/services/core/java/com/android/server/wm/SnapshotController.java b/services/core/java/com/android/server/wm/SnapshotController.java
index b6f040a6cb56..3014f979aa70 100644
--- a/services/core/java/com/android/server/wm/SnapshotController.java
+++ b/services/core/java/com/android/server/wm/SnapshotController.java
@@ -160,9 +160,7 @@ class SnapshotController {
                 if (!allOpensOptInOnBackInvoked() || mCloseActivities.isEmpty()) {
                     return;
                 }
-                for (int i = mCloseActivities.size() - 1; i >= 0; --i) {
-                    controller.recordSnapshot(mCloseActivities.get(i));
-                }
+                controller.recordSnapshot(mCloseActivities);
             }
         }
     }
diff --git a/services/core/java/com/android/server/wm/SnapshotPersistQueue.java b/services/core/java/com/android/server/wm/SnapshotPersistQueue.java
index bffdf54e17ce..e4379b5343f3 100644
--- a/services/core/java/com/android/server/wm/SnapshotPersistQueue.java
+++ b/services/core/java/com/android/server/wm/SnapshotPersistQueue.java
@@ -274,7 +274,9 @@ class SnapshotPersistQueue {
 
         @Override
         void write() {
-            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "StoreWriteQueueItem");
+            if (Trace.isTagEnabled(TRACE_TAG_WINDOW_MANAGER)) {
+                Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "StoreWriteQueueItem#" + mId);
+            }
             if (!mPersistInfoProvider.createDirectory(mUserId)) {
                 Slog.e(TAG, "Unable to create snapshot directory for user dir="
                         + mPersistInfoProvider.getDirectory(mUserId));
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 98f18433e53d..03d30294e1d8 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivitySnapshotControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivitySnapshotControllerTests.java
@@ -17,13 +17,31 @@
 package com.android.server.wm;
 
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
+import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
+import static android.content.res.Configuration.ORIENTATION_PORTRAIT;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 
+import android.content.ComponentName;
+import android.graphics.ColorSpace;
+import android.graphics.Point;
+import android.graphics.Rect;
+import android.hardware.HardwareBuffer;
 import android.platform.test.annotations.Presubmit;
+import android.util.ArraySet;
+import android.view.Surface;
+import android.window.TaskSnapshot;
 
 import androidx.test.filters.SmallTest;
 
@@ -32,6 +50,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 
 /**
  * Test class for {@link ActivitySnapshotController}.
@@ -45,6 +64,7 @@ import java.util.ArrayList;
 public class ActivitySnapshotControllerTests extends WindowTestsBase {
 
     private ActivitySnapshotController mActivitySnapshotController;
+
     @Before
     public void setUp() throws Exception {
         spyOn(mWm.mSnapshotController.mActivitySnapshotController);
@@ -154,4 +174,90 @@ public class ActivitySnapshotControllerTests extends WindowTestsBase {
         assertEquals(openingWindowBelow.mActivityRecord,
                 mActivitySnapshotController.mPendingLoadActivity.valueAt(0));
     }
+
+    /**
+     * Simulate multiple TaskFragments inside a task.
+     */
+    @Test
+    public void testMultipleActivitiesLoadSnapshot() {
+        final Task testTask = createTask(mDisplayContent);
+        final ActivityRecord activityA = createActivityRecord(testTask);
+        final ActivityRecord activityB = createActivityRecord(testTask);
+        final ActivityRecord activityC = createActivityRecord(testTask);
+        final TaskSnapshot taskSnapshot = createSnapshot();
+
+        final int[] mixedCode = new int[3];
+        mixedCode[0] = ActivitySnapshotController.getSystemHashCode(activityA);
+        mixedCode[1] = ActivitySnapshotController.getSystemHashCode(activityB);
+        mixedCode[2] = ActivitySnapshotController.getSystemHashCode(activityC);
+
+        mActivitySnapshotController.addUserSavedFile(testTask.mUserId, taskSnapshot, mixedCode);
+        mActivitySnapshotController.mCache.putSnapshot(activityA, taskSnapshot);
+        mActivitySnapshotController.mCache.putSnapshot(activityB, taskSnapshot);
+        mActivitySnapshotController.mCache.putSnapshot(activityC, taskSnapshot);
+
+        assertTrue(mActivitySnapshotController.hasRecord(activityA));
+        assertTrue(mActivitySnapshotController.hasRecord(activityB));
+
+        // If A is removed, B and C should also be removed because they share the same snapshot.
+        mActivitySnapshotController.onAppRemoved(activityA);
+        assertFalse(mActivitySnapshotController.hasRecord(activityA));
+        assertFalse(mActivitySnapshotController.hasRecord(activityB));
+        final ActivityRecord[] singleActivityList = new ActivityRecord[1];
+        singleActivityList[0] = activityA;
+        assertNull(mActivitySnapshotController.getSnapshot(singleActivityList));
+        singleActivityList[0] = activityB;
+        assertNull(mActivitySnapshotController.getSnapshot(singleActivityList));
+        final ActivityRecord[] activities = new ActivityRecord[3];
+        activities[0] = activityA;
+        activities[1] = activityB;
+        activities[2] = activityC;
+        assertNull(mActivitySnapshotController.getSnapshot(activities));
+
+        // Reset and test load snapshot
+        mActivitySnapshotController.addUserSavedFile(testTask.mUserId, taskSnapshot, mixedCode);
+        // Request to load by B, nothing will be loaded because the snapshot was [A,B,C].
+        mActivitySnapshotController.mPendingLoadActivity.add(activityB);
+        mActivitySnapshotController.loadActivitySnapshot();
+        verify(mActivitySnapshotController, never()).loadSnapshotInner(any(), any());
+
+        // Able to load snapshot when requesting for all A, B, C
+        mActivitySnapshotController.mPendingLoadActivity.clear();
+        mActivitySnapshotController.mPendingLoadActivity.add(activityA);
+        mActivitySnapshotController.mPendingLoadActivity.add(activityB);
+        mActivitySnapshotController.mPendingLoadActivity.add(activityC);
+        final ArraySet<ActivityRecord> verifyList = new ArraySet<>();
+        verifyList.add(activityA);
+        verifyList.add(activityB);
+        verifyList.add(activityC);
+        mActivitySnapshotController.loadActivitySnapshot();
+        verify(mActivitySnapshotController).loadSnapshotInner(argThat(
+                argument -> {
+                    final ArrayList<ActivityRecord> argumentList = new ArrayList<>(
+                            Arrays.asList(argument));
+                    return verifyList.containsAll(argumentList)
+                            && argumentList.containsAll(verifyList);
+                }),
+                any());
+
+        for (int i = activities.length - 1; i >= 0; --i) {
+            mActivitySnapshotController.mCache.putSnapshot(activities[i], taskSnapshot);
+        }
+        // The loaded snapshot can be retrieved only if the activities match exactly.
+        singleActivityList[0] = activityB;
+        assertNull(mActivitySnapshotController.getSnapshot(singleActivityList));
+        assertEquals(taskSnapshot, mActivitySnapshotController.getSnapshot(activities));
+    }
+
+    private TaskSnapshot createSnapshot() {
+        HardwareBuffer buffer = mock(HardwareBuffer.class);
+        doReturn(100).when(buffer).getWidth();
+        doReturn(100).when(buffer).getHeight();
+        return new TaskSnapshot(1, 0 /* captureTime */, new ComponentName("", ""), buffer,
+                ColorSpace.get(ColorSpace.Named.SRGB), ORIENTATION_PORTRAIT,
+                Surface.ROTATION_0, new Point(100, 100), new Rect() /* contentInsets */,
+                new Rect() /* letterboxInsets*/, false /* isLowResolution */,
+                true /* isRealSnapshot */, WINDOWING_MODE_FULLSCREEN, 0 /* mSystemUiVisibility */,
+                false /* isTranslucent */, false /* hasImeSurface */);
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
index ac18f802d1c6..3378ce526a5e 100644
--- a/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/BackNavigationControllerTests.java
@@ -741,7 +741,7 @@ public class BackNavigationControllerTests extends WindowTestsBase {
 
         MockitoSession mockitoSession = mockitoSession().mockStatic(BackNavigationController.class)
                 .strictness(Strictness.LENIENT).startMocking();
-        doReturn(taskSnapshot).when(() -> BackNavigationController.getSnapshot(any()));
+        doReturn(taskSnapshot).when(() -> BackNavigationController.getSnapshot(any(), any()));
         when(resourcesSpy.getBoolean(
                 com.android.internal.R.bool.config_predictShowStartingSurface))
                 .thenReturn(preferWindowlessSurface);