fix app prediction service doesn't clean-up properly.

Currently new app prediction session will be created each time
launcher crashes, this introduces some memory impact since the
existing callbacks are not released after the crash.

The root cause being random UUID was used to create the session,
as a result app prediction service is not able to dedupe it with
against existing sessions.

Bug: 161391868
Test: manual
Change-Id: I937cdaa6081d2b85d27ee0192f6af7f197b8e102

4 files changed, 92 insertions, 56 deletions

diff --git a/core/java/android/app/prediction/AppPredictor.java b/core/java/android/app/prediction/AppPredictor.java
index 7f436401dbf4..fa135b10ae1f 100644
--- a/core/java/android/app/prediction/AppPredictor.java
+++ b/core/java/android/app/prediction/AppPredictor.java
@@ -83,6 +83,8 @@ public final class AppPredictor {
     private final AppPredictionSessionId mSessionId;
     private final ArrayMap<Callback, CallbackWrapper> mRegisteredCallbacks = new ArrayMap<>();
 
+    private final IBinder mToken = new Binder();
+
     /**
      * Creates a new Prediction client.
      * <p>
@@ -98,7 +100,7 @@ public final class AppPredictor {
         mSessionId = new AppPredictionSessionId(
                 context.getPackageName() + ":" + UUID.randomUUID().toString(), context.getUserId());
         try {
-            mPredictionManager.createPredictionSession(predictionContext, mSessionId);
+            mPredictionManager.createPredictionSession(predictionContext, mSessionId, mToken);
         } catch (RemoteException e) {
             Log.e(TAG, "Failed to create predictor", e);
             e.rethrowAsRuntimeException();
diff --git a/core/java/android/app/prediction/IPredictionManager.aidl b/core/java/android/app/prediction/IPredictionManager.aidl
index 587e3fd52377..863fc6f952dd 100644
--- a/core/java/android/app/prediction/IPredictionManager.aidl
+++ b/core/java/android/app/prediction/IPredictionManager.aidl
@@ -29,7 +29,7 @@ import android.content.pm.ParceledListSlice;
 interface IPredictionManager {
 
     void createPredictionSession(in AppPredictionContext context,
-            in AppPredictionSessionId sessionId);
+            in AppPredictionSessionId sessionId, in IBinder token);
 
     void notifyAppTargetEvent(in AppPredictionSessionId sessionId, in AppTargetEvent event);
 
diff --git a/services/appprediction/java/com/android/server/appprediction/AppPredictionManagerService.java b/services/appprediction/java/com/android/server/appprediction/AppPredictionManagerService.java
index 1c4db1214d3b..59ba82e4616a 100644
--- a/services/appprediction/java/com/android/server/appprediction/AppPredictionManagerService.java
+++ b/services/appprediction/java/com/android/server/appprediction/AppPredictionManagerService.java
@@ -34,6 +34,7 @@ import android.app.prediction.IPredictionManager;
 import android.content.Context;
 import android.content.pm.ParceledListSlice;
 import android.os.Binder;
+import android.os.IBinder;
 import android.os.ResultReceiver;
 import android.os.ShellCallback;
 import android.util.Slog;
@@ -108,9 +109,9 @@ public class AppPredictionManagerService extends
 
         @Override
         public void createPredictionSession(@NonNull AppPredictionContext context,
-                @NonNull AppPredictionSessionId sessionId) {
-            runForUserLocked("createPredictionSession", sessionId,
-                    (service) -> service.onCreatePredictionSessionLocked(context, sessionId));
+                @NonNull AppPredictionSessionId sessionId, @NonNull IBinder token) {
+            runForUserLocked("createPredictionSession", sessionId, (service) ->
+                    service.onCreatePredictionSessionLocked(context, sessionId, token));
         }
 
         @Override
diff --git a/services/appprediction/java/com/android/server/appprediction/AppPredictionPerUserService.java b/services/appprediction/java/com/android/server/appprediction/AppPredictionPerUserService.java
index 103151dcdda5..f9cd60a43524 100644
--- a/services/appprediction/java/com/android/server/appprediction/AppPredictionPerUserService.java
+++ b/services/appprediction/java/com/android/server/appprediction/AppPredictionPerUserService.java
@@ -30,6 +30,7 @@ import android.content.pm.PackageManager;
 import android.content.pm.PackageManager.NameNotFoundException;
 import android.content.pm.ParceledListSlice;
 import android.content.pm.ServiceInfo;
+import android.os.IBinder;
 import android.os.RemoteCallbackList;
 import android.os.RemoteException;
 import android.provider.DeviceConfig;
@@ -44,8 +45,6 @@ import com.android.server.LocalServices;
 import com.android.server.infra.AbstractPerUserSystemService;
 import com.android.server.people.PeopleServiceInternal;
 
-import java.util.function.Consumer;
-
 /**
  * Per-user instance of {@link AppPredictionManagerService}.
  */
@@ -112,17 +111,24 @@ public class AppPredictionPerUserService extends
      */
     @GuardedBy("mLock")
     public void onCreatePredictionSessionLocked(@NonNull AppPredictionContext context,
-            @NonNull AppPredictionSessionId sessionId) {
-        if (!mSessionInfos.containsKey(sessionId)) {
-            mSessionInfos.put(sessionId, new AppPredictionSessionInfo(sessionId, context,
-                    DeviceConfig.getBoolean(NAMESPACE_SYSTEMUI,
-                            PREDICT_USING_PEOPLE_SERVICE_PREFIX + context.getUiSurface(), false),
-                    this::removeAppPredictionSessionInfo));
-        }
-        final boolean serviceExists = resolveService(sessionId, s ->
-                s.onCreatePredictionSession(context, sessionId), true);
-        if (!serviceExists) {
-            mSessionInfos.remove(sessionId);
+            @NonNull AppPredictionSessionId sessionId, @NonNull IBinder token) {
+        final boolean usesPeopleService = DeviceConfig.getBoolean(NAMESPACE_SYSTEMUI,
+                PREDICT_USING_PEOPLE_SERVICE_PREFIX + context.getUiSurface(), false);
+        final boolean serviceExists = resolveService(sessionId, false,
+                usesPeopleService, s -> s.onCreatePredictionSession(context, sessionId));
+        if (serviceExists && !mSessionInfos.containsKey(sessionId)) {
+            final AppPredictionSessionInfo sessionInfo = new AppPredictionSessionInfo(
+                    sessionId, context, usesPeopleService, token, () -> {
+                synchronized (mLock) {
+                    onDestroyPredictionSessionLocked(sessionId);
+                }
+            });
+            if (sessionInfo.linkToDeath()) {
+                mSessionInfos.put(sessionId, sessionInfo);
+            } else {
+                // destroy the session if calling process is already dead
+                onDestroyPredictionSessionLocked(sessionId);
+            }
         }
     }
 
@@ -132,7 +138,10 @@ public class AppPredictionPerUserService extends
     @GuardedBy("mLock")
     public void notifyAppTargetEventLocked(@NonNull AppPredictionSessionId sessionId,
             @NonNull AppTargetEvent event) {
-        resolveService(sessionId, s -> s.notifyAppTargetEvent(sessionId, event), false);
+        final AppPredictionSessionInfo sessionInfo = mSessionInfos.get(sessionId);
+        if (sessionInfo == null) return;
+        resolveService(sessionId, false, sessionInfo.mUsesPeopleService,
+                s -> s.notifyAppTargetEvent(sessionId, event));
     }
 
     /**
@@ -141,8 +150,10 @@ public class AppPredictionPerUserService extends
     @GuardedBy("mLock")
     public void notifyLaunchLocationShownLocked(@NonNull AppPredictionSessionId sessionId,
             @NonNull String launchLocation, @NonNull ParceledListSlice targetIds) {
-        resolveService(sessionId, s ->
-                s.notifyLaunchLocationShown(sessionId, launchLocation, targetIds), false);
+        final AppPredictionSessionInfo sessionInfo = mSessionInfos.get(sessionId);
+        if (sessionInfo == null) return;
+        resolveService(sessionId, false, sessionInfo.mUsesPeopleService,
+                s -> s.notifyLaunchLocationShown(sessionId, launchLocation, targetIds));
     }
 
     /**
@@ -151,7 +162,10 @@ public class AppPredictionPerUserService extends
     @GuardedBy("mLock")
     public void sortAppTargetsLocked(@NonNull AppPredictionSessionId sessionId,
             @NonNull ParceledListSlice targets, @NonNull IPredictionCallback callback) {
-        resolveService(sessionId, s -> s.sortAppTargets(sessionId, targets, callback), true);
+        final AppPredictionSessionInfo sessionInfo = mSessionInfos.get(sessionId);
+        if (sessionInfo == null) return;
+        resolveService(sessionId, true, sessionInfo.mUsesPeopleService,
+                s -> s.sortAppTargets(sessionId, targets, callback));
     }
 
     /**
@@ -160,10 +174,12 @@ public class AppPredictionPerUserService extends
     @GuardedBy("mLock")
     public void registerPredictionUpdatesLocked(@NonNull AppPredictionSessionId sessionId,
             @NonNull IPredictionCallback callback) {
-        final boolean serviceExists = resolveService(sessionId, s ->
-                s.registerPredictionUpdates(sessionId, callback), false);
         final AppPredictionSessionInfo sessionInfo = mSessionInfos.get(sessionId);
-        if (serviceExists && sessionInfo != null) {
+        if (sessionInfo == null) return;
+        final boolean serviceExists = resolveService(sessionId, false,
+                sessionInfo.mUsesPeopleService,
+                s -> s.registerPredictionUpdates(sessionId, callback));
+        if (serviceExists) {
             sessionInfo.addCallbackLocked(callback);
         }
     }
@@ -174,10 +190,12 @@ public class AppPredictionPerUserService extends
     @GuardedBy("mLock")
     public void unregisterPredictionUpdatesLocked(@NonNull AppPredictionSessionId sessionId,
             @NonNull IPredictionCallback callback) {
-        final boolean serviceExists = resolveService(sessionId, s ->
-                s.unregisterPredictionUpdates(sessionId, callback), false);
         final AppPredictionSessionInfo sessionInfo = mSessionInfos.get(sessionId);
-        if (serviceExists && sessionInfo != null) {
+        if (sessionInfo == null) return;
+        final boolean serviceExists = resolveService(sessionId, false,
+                sessionInfo.mUsesPeopleService,
+                s -> s.unregisterPredictionUpdates(sessionId, callback));
+        if (serviceExists) {
             sessionInfo.removeCallbackLocked(callback);
         }
     }
@@ -187,7 +205,10 @@ public class AppPredictionPerUserService extends
      */
     @GuardedBy("mLock")
     public void requestPredictionUpdateLocked(@NonNull AppPredictionSessionId sessionId) {
-        resolveService(sessionId, s -> s.requestPredictionUpdate(sessionId), true);
+        final AppPredictionSessionInfo sessionInfo = mSessionInfos.get(sessionId);
+        if (sessionInfo == null) return;
+        resolveService(sessionId, true, sessionInfo.mUsesPeopleService,
+                s -> s.requestPredictionUpdate(sessionId));
     }
 
     /**
@@ -195,12 +216,14 @@ public class AppPredictionPerUserService extends
      */
     @GuardedBy("mLock")
     public void onDestroyPredictionSessionLocked(@NonNull AppPredictionSessionId sessionId) {
-        final boolean serviceExists = resolveService(sessionId, s ->
-                s.onDestroyPredictionSession(sessionId), false);
-        final AppPredictionSessionInfo sessionInfo = mSessionInfos.get(sessionId);
-        if (serviceExists && sessionInfo != null) {
-            sessionInfo.destroy();
+        if (isDebug()) {
+            Slog.d(TAG, "onDestroyPredictionSessionLocked(): sessionId=" + sessionId);
         }
+        final AppPredictionSessionInfo sessionInfo = mSessionInfos.remove(sessionId);
+        if (sessionInfo == null) return;
+        resolveService(sessionId, false, sessionInfo.mUsesPeopleService,
+                s -> s.onDestroyPredictionSession(sessionId));
+        sessionInfo.destroy();
     }
 
     @Override
@@ -291,27 +314,18 @@ public class AppPredictionPerUserService extends
         }
 
         for (AppPredictionSessionInfo sessionInfo : mSessionInfos.values()) {
-            sessionInfo.resurrectSessionLocked(this);
-        }
-    }
-
-    private void removeAppPredictionSessionInfo(AppPredictionSessionId sessionId) {
-        if (isDebug()) {
-            Slog.d(TAG, "removeAppPredictionSessionInfo(): sessionId=" + sessionId);
-        }
-        synchronized (mLock) {
-            mSessionInfos.remove(sessionId);
+            sessionInfo.resurrectSessionLocked(this, sessionInfo.mToken);
         }
     }
 
     @GuardedBy("mLock")
     @Nullable
-    protected boolean resolveService(@NonNull final AppPredictionSessionId sessionId,
-            @NonNull final AbstractRemoteService.AsyncRequest<IPredictionService> cb,
-            boolean sendImmediately) {
-        final AppPredictionSessionInfo sessionInfo = mSessionInfos.get(sessionId);
-        if (sessionInfo == null) return false;
-        if (sessionInfo.mUsesPeopleService) {
+    protected boolean resolveService(
+            @NonNull final AppPredictionSessionId sessionId,
+            boolean sendImmediately,
+            boolean usesPeopleService,
+            @NonNull final AbstractRemoteService.AsyncRequest<IPredictionService> cb) {
+        if (usesPeopleService) {
             final IPredictionService service =
                     LocalServices.getService(PeopleServiceInternal.class);
             if (service != null) {
@@ -368,7 +382,9 @@ public class AppPredictionPerUserService extends
         private final AppPredictionContext mPredictionContext;
         private final boolean mUsesPeopleService;
         @NonNull
-        private final Consumer<AppPredictionSessionId> mRemoveSessionInfoAction;
+        final IBinder mToken;
+        @NonNull
+        final IBinder.DeathRecipient mDeathRecipient;
 
         private final RemoteCallbackList<IPredictionCallback> mCallbacks =
                 new RemoteCallbackList<IPredictionCallback>() {
@@ -388,14 +404,16 @@ public class AppPredictionPerUserService extends
                 @NonNull final AppPredictionSessionId id,
                 @NonNull final AppPredictionContext predictionContext,
                 final boolean usesPeopleService,
-                @NonNull final Consumer<AppPredictionSessionId> removeSessionInfoAction) {
+                @NonNull final IBinder token,
+                @NonNull final IBinder.DeathRecipient deathRecipient) {
             if (DEBUG) {
                 Slog.d(TAG, "Creating AppPredictionSessionInfo for session Id=" + id);
             }
             mSessionId = id;
             mPredictionContext = predictionContext;
             mUsesPeopleService = usesPeopleService;
-            mRemoveSessionInfoAction = removeSessionInfoAction;
+            mToken = token;
+            mDeathRecipient = deathRecipient;
         }
 
         void addCallbackLocked(IPredictionCallback callback) {
@@ -414,23 +432,38 @@ public class AppPredictionPerUserService extends
             mCallbacks.unregister(callback);
         }
 
+        boolean linkToDeath() {
+            try {
+                mToken.linkToDeath(mDeathRecipient, 0);
+            } catch (RemoteException e) {
+                if (DEBUG) {
+                    Slog.w(TAG, "Caller is dead before session can be started, sessionId: "
+                            + mSessionId);
+                }
+                return false;
+            }
+            return true;
+        }
+
         void destroy() {
             if (DEBUG) {
                 Slog.d(TAG, "Removing all callbacks for session Id=" + mSessionId
                         + " and " + mCallbacks.getRegisteredCallbackCount() + " callbacks.");
             }
+            if (mToken != null) {
+                mToken.unlinkToDeath(mDeathRecipient, 0);
+            }
             mCallbacks.kill();
-            mRemoveSessionInfoAction.accept(mSessionId);
         }
 
-        void resurrectSessionLocked(AppPredictionPerUserService service) {
+        void resurrectSessionLocked(AppPredictionPerUserService service, IBinder token) {
             int callbackCount = mCallbacks.getRegisteredCallbackCount();
             if (DEBUG) {
                 Slog.d(TAG, "Resurrecting remote service (" + service.getRemoteServiceLocked()
                         + ") for session Id=" + mSessionId + " and "
                         + callbackCount + " callbacks.");
             }
-            service.onCreatePredictionSessionLocked(mPredictionContext, mSessionId);
+            service.onCreatePredictionSessionLocked(mPredictionContext, mSessionId, token);
             mCallbacks.broadcast(
                     callback -> service.registerPredictionUpdatesLocked(mSessionId, callback));
         }