Merge "Defer close snapshot's hardware buffer after remove snapshot from cache" into main

7 files changed, 80 insertions, 1 deletions

diff --git a/core/java/android/window/TaskSnapshot.java b/core/java/android/window/TaskSnapshot.java
index 53c64bd6e664..4ddd73319840 100644
--- a/core/java/android/window/TaskSnapshot.java
+++ b/core/java/android/window/TaskSnapshot.java
@@ -37,6 +37,7 @@ import com.android.window.flags.Flags;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
+import java.util.function.Consumer;
 
 /**
  * Represents a task snapshot.
@@ -76,6 +77,7 @@ public class TaskSnapshot implements Parcelable {
     // Must be one of the named color spaces, otherwise, always use SRGB color space.
     private final ColorSpace mColorSpace;
     private int mInternalReferences;
+    private Consumer<HardwareBuffer> mSafeSnapshotReleaser;
 
     /** Keep in cache, doesn't need reference. */
     public static final int REFERENCE_NONE = 0;
@@ -365,8 +367,24 @@ public class TaskSnapshot implements Parcelable {
         mInternalReferences &= ~usage;
         if (Flags.releaseSnapshotAggressively() && mInternalReferences == 0 && mSnapshot != null
                 && !mSnapshot.isClosed()) {
-            mSnapshot.close();
+            if (mSafeSnapshotReleaser != null) {
+                mSafeSnapshotReleaser.accept(mSnapshot);
+            } else {
+                mSnapshot.close();
+            }
+        }
+    }
+
+    /**
+     * Register a safe release callback, instead of immediately closing the hardware buffer when
+     * no more reference, to let the system server decide when to close it.
+     * Only used in core.
+     */
+    public synchronized void setSafeRelease(Consumer<HardwareBuffer> releaser) {
+        if (!Flags.safeReleaseSnapshotAggressively()) {
+            return;
         }
+        mSafeSnapshotReleaser = releaser;
     }
 
     public static final @NonNull Creator<TaskSnapshot> CREATOR = new Creator<TaskSnapshot>() {
diff --git a/core/java/android/window/flags/windowing_frontend.aconfig b/core/java/android/window/flags/windowing_frontend.aconfig
index b19ab4527f06..f3c00f27072e 100644
--- a/core/java/android/window/flags/windowing_frontend.aconfig
+++ b/core/java/android/window/flags/windowing_frontend.aconfig
@@ -516,6 +516,17 @@ flag {
 }
 
 flag {
+    name: "safe_release_snapshot_aggressively"
+    namespace: "windowing_frontend"
+    description: "Protect task snapshot memory from premature release, which can occur when a local variable holds a reference while the snapshot is removed from the cache."
+    bug: "238206323"
+    is_fixed_read_only: true
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
+
+flag {
     name: "scramble_snapshot_file_name"
     namespace: "windowing_frontend"
     description: "Scramble the file name of task snapshot."
diff --git a/services/core/java/com/android/server/wm/AbsAppSnapshotController.java b/services/core/java/com/android/server/wm/AbsAppSnapshotController.java
index 70fc6bace868..ee173a2a6a35 100644
--- a/services/core/java/com/android/server/wm/AbsAppSnapshotController.java
+++ b/services/core/java/com/android/server/wm/AbsAppSnapshotController.java
@@ -131,6 +131,10 @@ abstract class AbsAppSnapshotController<TYPE extends WindowContainer,
         mCache = cache;
     }
 
+    void setSnapshotReleaser(Consumer<HardwareBuffer> releaser) {
+        mCache.setSafeSnapshotReleaser(releaser);
+    }
+
     void setSnapshotEnabled(boolean enabled) {
         mSnapshotEnabled = enabled;
     }
diff --git a/services/core/java/com/android/server/wm/ActivitySnapshotCache.java b/services/core/java/com/android/server/wm/ActivitySnapshotCache.java
index ed07afd2eab5..a0e537231071 100644
--- a/services/core/java/com/android/server/wm/ActivitySnapshotCache.java
+++ b/services/core/java/com/android/server/wm/ActivitySnapshotCache.java
@@ -31,6 +31,7 @@ class ActivitySnapshotCache extends SnapshotCache<ActivityRecord> {
     void putSnapshot(ActivityRecord ar, TaskSnapshot snapshot) {
         final int hasCode = System.identityHashCode(ar);
         snapshot.addReference(TaskSnapshot.REFERENCE_CACHE);
+        snapshot.setSafeRelease(mSafeSnapshotReleaser);
         synchronized (mLock) {
             final CacheEntry entry = mRunningCache.get(hasCode);
             if (entry != null) {
diff --git a/services/core/java/com/android/server/wm/SnapshotCache.java b/services/core/java/com/android/server/wm/SnapshotCache.java
index 9812a88bfa5a..685bd3b355ed 100644
--- a/services/core/java/com/android/server/wm/SnapshotCache.java
+++ b/services/core/java/com/android/server/wm/SnapshotCache.java
@@ -16,12 +16,14 @@
 package com.android.server.wm;
 
 import android.annotation.Nullable;
+import android.hardware.HardwareBuffer;
 import android.util.ArrayMap;
 import android.window.TaskSnapshot;
 
 import com.android.internal.annotations.GuardedBy;
 
 import java.io.PrintWriter;
+import java.util.function.Consumer;
 
 /**
  * Base class for an app snapshot cache
@@ -38,12 +40,18 @@ abstract class SnapshotCache<TYPE extends WindowContainer> {
     @GuardedBy("mLock")
     protected final ArrayMap<Integer, CacheEntry> mRunningCache = new ArrayMap<>();
 
+    protected Consumer<HardwareBuffer> mSafeSnapshotReleaser;
+
     SnapshotCache(String name) {
         mName = name;
     }
 
     abstract void putSnapshot(TYPE window, TaskSnapshot snapshot);
 
+    void setSafeSnapshotReleaser(Consumer<HardwareBuffer> safeSnapshotReleaser) {
+        mSafeSnapshotReleaser = safeSnapshotReleaser;
+    }
+
     void clearRunningCache() {
         synchronized (mLock) {
             mRunningCache.clear();
diff --git a/services/core/java/com/android/server/wm/SnapshotController.java b/services/core/java/com/android/server/wm/SnapshotController.java
index 3a7222ae6d51..48f0959214c1 100644
--- a/services/core/java/com/android/server/wm/SnapshotController.java
+++ b/services/core/java/com/android/server/wm/SnapshotController.java
@@ -26,13 +26,16 @@ import static android.view.WindowManager.TRANSIT_PREPARE_BACK_NAVIGATION;
 import static android.view.WindowManager.TRANSIT_TO_BACK;
 import static android.view.WindowManager.TRANSIT_TO_FRONT;
 
+import android.hardware.HardwareBuffer;
 import android.os.Trace;
 import android.util.ArrayMap;
 import android.view.WindowManager;
 import android.window.TaskSnapshot;
 
 import java.io.PrintWriter;
+import java.lang.ref.WeakReference;
 import java.util.ArrayList;
+import java.util.function.Consumer;
 
 /**
  * Integrates common functionality from TaskSnapshotController and ActivitySnapshotController.
@@ -41,11 +44,30 @@ class SnapshotController {
     private final SnapshotPersistQueue mSnapshotPersistQueue;
     final TaskSnapshotController mTaskSnapshotController;
     final ActivitySnapshotController mActivitySnapshotController;
+    private final WindowManagerService mService;
+    private final ArrayList<WeakReference<HardwareBuffer>> mObsoleteSnapshots = new ArrayList<>();
 
     SnapshotController(WindowManagerService wms) {
+        mService = wms;
         mSnapshotPersistQueue = new SnapshotPersistQueue();
         mTaskSnapshotController = new TaskSnapshotController(wms, mSnapshotPersistQueue);
         mActivitySnapshotController = new ActivitySnapshotController(wms, mSnapshotPersistQueue);
+        final Consumer<HardwareBuffer> releaser = hb -> {
+            mService.mH.post(() -> {
+                synchronized (mService.mGlobalLock) {
+                    if (hb.isClosed()) {
+                        return;
+                    }
+                    if (mService.mAtmService.getTransitionController().inTransition()) {
+                        mObsoleteSnapshots.add(new WeakReference<>(hb));
+                    } else {
+                        hb.close();
+                    }
+                }
+            });
+        };
+        mTaskSnapshotController.setSnapshotReleaser(releaser);
+        mActivitySnapshotController.setSnapshotReleaser(releaser);
     }
 
     void systemReady() {
@@ -168,6 +190,7 @@ class SnapshotController {
         final boolean isTransitionClose = isTransitionClose(type);
         if (!isTransitionOpen && !isTransitionClose && type < TRANSIT_FIRST_CUSTOM
                 || (changeInfos.isEmpty())) {
+            closeObsoleteSnapshots();
             return;
         }
         Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "SnapshotController_analysis");
@@ -195,9 +218,22 @@ class SnapshotController {
                 }
             }
         }
+        closeObsoleteSnapshots();
         Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
     }
 
+    private void closeObsoleteSnapshots() {
+        if (mObsoleteSnapshots.isEmpty()) {
+            return;
+        }
+        for (int i = mObsoleteSnapshots.size() - 1; i >= 0; --i) {
+            final HardwareBuffer hb = mObsoleteSnapshots.remove(i).get();
+            if (hb != null && !hb.isClosed()) {
+                hb.close();
+            }
+        }
+    }
+
     private static boolean isTransitionOpen(int type) {
         return type == TRANSIT_OPEN || type == TRANSIT_TO_FRONT
                 || type == TRANSIT_PREPARE_BACK_NAVIGATION;
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotCache.java b/services/core/java/com/android/server/wm/TaskSnapshotCache.java
index cc957bd9ee42..1988f2630b6e 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotCache.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotCache.java
@@ -36,6 +36,7 @@ class TaskSnapshotCache extends SnapshotCache<Task> {
     void putSnapshot(Task task, TaskSnapshot snapshot) {
         synchronized (mLock) {
             snapshot.addReference(TaskSnapshot.REFERENCE_CACHE);
+            snapshot.setSafeRelease(mSafeSnapshotReleaser);
             final CacheEntry entry = mRunningCache.get(task.mTaskId);
             if (entry != null) {
                 mAppIdMap.remove(entry.topApp);