Unlink to death before removing the callback wrapper in

SmartspaceService/AppPredictionServic3e

Bug: 232548271
Test: ahat

Change-Id: Ia86343721c4c8c57e52dd7d3e76de0eba6809058

2 files changed, 41 insertions, 23 deletions

diff --git a/core/java/android/service/appprediction/AppPredictionService.java b/core/java/android/service/appprediction/AppPredictionService.java
index 2d8aee567010..4f37cd91b11f 100644
--- a/core/java/android/service/appprediction/AppPredictionService.java
+++ b/core/java/android/service/appprediction/AppPredictionService.java
@@ -224,17 +224,16 @@ public abstract class AppPredictionService extends Service {
         }
 
         final CallbackWrapper wrapper = findCallbackWrapper(callbacks, callback);
-        if (wrapper != null) {
-            removeCallbackWrapper(callbacks, wrapper);
-        }
+        removeCallbackWrapper(callbacks, wrapper);
     }
 
-    private void removeCallbackWrapper(
-                ArrayList<CallbackWrapper> callbacks, CallbackWrapper wrapper) {
-        if (callbacks == null) {
+    private void removeCallbackWrapper(@Nullable ArrayList<CallbackWrapper> callbacks,
+            @Nullable CallbackWrapper wrapper) {
+        if (callbacks == null || wrapper == null) {
             return;
         }
         callbacks.remove(wrapper);
+        wrapper.destroy();
         if (callbacks.isEmpty()) {
             onStopPredictionUpdates();
         }
@@ -264,7 +263,8 @@ public abstract class AppPredictionService extends Service {
     public abstract void onRequestPredictionUpdate(@NonNull AppPredictionSessionId sessionId);
 
     private void doDestroyPredictionSession(@NonNull AppPredictionSessionId sessionId) {
-        mSessionCallbacks.remove(sessionId);
+        final ArrayList<CallbackWrapper> callbacks = mSessionCallbacks.remove(sessionId);
+        if (callbacks != null) callbacks.forEach(CallbackWrapper::destroy);
         onDestroyPredictionSession(sessionId);
     }
 
@@ -314,10 +314,12 @@ public abstract class AppPredictionService extends Service {
                 @Nullable Consumer<CallbackWrapper> onBinderDied) {
             mCallback = callback;
             mOnBinderDied = onBinderDied;
-            try {
-                mCallback.asBinder().linkToDeath(this, 0);
-            } catch (RemoteException e) {
-                Slog.e(TAG, "Failed to link to death: " + e);
+            if (mOnBinderDied != null) {
+                try {
+                    mCallback.asBinder().linkToDeath(this, 0);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Failed to link to death: " + e);
+                }
             }
         }
 
@@ -329,6 +331,12 @@ public abstract class AppPredictionService extends Service {
             return mCallback.equals(callback);
         }
 
+        public void destroy() {
+            if (mCallback != null && mOnBinderDied != null) {
+                mCallback.asBinder().unlinkToDeath(this, 0);
+            }
+        }
+
         @Override
         public void accept(List<AppTarget> ts) {
             try {
@@ -342,6 +350,7 @@ public abstract class AppPredictionService extends Service {
 
         @Override
         public void binderDied() {
+            destroy();
             mCallback = null;
             if (mOnBinderDied != null) {
                 mOnBinderDied.accept(this);
diff --git a/core/java/android/service/smartspace/SmartspaceService.java b/core/java/android/service/smartspace/SmartspaceService.java
index b903fbeb035c..3a148dffe6d6 100644
--- a/core/java/android/service/smartspace/SmartspaceService.java
+++ b/core/java/android/service/smartspace/SmartspaceService.java
@@ -202,9 +202,7 @@ public abstract class SmartspaceService extends Service {
         }
 
         final CallbackWrapper wrapper = findCallbackWrapper(callbacks, callback);
-        if (wrapper != null) {
-            removeCallbackWrapper(callbacks, wrapper);
-        }
+        removeCallbackWrapper(callbacks, wrapper);
     }
 
     private void doRequestPredictionUpdate(@NonNull SmartspaceSessionId sessionId) {
@@ -231,12 +229,13 @@ public abstract class SmartspaceService extends Service {
         return null;
     }
 
-    private void removeCallbackWrapper(
-            ArrayList<CallbackWrapper> callbacks, CallbackWrapper wrapper) {
-        if (callbacks == null) {
+    private void removeCallbackWrapper(@Nullable ArrayList<CallbackWrapper> callbacks,
+            @Nullable CallbackWrapper wrapper) {
+        if (callbacks == null || wrapper == null) {
             return;
         }
         callbacks.remove(wrapper);
+        wrapper.destroy();
     }
 
     /**
@@ -249,7 +248,9 @@ public abstract class SmartspaceService extends Service {
             Log.d(TAG, "doDestroy mSessionCallbacks: " + mSessionCallbacks);
         }
         super.onDestroy();
-        mSessionCallbacks.remove(sessionId);
+
+        final ArrayList<CallbackWrapper> callbacks = mSessionCallbacks.remove(sessionId);
+        if (callbacks != null) callbacks.forEach(CallbackWrapper::destroy);
         onDestroySmartspaceSession(sessionId);
     }
 
@@ -287,10 +288,12 @@ public abstract class SmartspaceService extends Service {
                 @Nullable Consumer<CallbackWrapper> onBinderDied) {
             mCallback = callback;
             mOnBinderDied = onBinderDied;
-            try {
-                mCallback.asBinder().linkToDeath(this, 0);
-            } catch (RemoteException e) {
-                Slog.e(TAG, "Failed to link to death: " + e);
+            if (mOnBinderDied != null) {
+                try {
+                    mCallback.asBinder().linkToDeath(this, 0);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Failed to link to death: " + e);
+                }
             }
         }
 
@@ -317,9 +320,15 @@ public abstract class SmartspaceService extends Service {
             }
         }
 
+        public void destroy() {
+            if (mCallback != null && mOnBinderDied != null) {
+                mCallback.asBinder().unlinkToDeath(this, 0);
+            }
+        }
+
         @Override
         public void binderDied() {
-            mCallback.asBinder().unlinkToDeath(this, 0);
+            destroy();
             mCallback = null;
             if (mOnBinderDied != null) {
                 mOnBinderDied.accept(this);