Merge "Add parameter to save task snapshot when it is requested" into udc-dev

7 files changed, 29 insertions, 11 deletions

diff --git a/core/java/android/app/IActivityTaskManager.aidl b/core/java/android/app/IActivityTaskManager.aidl
index 3249b41da934..4e2b6fa56b17 100644
--- a/core/java/android/app/IActivityTaskManager.aidl
+++ b/core/java/android/app/IActivityTaskManager.aidl
@@ -261,6 +261,9 @@ interface IActivityTaskManager {
     void cancelTaskWindowTransition(int taskId);
 
     /**
+     * Fetches the snapshot for the task with the given id, taking a new snapshot if it is not in
+     * the task snapshot cache and it is requested.
+     *
      * @param taskId the id of the task to retrieve the sAutoapshots for
      * @param isLowResolution if set, if the snapshot needs to be loaded from disk, this will load
      *                          a reduced resolution of it, which is much faster
@@ -272,10 +275,14 @@ interface IActivityTaskManager {
             int taskId, boolean isLowResolution, boolean takeSnapshotIfNeeded);
 
     /**
+     * Requests for a new snapshot to be taken for the task with the given id, storing it in the
+     * task snapshot cache only if requested.
+     *
      * @param taskId the id of the task to take a snapshot of
+     * @param updateCache whether to store the new snapshot in the system's task snapshot cache
      * @return a graphic buffer representing a screenshot of a task
      */
-    android.window.TaskSnapshot takeTaskSnapshot(int taskId);
+    android.window.TaskSnapshot takeTaskSnapshot(int taskId, boolean updateCache);
 
     /**
      * Return the user id of last resumed activity.
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
index 387d39056ffc..a9ad3c90949f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentsTransitionHandler.java
@@ -301,7 +301,7 @@ public class RecentsTransitionHandler implements Transitions.TransitionHandler {
                 try {
                     for (int i = 0; i < mPausingTasks.size(); ++i) {
                         snapshots[i] = ActivityTaskManager.getService().takeTaskSnapshot(
-                                mPausingTasks.get(0).mTaskInfo.taskId);
+                                mPausingTasks.get(0).mTaskInfo.taskId, false /* updateCache */);
                     }
                 } catch (RemoteException e) {
                     taskIds = null;
@@ -671,7 +671,8 @@ public class RecentsTransitionHandler implements Transitions.TransitionHandler {
             try {
                 ProtoLog.v(ShellProtoLogGroup.WM_SHELL_RECENTS_TRANSITION,
                         "[%d] RecentsController.screenshotTask: taskId=%d", mInstanceId, taskId);
-                return ActivityTaskManager.getService().takeTaskSnapshot(taskId);
+                return ActivityTaskManager.getService().takeTaskSnapshot(taskId,
+                        true /* updateCache */);
             } catch (RemoteException e) {
                 Slog.e(TAG, "Failed to screenshot task", e);
             }
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/system/RemoteTransitionCompat.java b/packages/SystemUI/shared/src/com/android/systemui/shared/system/RemoteTransitionCompat.java
index 1fbf743836a1..74c325dea15c 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/system/RemoteTransitionCompat.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/system/RemoteTransitionCompat.java
@@ -352,7 +352,8 @@ public class RemoteTransitionCompat {
 
         @Override public TaskSnapshot screenshotTask(int taskId) {
             try {
-                return ActivityTaskManager.getService().takeTaskSnapshot(taskId);
+                return ActivityTaskManager.getService().takeTaskSnapshot(taskId,
+                        true /* updateCache */);
             } catch (RemoteException e) {
                 Log.e(TAG, "Failed to screenshot task", e);
             }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
index 67e9a87e4f18..2e47ab6417ae 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ImageCaptureImpl.kt
@@ -48,7 +48,9 @@ open class ImageCaptureImpl @Inject constructor(
     }
 
     override suspend fun captureTask(taskId: Int): Bitmap? {
-        val snapshot = withContext(bgContext) { atmService.takeTaskSnapshot(taskId) } ?: return null
+        val snapshot = withContext(bgContext) {
+            atmService.takeTaskSnapshot(taskId, false /* updateCache */)
+        } ?: return null
         return Bitmap.wrapHardwareBuffer(snapshot.hardwareBuffer, snapshot.colorSpace)
     }
 }
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index cff65547c673..1360a956dc0b 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -3789,7 +3789,7 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
             TaskSnapshot taskSnapshot = mWindowManager.mTaskSnapshotController.getSnapshot(taskId,
                     task.mUserId, true /* restoreFromDisk */, isLowResolution);
             if (taskSnapshot == null && takeSnapshotIfNeeded) {
-                taskSnapshot = takeTaskSnapshot(taskId);
+                taskSnapshot = takeTaskSnapshot(taskId, false /* updateCache */);
             }
             return taskSnapshot;
         } finally {
@@ -3798,7 +3798,7 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
     }
 
     @Override
-    public TaskSnapshot takeTaskSnapshot(int taskId) {
+    public TaskSnapshot takeTaskSnapshot(int taskId, boolean updateCache) {
         mAmInternal.enforceCallingPermission(READ_FRAME_BUFFER, "takeTaskSnapshot()");
         final long ident = Binder.clearCallingIdentity();
         try {
@@ -3809,8 +3809,13 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
                     Slog.w(TAG, "takeTaskSnapshot: taskId=" + taskId + " not found or not visible");
                     return null;
                 }
-                return mWindowManager.mTaskSnapshotController.captureSnapshot(
-                        task, true /* snapshotHome */);
+                if (updateCache) {
+                    return mWindowManager.mTaskSnapshotController.recordSnapshot(task,
+                            true /* snapshotHome */);
+                } else {
+                    return mWindowManager.mTaskSnapshotController.captureSnapshot(task,
+                            true /* snapshotHome */);
+                }
             }
         } finally {
             Binder.restoreCallingIdentity(ident);
diff --git a/services/core/java/com/android/server/wm/InputMonitor.java b/services/core/java/com/android/server/wm/InputMonitor.java
index b67ccd2a7672..cea886ff6e87 100644
--- a/services/core/java/com/android/server/wm/InputMonitor.java
+++ b/services/core/java/com/android/server/wm/InputMonitor.java
@@ -440,7 +440,8 @@ final class InputMonitor {
                         if (app != null) {
                             mDisplayContent.removeImeSurfaceImmediately();
                             if (app.getTask() != null) {
-                                mDisplayContent.mAtmService.takeTaskSnapshot(app.getTask().mTaskId);
+                                mDisplayContent.mAtmService.takeTaskSnapshot(app.getTask().mTaskId,
+                                        true /* updateCache */);
                             }
                         }
                     } else {
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index 4d0bff961e60..20ef0171358b 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -171,13 +171,14 @@ class TaskSnapshotController extends AbsAppSnapshotController<Task, TaskSnapshot
         snapshotTasks(tasks, false /* allowSnapshotHome */);
     }
 
-    void recordSnapshot(Task task, boolean allowSnapshotHome) {
+    TaskSnapshot recordSnapshot(Task task, boolean allowSnapshotHome) {
         final boolean snapshotHome = allowSnapshotHome && task.isActivityTypeHome();
         final TaskSnapshot snapshot = recordSnapshotInner(task, allowSnapshotHome);
         if (!snapshotHome && snapshot != null) {
             mPersister.persistSnapshot(task.mTaskId, task.mUserId, snapshot);
             task.onSnapshotChanged(snapshot);
         }
+        return snapshot;
     }
 
     private void snapshotTasks(ArraySet<Task> tasks, boolean allowSnapshotHome) {