Workaround for app predictor callback leaking

Test: atest ResolverAppPredictorCallbackTest
Bug: 290971946

Change-Id: If402e752b1eff3b19b6222ae3a1b2e2ef619f9f0

5 files changed, 205 insertions, 23 deletions

diff --git a/core/java/com/android/internal/app/AppPredictionServiceResolverComparator.java b/core/java/com/android/internal/app/AppPredictionServiceResolverComparator.java
index b9f02365bbe7..d2fdc65b2c36 100644
--- a/core/java/com/android/internal/app/AppPredictionServiceResolverComparator.java
+++ b/core/java/com/android/internal/app/AppPredictionServiceResolverComparator.java
@@ -61,6 +61,8 @@ class AppPredictionServiceResolverComparator extends AbstractResolverComparator
     private final ModelBuilder mModelBuilder;
     private ResolverComparatorModel mComparatorModel;
 
+    private ResolverAppPredictorCallback mSortingCallback;
+
     // If this is non-null (and this is not destroyed), it means APS is disabled and we should fall
     // back to using the ResolverRankerService.
     // TODO: responsibility for this fallback behavior can live outside of the AppPrediction client.
@@ -94,6 +96,9 @@ class AppPredictionServiceResolverComparator extends AbstractResolverComparator
             // TODO: may not be necessary to build a new model, since we're destroying anyways.
             mComparatorModel = mModelBuilder.buildFallbackModel(mResolverRankerService);
         }
+        if (mSortingCallback != null) {
+            mSortingCallback.destroy();
+        }
     }
 
     @Override
@@ -140,22 +145,27 @@ class AppPredictionServiceResolverComparator extends AbstractResolverComparator
                     .setClassName(target.name.getClassName())
                     .build());
         }
+
+        if (mSortingCallback != null) {
+            mSortingCallback.destroy();
+        }
+        mSortingCallback = new ResolverAppPredictorCallback(sortedAppTargets -> {
+            if (sortedAppTargets.isEmpty()) {
+                Log.i(TAG, "AppPredictionService disabled. Using resolver.");
+                setupFallbackModel(targets);
+            } else {
+                Log.i(TAG, "AppPredictionService response received");
+                // Skip sending to Handler which takes extra time to dispatch messages.
+                // TODO: the Handler guards some concurrency conditions, so this could
+                // probably result in a race (we're not currently on the Handler thread?).
+                // We'll leave this as-is since we intend to remove the Handler design
+                // shortly, but this is still an unsound shortcut.
+                handleResult(sortedAppTargets);
+            }
+        });
+
         mAppPredictor.sortTargets(appTargets, Executors.newSingleThreadExecutor(),
-                sortedAppTargets -> {
-                    if (sortedAppTargets.isEmpty()) {
-                        Log.i(TAG, "AppPredictionService disabled. Using resolver.");
-                        setupFallbackModel(targets);
-                    } else {
-                        Log.i(TAG, "AppPredictionService response received");
-                        // Skip sending to Handler which takes extra time to dispatch messages.
-                        // TODO: the Handler guards some concurrency conditions, so this could
-                        // probably result in a race (we're not currently on the Handler thread?).
-                        // We'll leave this as-is since we intend to remove the Handler design
-                        // shortly, but this is still an unsound shortcut.
-                        handleResult(sortedAppTargets);
-                    }
-                }
-        );
+                mSortingCallback.asConsumer());
     }
 
     private void setupFallbackModel(List<ResolvedComponentInfo> targets) {
diff --git a/core/java/com/android/internal/app/ChooserActivity.java b/core/java/com/android/internal/app/ChooserActivity.java
index 2b39bb4eb7a5..6d8512c9d07c 100644
--- a/core/java/com/android/internal/app/ChooserActivity.java
+++ b/core/java/com/android/internal/app/ChooserActivity.java
@@ -24,9 +24,7 @@ import static android.app.admin.DevicePolicyResources.Strings.Core.RESOLVER_CROS
 import static android.content.ContentProvider.getUserIdFromUri;
 import static android.stats.devicepolicy.DevicePolicyEnums.RESOLVER_EMPTY_STATE_NO_SHARING_TO_PERSONAL;
 import static android.stats.devicepolicy.DevicePolicyEnums.RESOLVER_EMPTY_STATE_NO_SHARING_TO_WORK;
-
 import static com.android.internal.util.LatencyTracker.ACTION_LOAD_SHARE_SHEET;
-
 import static java.lang.annotation.RetentionPolicy.SOURCE;
 
 import android.animation.Animator;
@@ -777,9 +775,9 @@ public class ChooserActivity extends ResolverActivity implements
         return appPredictor;
     }
 
-    private AppPredictor.Callback createAppPredictorCallback(
+    private ResolverAppPredictorCallback createAppPredictorCallback(
             ChooserListAdapter chooserListAdapter) {
-        return resultList -> {
+        return new ResolverAppPredictorCallback(resultList -> {
             if (isFinishing() || isDestroyed()) {
                 return;
             }
@@ -811,7 +809,7 @@ public class ChooserActivity extends ResolverActivity implements
             }
             sendShareShortcutInfoList(shareShortcutInfos, chooserListAdapter, resultList,
                     chooserListAdapter.getUserHandle());
-        };
+        });
     }
 
     static SharedPreferences getPinnedSharedPrefs(Context context) {
@@ -2559,10 +2557,13 @@ public class ChooserActivity extends ResolverActivity implements
             boolean filterLastUsed, UserHandle userHandle) {
         ChooserListAdapter chooserListAdapter = createChooserListAdapter(context, payloadIntents,
                 initialIntents, rList, filterLastUsed, userHandle);
-        AppPredictor.Callback appPredictorCallback = createAppPredictorCallback(chooserListAdapter);
+        ResolverAppPredictorCallback appPredictorCallbackWrapper =
+                createAppPredictorCallback(chooserListAdapter);
+        AppPredictor.Callback appPredictorCallback = appPredictorCallbackWrapper.asCallback();
         AppPredictor appPredictor = setupAppPredictorForUser(userHandle, appPredictorCallback);
         chooserListAdapter.setAppPredictor(appPredictor);
-        chooserListAdapter.setAppPredictorCallback(appPredictorCallback);
+        chooserListAdapter.setAppPredictorCallback(
+                appPredictorCallback, appPredictorCallbackWrapper);
         return new ChooserGridAdapter(chooserListAdapter);
     }
 
diff --git a/core/java/com/android/internal/app/ChooserListAdapter.java b/core/java/com/android/internal/app/ChooserListAdapter.java
index 1eecb413adcb..f77e71863125 100644
--- a/core/java/com/android/internal/app/ChooserListAdapter.java
+++ b/core/java/com/android/internal/app/ChooserListAdapter.java
@@ -103,6 +103,7 @@ public class ChooserListAdapter extends ResolverListAdapter {
     // Sorted list of DisplayResolveInfos for the alphabetical app section.
     private List<DisplayResolveInfo> mSortedList = new ArrayList<>();
     private AppPredictor mAppPredictor;
+    private ResolverAppPredictorCallback mAppPredictorCallbackWrapper;
     private AppPredictor.Callback mAppPredictorCallback;
 
     // Represents the UserSpace in which the Initial Intents should be resolved.
@@ -747,8 +748,11 @@ public class ChooserListAdapter extends ResolverListAdapter {
         mAppPredictor = appPredictor;
     }
 
-    public void setAppPredictorCallback(AppPredictor.Callback appPredictorCallback) {
+    public void setAppPredictorCallback(
+            AppPredictor.Callback appPredictorCallback,
+            ResolverAppPredictorCallback appPredictorCallbackWrapper) {
         mAppPredictorCallback = appPredictorCallback;
+        mAppPredictorCallbackWrapper = appPredictorCallbackWrapper;
     }
 
     public void destroyAppPredictor() {
@@ -757,6 +761,10 @@ public class ChooserListAdapter extends ResolverListAdapter {
             getAppPredictor().destroy();
             setAppPredictor(null);
         }
+
+        if (mAppPredictorCallbackWrapper != null) {
+            mAppPredictorCallbackWrapper.destroy();
+        }
     }
 
     /**
diff --git a/core/java/com/android/internal/app/ResolverAppPredictorCallback.java b/core/java/com/android/internal/app/ResolverAppPredictorCallback.java
new file mode 100644
index 000000000000..c35e536275f0
--- /dev/null
+++ b/core/java/com/android/internal/app/ResolverAppPredictorCallback.java
@@ -0,0 +1,55 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.internal.app;
+
+import android.app.prediction.AppPredictor;
+import android.app.prediction.AppTarget;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.function.Consumer;
+
+/**
+ * Callback wrapper that works around potential memory leaks in app predictor.
+ *
+ * Nulls the callback itself when destroyed, so at worst you'll leak just this object.
+ */
+public class ResolverAppPredictorCallback {
+    private volatile Consumer<List<AppTarget>> mCallback;
+
+    public ResolverAppPredictorCallback(Consumer<List<AppTarget>> callback) {
+        mCallback = callback;
+    }
+
+    private void notifyCallback(List<AppTarget> list) {
+        Consumer<List<AppTarget>> callback = mCallback;
+        if (callback != null) {
+            callback.accept(Objects.requireNonNullElseGet(list, List::of));
+        }
+    }
+
+    public Consumer<List<AppTarget>> asConsumer() {
+        return this::notifyCallback;
+    }
+
+    public AppPredictor.Callback asCallback() {
+        return this::notifyCallback;
+    }
+
+    public void destroy() {
+        mCallback = null;
+    }
+}
diff --git a/core/tests/coretests/src/com/android/internal/app/ResolverAppPredictorCallbackTest.java b/core/tests/coretests/src/com/android/internal/app/ResolverAppPredictorCallbackTest.java
new file mode 100644
index 000000000000..4aca854469f2
--- /dev/null
+++ b/core/tests/coretests/src/com/android/internal/app/ResolverAppPredictorCallbackTest.java
@@ -0,0 +1,108 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.internal.app;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.app.prediction.AppTarget;
+import android.app.prediction.AppTargetId;
+import android.os.UserHandle;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.List;
+import java.util.function.Consumer;
+
+@RunWith(AndroidJUnit4.class)
+public class ResolverAppPredictorCallbackTest {
+    private class Callback implements Consumer<List<AppTarget>> {
+        public int count = 0;
+        public List<AppTarget> latest = null;
+        @Override
+        public void accept(List<AppTarget> appTargets) {
+            count++;
+            latest = appTargets;
+        }
+    };
+
+    @Test
+    public void testAsConsumer() {
+        Callback callback = new Callback();
+        ResolverAppPredictorCallback wrapped = new ResolverAppPredictorCallback(callback);
+        assertThat(callback.count).isEqualTo(0);
+
+        List<AppTarget> targets = createAppTargetList();
+        wrapped.asConsumer().accept(targets);
+
+        assertThat(callback.count).isEqualTo(1);
+        assertThat(callback.latest).isEqualTo(targets);
+
+        wrapped.destroy();
+
+        // Shouldn't do anything:
+        wrapped.asConsumer().accept(targets);
+
+        assertThat(callback.count).isEqualTo(1);
+    }
+
+    @Test
+    public void testAsCallback() {
+        Callback callback = new Callback();
+        ResolverAppPredictorCallback wrapped = new ResolverAppPredictorCallback(callback);
+        assertThat(callback.count).isEqualTo(0);
+
+        List<AppTarget> targets = createAppTargetList();
+        wrapped.asCallback().onTargetsAvailable(targets);
+
+        assertThat(callback.count).isEqualTo(1);
+        assertThat(callback.latest).isEqualTo(targets);
+
+        wrapped.destroy();
+
+        // Shouldn't do anything:
+        wrapped.asConsumer().accept(targets);
+
+        assertThat(callback.count).isEqualTo(1);
+    }
+
+    @Test
+    public void testAsConsumer_null() {
+        Callback callback = new Callback();
+        ResolverAppPredictorCallback wrapped = new ResolverAppPredictorCallback(callback);
+        assertThat(callback.count).isEqualTo(0);
+
+        wrapped.asConsumer().accept(null);
+
+        assertThat(callback.count).isEqualTo(1);
+        assertThat(callback.latest).isEmpty();
+
+        wrapped.destroy();
+
+        // Shouldn't do anything:
+        wrapped.asConsumer().accept(null);
+
+        assertThat(callback.count).isEqualTo(1);
+    }
+
+    private List<AppTarget> createAppTargetList() {
+        AppTarget.Builder builder = new AppTarget.Builder(
+                new AppTargetId("ID"), "package", UserHandle.CURRENT);
+        return List.of(builder.build());
+    }
+}