Pass TaskSnapshot object to content suggestion service.

To easier track where do access the HardwareBuffer from a TaskSnapshot
object. And release the reference as soon as possible.

Remove unused members.

Flag: com.android.window.flags.release_snapshot_aggressively
Bug: 238206323
Test: atest TaplOverviewIconTest
Test: atest FlickerTestsNotification
Change-Id: I8e633371153515013ebded52c93fe62d94f8bd30

5 files changed, 31 insertions, 44 deletions

diff --git a/core/java/android/service/contentsuggestions/ContentSuggestionsService.java b/core/java/android/service/contentsuggestions/ContentSuggestionsService.java
index 0b4739ea1afc..7f4e67a7a270 100644
--- a/core/java/android/service/contentsuggestions/ContentSuggestionsService.java
+++ b/core/java/android/service/contentsuggestions/ContentSuggestionsService.java
@@ -39,6 +39,7 @@ import android.os.Looper;
 import android.os.RemoteException;
 import android.util.Log;
 import android.util.Slog;
+import android.window.TaskSnapshot;
 
 /**
  * @hide
@@ -62,10 +63,10 @@ public abstract class ContentSuggestionsService extends Service {
 
     private final IContentSuggestionsService mInterface = new IContentSuggestionsService.Stub() {
         @Override
-        public void provideContextImage(int taskId, HardwareBuffer contextImage,
-                int colorSpaceId, Bundle imageContextRequestExtras) {
+        public void provideContextImage(int taskId, TaskSnapshot snapshot,
+                Bundle imageContextRequestExtras) {
             if (imageContextRequestExtras.containsKey(ContentSuggestionsManager.EXTRA_BITMAP)
-                    && contextImage != null) {
+                    && snapshot != null) {
                 throw new IllegalArgumentException("Two bitmaps provided; expected one.");
             }
 
@@ -74,13 +75,18 @@ public abstract class ContentSuggestionsService extends Service {
                 wrappedBuffer = imageContextRequestExtras.getParcelable(
                         ContentSuggestionsManager.EXTRA_BITMAP, android.graphics.Bitmap.class);
             } else {
-                if (contextImage != null) {
-                    ColorSpace colorSpace = null;
+                if (snapshot != null) {
+                    final HardwareBuffer snapshotBuffer = snapshot.getHardwareBuffer();
+                    ColorSpace colorSpace = snapshot.getColorSpace();
+                    int colorSpaceId = 0;
+                    if (colorSpace != null) {
+                        colorSpaceId = colorSpace.getId();
+                    }
                     if (colorSpaceId >= 0 && colorSpaceId < ColorSpace.Named.values().length) {
                         colorSpace = ColorSpace.get(ColorSpace.Named.values()[colorSpaceId]);
                     }
-                    wrappedBuffer = Bitmap.wrapHardwareBuffer(contextImage, colorSpace);
-                    contextImage.close();
+                    wrappedBuffer = Bitmap.wrapHardwareBuffer(snapshotBuffer, colorSpace);
+                    snapshotBuffer.close();
                 }
             }
 
diff --git a/core/java/android/service/contentsuggestions/IContentSuggestionsService.aidl b/core/java/android/service/contentsuggestions/IContentSuggestionsService.aidl
index d8f23e78dd65..398fa37f5e90 100644
--- a/core/java/android/service/contentsuggestions/IContentSuggestionsService.aidl
+++ b/core/java/android/service/contentsuggestions/IContentSuggestionsService.aidl
@@ -20,7 +20,7 @@ import android.app.contentsuggestions.IClassificationsCallback;
 import android.app.contentsuggestions.ISelectionsCallback;
 import android.app.contentsuggestions.ClassificationsRequest;
 import android.app.contentsuggestions.SelectionsRequest;
-import android.hardware.HardwareBuffer;
+import android.window.TaskSnapshot;
 import android.os.Bundle;
 
 /**
@@ -31,8 +31,7 @@ import android.os.Bundle;
 oneway interface IContentSuggestionsService {
     void provideContextImage(
             int taskId,
-            in HardwareBuffer contextImage,
-            int colorSpaceId,
+            in TaskSnapshot snapshot,
             in Bundle imageContextRequestExtras);
     void suggestContentSelections(
             in SelectionsRequest request,
diff --git a/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsManagerService.java b/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsManagerService.java
index a7aab49cde56..84aa331efc1a 100644
--- a/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsManagerService.java
+++ b/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsManagerService.java
@@ -30,8 +30,6 @@ import android.app.contentsuggestions.ISelectionsCallback;
 import android.app.contentsuggestions.SelectionsRequest;
 import android.content.Context;
 import android.graphics.Bitmap;
-import android.graphics.ColorSpace;
-import android.hardware.HardwareBuffer;
 import android.os.Binder;
 import android.os.Bundle;
 import android.os.RemoteException;
@@ -147,6 +145,7 @@ public class ContentSuggestionsManagerService extends
             }
         }
 
+        @SuppressWarnings("GuardedBy")
         @Override
         public void provideContextImage(
                 int userId,
@@ -157,9 +156,6 @@ public class ContentSuggestionsManagerService extends
             }
             enforceCaller(UserHandle.getCallingUserId(), "provideContextImage");
 
-            HardwareBuffer snapshotBuffer = null;
-            int colorSpaceId = 0;
-
             TaskSnapshot snapshot = null;
             // Skip taking TaskSnapshot when bitmap is provided.
             if (!imageContextRequestExtras.containsKey(ContentSuggestionsManager.EXTRA_BITMAP)) {
@@ -167,29 +163,18 @@ public class ContentSuggestionsManagerService extends
                 snapshot = mActivityTaskManagerInternal.getTaskSnapshotBlocking(
                         taskId, false /* isLowResolution */,
                         TaskSnapshot.REFERENCE_CONTENT_SUGGESTION);
-                if (snapshot != null) {
-                    snapshotBuffer = snapshot.getHardwareBuffer();
-                    ColorSpace colorSpace = snapshot.getColorSpace();
-                    if (colorSpace != null) {
-                        colorSpaceId = colorSpace.getId();
-                    }
-                }
             }
 
             synchronized (mLock) {
                 final ContentSuggestionsPerUserService service = getServiceForUserLocked(userId);
                 if (service != null) {
-                    service.provideContextImageLocked(taskId, snapshotBuffer, colorSpaceId,
-                            imageContextRequestExtras);
+                    service.provideContextImageLocked(taskId, snapshot, imageContextRequestExtras);
                 } else {
                     if (VERBOSE) {
                         Slog.v(TAG, "provideContextImageLocked: no service for " + userId);
                     }
                 }
             }
-            if (snapshot != null) {
-                snapshot.removeReference(TaskSnapshot.REFERENCE_CONTENT_SUGGESTION);
-            }
         }
 
         @Override
diff --git a/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsPerUserService.java b/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsPerUserService.java
index c8588e2d98e6..6f543a5a0916 100644
--- a/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsPerUserService.java
+++ b/services/contentsuggestions/java/com/android/server/contentsuggestions/ContentSuggestionsPerUserService.java
@@ -27,15 +27,13 @@ import android.app.contentsuggestions.SelectionsRequest;
 import android.content.ComponentName;
 import android.content.pm.PackageManager;
 import android.content.pm.ServiceInfo;
-import android.hardware.HardwareBuffer;
 import android.os.Bundle;
 import android.os.RemoteException;
 import android.util.Slog;
+import android.window.TaskSnapshot;
 
 import com.android.internal.annotations.GuardedBy;
-import com.android.server.LocalServices;
 import com.android.server.infra.AbstractPerUserSystemService;
-import com.android.server.wm.ActivityTaskManagerInternal;
 
 /**
  * Per user delegate of {@link ContentSuggestionsManagerService}.
@@ -52,13 +50,9 @@ public final class ContentSuggestionsPerUserService extends
     @GuardedBy("mLock")
     private RemoteContentSuggestionsService mRemoteService;
 
-    @NonNull
-    private final ActivityTaskManagerInternal mActivityTaskManagerInternal;
-
     ContentSuggestionsPerUserService(
             ContentSuggestionsManagerService master, Object lock, int userId) {
         super(master, lock, userId);
-        mActivityTaskManagerInternal = LocalServices.getService(ActivityTaskManagerInternal.class);
     }
 
     @GuardedBy("mLock")
@@ -94,16 +88,15 @@ public final class ContentSuggestionsPerUserService extends
     @GuardedBy("mLock")
     void provideContextImageFromBitmapLocked(@NonNull Bundle bitmapContainingExtras) {
         // No task or snapshot provided, the bitmap is contained in the extras
-        provideContextImageLocked(-1, null, 0, bitmapContainingExtras);
+        provideContextImageLocked(-1, null, bitmapContainingExtras);
     }
 
     @GuardedBy("mLock")
-    void provideContextImageLocked(int taskId, @Nullable HardwareBuffer snapshot,
-            int colorSpaceIdForSnapshot, @NonNull Bundle imageContextRequestExtras) {
+    void provideContextImageLocked(int taskId, @Nullable TaskSnapshot snapshot,
+            @NonNull Bundle imageContextRequestExtras) {
         RemoteContentSuggestionsService service = ensureRemoteServiceLocked();
         if (service != null) {
-            service.provideContextImage(taskId, snapshot, colorSpaceIdForSnapshot,
-                    imageContextRequestExtras);
+            service.provideContextImage(taskId, snapshot, imageContextRequestExtras);
         }
     }
 
diff --git a/services/contentsuggestions/java/com/android/server/contentsuggestions/RemoteContentSuggestionsService.java b/services/contentsuggestions/java/com/android/server/contentsuggestions/RemoteContentSuggestionsService.java
index 6e0a0da76d05..92fa3b256646 100644
--- a/services/contentsuggestions/java/com/android/server/contentsuggestions/RemoteContentSuggestionsService.java
+++ b/services/contentsuggestions/java/com/android/server/contentsuggestions/RemoteContentSuggestionsService.java
@@ -24,12 +24,12 @@ import android.app.contentsuggestions.ISelectionsCallback;
 import android.app.contentsuggestions.SelectionsRequest;
 import android.content.ComponentName;
 import android.content.Context;
-import android.hardware.HardwareBuffer;
 import android.os.Bundle;
 import android.os.IBinder;
 import android.service.contentsuggestions.ContentSuggestionsService;
 import android.service.contentsuggestions.IContentSuggestionsService;
 import android.text.format.DateUtils;
+import android.window.TaskSnapshot;
 
 import com.android.internal.infra.AbstractMultiplePendingRequestsRemoteService;
 
@@ -67,10 +67,14 @@ public class RemoteContentSuggestionsService extends
         return TIMEOUT_REMOTE_REQUEST_MILLIS;
     }
 
-    void provideContextImage(int taskId, @Nullable HardwareBuffer contextImage,
-            int colorSpaceId, @NonNull Bundle imageContextRequestExtras) {
-        scheduleAsyncRequest((s) -> s.provideContextImage(taskId, contextImage,
-                colorSpaceId, imageContextRequestExtras));
+    void provideContextImage(int taskId, @Nullable TaskSnapshot snapshot,
+            @NonNull Bundle imageContextRequestExtras) {
+        scheduleAsyncRequest((s) -> {
+            s.provideContextImage(taskId, snapshot, imageContextRequestExtras);
+            if (snapshot != null) {
+                snapshot.removeReference(TaskSnapshot.REFERENCE_CONTENT_SUGGESTION);
+            }
+        });
     }
 
     void suggestContentSelections(