Merge cherrypicks of ['googleplex-android-review.googlesource.com/29573568'] into 24Q4-release.

Change-Id: I32f46d6425e475f8cdc4b76c7ad4d6a00ca2b100
diff --git a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
index 531cdfd..27ec838 100644
--- a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
+++ b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
@@ -30,14 +30,11 @@
 import android.os.Binder;
 import android.os.Bundle;
 import android.os.Handler;
-import android.os.IBinder;
 import android.os.IBinder.DeathRecipient;
 import android.os.Message;
 import android.os.Messenger;
 import android.text.TextUtils;
-import android.util.ArrayMap;
 import android.util.Log;
-import android.util.Pair;
 
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile.GridOption;
@@ -47,8 +44,12 @@
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.Preconditions;
+import com.android.launcher3.util.RunnableList;
 import com.android.systemui.shared.Flags;
 
+import java.util.Collections;
+import java.util.Set;
+import java.util.WeakHashMap;
 import java.util.concurrent.ExecutionException;
 
 /**
@@ -95,11 +96,9 @@
     private static final int MESSAGE_ID_UPDATE_PREVIEW = 1337;
     private static final int MESSAGE_ID_UPDATE_GRID = 7414;
 
-    /**
-     * Here we use the IBinder and the screen ID as the key of the active previews.
-     */
-    private final ArrayMap<Pair<IBinder, Integer>, PreviewLifecycleObserver> mActivePreviews =
-            new ArrayMap<>();
+    // Set of all active previews used to track duplicate memory allocations
+    private final Set<PreviewLifecycleObserver> mActivePreviews =
+            Collections.newSetFromMap(new WeakHashMap<>());
 
     @Override
     public boolean onCreate() {
@@ -231,16 +230,19 @@
     }
 
     private synchronized Bundle getPreview(Bundle request) {
-        PreviewLifecycleObserver observer = null;
+        RunnableList lifeCycleTracker = new RunnableList();
         try {
-            PreviewSurfaceRenderer renderer = new PreviewSurfaceRenderer(getContext(), request);
+            PreviewSurfaceRenderer renderer = new PreviewSurfaceRenderer(
+                    getContext(), lifeCycleTracker, request);
+            PreviewLifecycleObserver observer =
+                    new PreviewLifecycleObserver(lifeCycleTracker, renderer);
 
-            observer = new PreviewLifecycleObserver(renderer);
-            // Destroy previous
-            destroyObserver(mActivePreviews.get(observer.getIdentifier()));
-            mActivePreviews.put(observer.getIdentifier(), observer);
+            // Destroy previous renderers to avoid any duplicate memory
+            mActivePreviews.stream().filter(observer::isSameRenderer).forEach(o ->
+                    MAIN_EXECUTOR.execute(o.lifeCycleTracker::executeAllAndDestroy));
 
             renderer.loadAsync();
+            lifeCycleTracker.add(() -> renderer.getHostToken().unlinkToDeath(observer, 0));
             renderer.getHostToken().linkToDeath(observer, 0);
 
             Bundle result = new Bundle();
@@ -254,33 +256,21 @@
             return result;
         } catch (Exception e) {
             Log.e(TAG, "Unable to generate preview", e);
-            if (observer != null) {
-                destroyObserver(observer);
-            }
+            MAIN_EXECUTOR.execute(lifeCycleTracker::executeAllAndDestroy);
             return null;
         }
     }
 
-    private synchronized void destroyObserver(PreviewLifecycleObserver observer) {
-        if (observer == null || observer.destroyed) {
-            return;
-        }
-        observer.destroyed = true;
-        observer.renderer.getHostToken().unlinkToDeath(observer, 0);
-        MAIN_EXECUTOR.execute(observer.renderer::destroy);
-        PreviewLifecycleObserver cached = mActivePreviews.get(observer.getIdentifier());
-        if (cached == observer) {
-            mActivePreviews.remove(observer.getIdentifier());
-        }
-    }
+    private static class PreviewLifecycleObserver implements Handler.Callback, DeathRecipient {
 
-    private class PreviewLifecycleObserver implements Handler.Callback, DeathRecipient {
-
+        public final RunnableList lifeCycleTracker;
         public final PreviewSurfaceRenderer renderer;
         public boolean destroyed = false;
 
-        PreviewLifecycleObserver(PreviewSurfaceRenderer renderer) {
+        PreviewLifecycleObserver(RunnableList lifeCycleTracker, PreviewSurfaceRenderer renderer) {
+            this.lifeCycleTracker = lifeCycleTracker;
             this.renderer = renderer;
+            lifeCycleTracker.add(() -> destroyed = true);
         }
 
         @Override
@@ -300,7 +290,9 @@
                     }
                     break;
                 default:
-                    destroyObserver(this);
+                    // Unknown command, destroy lifecycle
+                    Log.d(TAG, "Unknown preview command: " + message.what + ", destroying preview");
+                    MAIN_EXECUTOR.execute(lifeCycleTracker::executeAllAndDestroy);
                     break;
             }
 
@@ -309,16 +301,16 @@
 
         @Override
         public void binderDied() {
-            destroyObserver(this);
+            MAIN_EXECUTOR.execute(lifeCycleTracker::executeAllAndDestroy);
         }
 
         /**
-         * Returns a key that should make the PreviewSurfaceRenderer unique and if two of them have
-         * the same key they will be treated as the same PreviewSurfaceRenderer. Primary this is
-         * used to prevent memory leaks by removing the old PreviewSurfaceRenderer.
+         * Two renderers are considered same if they have the same host token and display Id
          */
-        public Pair<IBinder, Integer> getIdentifier() {
-            return new Pair<>(renderer.getHostToken(), renderer.getDisplayId());
+        public boolean isSameRenderer(PreviewLifecycleObserver plo) {
+            return plo != null
+                    && plo.renderer.getHostToken().equals(renderer.getHostToken())
+                    && plo.renderer.getDisplayId() == renderer.getDisplayId();
         }
     }
 }
diff --git a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
index 56c4ca4..1b23d75 100644
--- a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
+++ b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
@@ -91,7 +91,7 @@
     private final int mDisplayId;
     private final Display mDisplay;
     private final WallpaperColors mWallpaperColors;
-    private final RunnableList mOnDestroyCallbacks = new RunnableList();
+    private final RunnableList mLifeCycleTracker;
 
     private final SurfaceControlViewHost mSurfaceControlViewHost;
 
@@ -100,8 +100,10 @@
     private boolean mHideQsb;
     @Nullable private FrameLayout mViewRoot = null;
 
-    public PreviewSurfaceRenderer(Context context, Bundle bundle) throws Exception {
+    public PreviewSurfaceRenderer(
+            Context context, RunnableList lifecycleTracker, Bundle bundle) throws Exception {
         mContext = context;
+        mLifeCycleTracker = lifecycleTracker;
         mGridName = bundle.getString("name");
         bundle.remove("name");
         if (mGridName == null) {
@@ -120,11 +122,13 @@
             throw new IllegalArgumentException("Display ID does not match any displays.");
         }
 
-        mSurfaceControlViewHost = MAIN_EXECUTOR.submit(() ->
-                new SurfaceControlViewHost(mContext, context.getSystemService(DisplayManager.class)
-                        .getDisplay(DEFAULT_DISPLAY), mHostToken)
-        ).get(5, TimeUnit.SECONDS);
-        mOnDestroyCallbacks.add(mSurfaceControlViewHost::release);
+        mSurfaceControlViewHost = MAIN_EXECUTOR.submit(() -> new MySurfaceControlViewHost(
+                mContext,
+                context.getSystemService(DisplayManager.class).getDisplay(DEFAULT_DISPLAY),
+                mHostToken,
+                mLifeCycleTracker))
+                .get(5, TimeUnit.SECONDS);
+        mLifeCycleTracker.add(this::destroy);
     }
 
     public int getDisplayId() {
@@ -139,25 +143,18 @@
         return mSurfaceControlViewHost.getSurfacePackage();
     }
 
-    /**
-     * Destroys the preview and all associated data
-     */
-    @UiThread
-    public void destroy() {
+    private void destroy() {
         mDestroyed = true;
-        mOnDestroyCallbacks.executeAllAndDestroy();
     }
 
     /**
      * A function that queries for the launcher app widget span info
      *
-     * @param context The context to get the content resolver from, should be related to launcher
      * @return A SparseArray with the app widget id being the key and the span info being the values
      */
     @WorkerThread
     @Nullable
-    public SparseArray<Size> getLoadedLauncherWidgetInfo(
-            @NonNull final Context context) {
+    public SparseArray<Size> getLoadedLauncherWidgetInfo() {
         final SparseArray<Size> widgetInfo = new SparseArray<>();
         final String query = LauncherSettings.Favorites.ITEM_TYPE + " = "
                 + LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET;
@@ -276,13 +273,11 @@
                     }
                     loadWorkspace(new ArrayList<>(), query, null, null);
 
-                    final SparseArray<Size> spanInfo =
-                            getLoadedLauncherWidgetInfo(previewContext.getBaseContext());
-
+                    final SparseArray<Size> spanInfo = getLoadedLauncherWidgetInfo();
                     MAIN_EXECUTOR.execute(() -> {
                         renderView(previewContext, mBgDataModel, mWidgetProvidersMap, spanInfo,
                                 idp);
-                        mOnDestroyCallbacks.add(previewContext::onDestroy);
+                        mLifeCycleTracker.add(previewContext::onDestroy);
                     });
                 }
             }.run();
@@ -355,4 +350,24 @@
             mViewRoot.addView(view);
         }
     }
+
+    private static class MySurfaceControlViewHost extends SurfaceControlViewHost {
+
+        private final RunnableList mLifecycleTracker;
+
+        MySurfaceControlViewHost(Context context, Display display, IBinder hostToken,
+                RunnableList lifeCycleTracker) {
+            super(context, display, hostToken);
+            mLifecycleTracker = lifeCycleTracker;
+            mLifecycleTracker.add(this::release);
+        }
+
+        @Override
+        public void release() {
+            super.release();
+            // RunnableList ensures that the callback is only called once
+            MAIN_EXECUTOR.execute(mLifecycleTracker::executeAllAndDestroy);
+        }
+    }
+
 }