ContentSuggestions - restore snapshots when needed

Make ContentSuggestions restore snapshots from disk where needed.
Previously this wasn't needed as only the running task would be
used for content suggestions. Now it can be invoked for any task
in Overview so need to be able to restore snapshots from disk.

Bug: 149595462
Test: manual
Change-Id: Ia67efc122a44b27a9049e48afaf2ddb789bc31b9

4 files changed, 40 insertions, 27 deletions

diff --git a/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsManagerService.java b/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsManagerService.java
index b54ec4ea9441..0fdabd09f9d4 100644
--- a/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsManagerService.java
+++ b/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsManagerService.java
@@ -21,13 +21,17 @@ import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.app.ActivityManager;
 import android.app.contentsuggestions.ClassificationsRequest;
+import android.app.contentsuggestions.ContentSuggestionsManager;
 import android.app.contentsuggestions.IClassificationsCallback;
 import android.app.contentsuggestions.IContentSuggestionsManager;
 import android.app.contentsuggestions.ISelectionsCallback;
 import android.app.contentsuggestions.SelectionsRequest;
 import android.content.Context;
 import android.graphics.Bitmap;
+import android.graphics.ColorSpace;
+import android.graphics.GraphicBuffer;
 import android.os.Binder;
 import android.os.Bundle;
 import android.os.RemoteException;
@@ -133,7 +137,7 @@ public class ContentSuggestionsManagerService extends
                 if (service != null) {
                     // TODO(b/147324195): Temporarily pass bitmap until we change the service API.
                     imageContextRequestExtras.putParcelable(EXTRA_BITMAP, bitmap);
-                    service.provideContextImageLocked(/* taskId = */ -1, imageContextRequestExtras);
+                    service.provideContextImageFromBitmapLocked(imageContextRequestExtras);
                 } else {
                     if (VERBOSE) {
                         Slog.v(TAG, "provideContextImageLocked: no service for " + userId);
@@ -152,10 +156,28 @@ public class ContentSuggestionsManagerService extends
             }
             enforceCaller(UserHandle.getCallingUserId(), "provideContextImage");
 
+            GraphicBuffer snapshotBuffer = null;
+            int colorSpaceId = 0;
+
+            // Skip taking TaskSnapshot when bitmap is provided.
+            if (!imageContextRequestExtras.containsKey(ContentSuggestionsManager.EXTRA_BITMAP)) {
+                // Can block, so call before acquiring the lock.
+                ActivityManager.TaskSnapshot snapshot =
+                        mActivityTaskManagerInternal.getTaskSnapshotBlocking(taskId, false);
+                if (snapshot != null) {
+                    snapshotBuffer = snapshot.getSnapshot();
+                    ColorSpace colorSpace = snapshot.getColorSpace();
+                    if (colorSpace != null) {
+                        colorSpaceId = colorSpace.getId();
+                    }
+                }
+            }
+
             synchronized (mLock) {
                 final ContentSuggestionsPerUserService service = getServiceForUserLocked(userId);
                 if (service != null) {
-                    service.provideContextImageLocked(taskId, imageContextRequestExtras);
+                    service.provideContextImageLocked(taskId, snapshotBuffer, colorSpaceId,
+                            imageContextRequestExtras);
                 } else {
                     if (VERBOSE) {
                         Slog.v(TAG, "provideContextImageLocked: no service for " + userId);
diff --git a/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsPerUserService.java b/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsPerUserService.java
index 7828050223f4..cf53b169515e 100644
--- a/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsPerUserService.java
+++ b/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsPerUserService.java
@@ -19,17 +19,14 @@ package com.android.server.contentsuggestions;
 import android.Manifest;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.app.ActivityManager;
 import android.app.AppGlobals;
 import android.app.contentsuggestions.ClassificationsRequest;
-import android.app.contentsuggestions.ContentSuggestionsManager;
 import android.app.contentsuggestions.IClassificationsCallback;
 import android.app.contentsuggestions.ISelectionsCallback;
 import android.app.contentsuggestions.SelectionsRequest;
 import android.content.ComponentName;
 import android.content.pm.PackageManager;
 import android.content.pm.ServiceInfo;
-import android.graphics.ColorSpace;
 import android.graphics.GraphicBuffer;
 import android.os.Bundle;
 import android.os.RemoteException;
@@ -95,26 +92,17 @@ public final class ContentSuggestionsPerUserService extends
     }
 
     @GuardedBy("mLock")
-    void provideContextImageLocked(int taskId, @NonNull Bundle imageContextRequestExtras) {
+    void provideContextImageFromBitmapLocked(@NonNull Bundle bitmapContainingExtras) {
+        // No task or snapshot provided, the bitmap is contained in the extras
+        provideContextImageLocked(-1, null, 0, bitmapContainingExtras);
+    }
+
+    @GuardedBy("mLock")
+    void provideContextImageLocked(int taskId, @Nullable GraphicBuffer snapshot,
+            int colorSpaceIdForSnapshot, @NonNull Bundle imageContextRequestExtras) {
         RemoteContentSuggestionsService service = ensureRemoteServiceLocked();
         if (service != null) {
-            GraphicBuffer snapshotBuffer = null;
-            int colorSpaceId = 0;
-
-            // Skip taking TaskSnapshot when bitmap is provided.
-            if (!imageContextRequestExtras.containsKey(ContentSuggestionsManager.EXTRA_BITMAP)) {
-                ActivityManager.TaskSnapshot snapshot =
-                        mActivityTaskManagerInternal.getTaskSnapshotNoRestore(taskId, false);
-                if (snapshot != null) {
-                    snapshotBuffer = snapshot.getSnapshot();
-                    ColorSpace colorSpace = snapshot.getColorSpace();
-                    if (colorSpace != null) {
-                        colorSpaceId = colorSpace.getId();
-                    }
-                }
-            }
-
-            service.provideContextImage(taskId, snapshotBuffer, colorSpaceId,
+            service.provideContextImage(taskId, snapshot, colorSpaceIdForSnapshot,
                     imageContextRequestExtras);
         }
     }
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java b/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
index edc87e5a4d88..2263795f9442 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
@@ -548,8 +548,11 @@ public abstract class ActivityTaskManagerInternal {
 
     /**
      * Gets bitmap snapshot of the provided task id.
+     *
+     * <p>Warning! this may restore the snapshot from disk so can block, don't call in a latency
+     * sensitive environment.
      */
-    public abstract ActivityManager.TaskSnapshot getTaskSnapshotNoRestore(int taskId,
+    public abstract ActivityManager.TaskSnapshot getTaskSnapshotBlocking(int taskId,
             boolean isLowResolution);
 
     /** Returns true if uid is considered foreground for activity start purposes. */
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 890b9454e6b6..659bb9956604 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -7476,10 +7476,10 @@ public class ActivityTaskManagerService extends IActivityTaskManager.Stub {
         }
 
         @Override
-        public ActivityManager.TaskSnapshot getTaskSnapshotNoRestore(int taskId,
-                boolean isLowResolution) {
+        public ActivityManager.TaskSnapshot getTaskSnapshotBlocking(
+                int taskId, boolean isLowResolution) {
             return ActivityTaskManagerService.this.getTaskSnapshot(taskId, isLowResolution,
-                    false /* restoreFromDisk */);
+                    true /* restoreFromDisk */);
         }
 
         @Override