Introducing a wrapper over binder to simplify lifecycle management for IPC

Fix: 400733213
Flag: NONE - release code
Test: manual
Change-Id: I5e143fce6a4cbc1dbb246bc6e0141ebbbc8bd274
diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index 7cf0605..aae8a56 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -352,7 +352,7 @@
                 new RemoteAnimationAdapter(runner, duration, statusBarTransitionDelay),
                 new RemoteTransition(runner.toRemoteTransition(),
                         mLauncher.getIApplicationThread(), "QuickstepLaunch"));
-        IRemoteCallback endCallback = completeRunnableListCallback(onEndCallback);
+        IRemoteCallback endCallback = completeRunnableListCallback(onEndCallback, mLauncher);
         options.setOnAnimationAbortListener(endCallback);
         options.setOnAnimationFinishedListener(endCallback);
         options.setLaunchCookie(StableViewInfo.toLaunchCookie(itemInfo));
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index fc8ea87..6063ed3 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -919,7 +919,7 @@
         options.setSplashScreenStyle(splashScreenStyle);
         options.setPendingIntentBackgroundActivityStartMode(
                 ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED);
-        IRemoteCallback endCallback = completeRunnableListCallback(callbacks);
+        IRemoteCallback endCallback = completeRunnableListCallback(callbacks, this);
         options.setOnAnimationAbortListener(endCallback);
         options.setOnAnimationFinishedListener(endCallback);
 
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 1a42d21..cd0a4f3 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -1274,7 +1274,7 @@
         options.setPendingIntentBackgroundActivityStartMode(
                 ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED);
 
-        IRemoteCallback endCallback = completeRunnableListCallback(callbacks);
+        IRemoteCallback endCallback = completeRunnableListCallback(callbacks, this);
         options.setOnAnimationAbortListener(endCallback);
         options.setOnAnimationFinishedListener(endCallback);
         return new ActivityOptionsWrapper(options, callbacks);
diff --git a/quickstep/src/com/android/quickstep/util/AnimUtils.java b/quickstep/src/com/android/quickstep/util/AnimUtils.java
index fda0c29..3492788 100644
--- a/quickstep/src/com/android/quickstep/util/AnimUtils.java
+++ b/quickstep/src/com/android/quickstep/util/AnimUtils.java
@@ -21,16 +21,22 @@
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
 import android.animation.AnimatorSet;
-import android.annotation.NonNull;
+import android.os.BinderUtils;
 import android.os.Bundle;
+import android.os.IBinder;
 import android.os.IRemoteCallback;
 import android.view.animation.Interpolator;
 
+import androidx.annotation.NonNull;
+import androidx.lifecycle.DefaultLifecycleObserver;
+import androidx.lifecycle.LifecycleOwner;
+
 import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.statemanager.BaseState;
 import com.android.launcher3.statemanager.StateManager;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.util.RunnableList;
+import com.android.launcher3.views.ActivityContext;
 import com.android.quickstep.views.RecentsViewContainer;
 
 /**
@@ -95,14 +101,30 @@
     }
 
     /**
-     * Returns a IRemoteCallback which completes the provided list as a result
+     * Returns a IRemoteCallback which completes the provided list as a result or when the owner
+     * is destroyed
      */
-    public static IRemoteCallback completeRunnableListCallback(RunnableList list) {
+    public static IRemoteCallback completeRunnableListCallback(
+            RunnableList list, ActivityContext owner) {
+        DefaultLifecycleObserver destroyObserver = new DefaultLifecycleObserver() {
+            @Override
+            public void onDestroy(@NonNull LifecycleOwner owner) {
+                list.executeAllAndClear();
+            }
+        };
+        MAIN_EXECUTOR.execute(() -> owner.getLifecycle().addObserver(destroyObserver));
+        list.add(() -> owner.getLifecycle().removeObserver(destroyObserver));
+
         return new IRemoteCallback.Stub() {
             @Override
             public void sendResult(Bundle bundle) {
                 MAIN_EXECUTOR.execute(list::executeAllAndDestroy);
             }
+
+            @Override
+            public IBinder asBinder() {
+                return BinderUtils.wrapLifecycle(this, owner.getOwnerCleanupSet());
+            }
         };
     }
 
diff --git a/src/android/os/BinderUtils.kt b/src/android/os/BinderUtils.kt
new file mode 100644
index 0000000..0536283
--- /dev/null
+++ b/src/android/os/BinderUtils.kt
@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2025 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 android.os
+
+import com.android.launcher3.util.Executors.MAIN_EXECUTOR
+import com.android.launcher3.util.WeakCleanupSet
+import com.android.launcher3.util.WeakCleanupSet.OnOwnerDestroyedCallback
+
+/** Utility methods related to Binder */
+object BinderUtils {
+
+    /** Creates a binder wrapper which is tied to the [lifecycle] */
+    @JvmStatic
+    fun <T : Binder> T.wrapLifecycle(cleanupSet: WeakCleanupSet): Binder =
+        LifecycleBinderWrapper(this, cleanupSet)
+
+    private class LifecycleBinderWrapper<T : Binder>(
+        private var realBinder: T?,
+        cleanupSet: WeakCleanupSet,
+    ) : Binder(realBinder?.interfaceDescriptor), OnOwnerDestroyedCallback {
+
+        init {
+            MAIN_EXECUTOR.execute { cleanupSet.addOnOwnerDestroyedCallback(this) }
+        }
+
+        override fun queryLocalInterface(descriptor: String): IInterface? =
+            realBinder?.queryLocalInterface(descriptor)
+
+        override fun onTransact(code: Int, data: Parcel, reply: Parcel?, flags: Int): Boolean =
+            realBinder?.transact(code, data, reply, flags)
+                ?: throw RemoteException("Original binder cleaned up")
+
+        override fun onOwnerDestroyed() {
+            realBinder = null
+        }
+    }
+}
diff --git a/src/com/android/launcher3/BaseActivity.java b/src/com/android/launcher3/BaseActivity.java
index 3e6b4dd..2426a61 100644
--- a/src/com/android/launcher3/BaseActivity.java
+++ b/src/com/android/launcher3/BaseActivity.java
@@ -55,6 +55,7 @@
 import com.android.launcher3.util.RunnableList;
 import com.android.launcher3.util.SystemUiController;
 import com.android.launcher3.util.ViewCache;
+import com.android.launcher3.util.WeakCleanupSet;
 import com.android.launcher3.util.WindowBounds;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.ScrimView;
@@ -105,6 +106,7 @@
     private final SavedStateRegistryController mSavedStateRegistryController =
             SavedStateRegistryController.create(this);
     private final LifecycleRegistry mLifecycleRegistry = new LifecycleRegistry(this);
+    private final WeakCleanupSet mCleanupSet = new WeakCleanupSet(this);
 
     protected DeviceProfile mDeviceProfile;
     protected SystemUiController mSystemUiController;
@@ -504,6 +506,11 @@
         return mLifecycleRegistry;
     }
 
+    @Override
+    public WeakCleanupSet getOwnerCleanupSet() {
+        return mCleanupSet;
+    }
+
     public static <T extends BaseActivity> T fromContext(Context context) {
         if (context instanceof BaseActivity) {
             return (T) context;
diff --git a/src/com/android/launcher3/util/BaseContext.kt b/src/com/android/launcher3/util/BaseContext.kt
index 819470b..8aa10f3 100644
--- a/src/com/android/launcher3/util/BaseContext.kt
+++ b/src/com/android/launcher3/util/BaseContext.kt
@@ -49,6 +49,7 @@
 
     private val savedStateRegistryController = SavedStateRegistryController.create(this)
     private val lifecycleRegistry = LifecycleRegistry(this)
+    private val cleanupSet = WeakCleanupSet(this)
 
     override val savedStateRegistry: SavedStateRegistry
         get() = savedStateRegistryController.savedStateRegistry
@@ -110,6 +111,8 @@
 
     override fun getViewCache() = viewCache
 
+    override fun getOwnerCleanupSet() = cleanupSet
+
     private fun updateState() {
         if (lifecycleRegistry.currentState.isAtLeast(CREATED)) {
             lifecycleRegistry.currentState =
diff --git a/src/com/android/launcher3/util/WeakCleanupSet.kt b/src/com/android/launcher3/util/WeakCleanupSet.kt
new file mode 100644
index 0000000..7bf3289
--- /dev/null
+++ b/src/com/android/launcher3/util/WeakCleanupSet.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2025 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.launcher3.util
+
+import androidx.lifecycle.DefaultLifecycleObserver
+import androidx.lifecycle.LifecycleOwner
+import com.android.launcher3.util.Executors.MAIN_EXECUTOR
+import java.util.Collections
+import java.util.WeakHashMap
+
+/**
+ * Utility class which maintains a list of cleanup callbacks using weak-references. These callbacks
+ * are called when the [owner] is destroyed, but can also be cleared when the caller is GCed
+ */
+class WeakCleanupSet(owner: LifecycleOwner) {
+
+    private val callbacks = Collections.newSetFromMap<OnOwnerDestroyedCallback>(WeakHashMap())
+    private var destroyed = false
+
+    init {
+        MAIN_EXECUTOR.execute {
+            owner.lifecycle.addObserver(
+                object : DefaultLifecycleObserver {
+
+                    override fun onDestroy(owner: LifecycleOwner) {
+                        destroyed = true
+                        callbacks.forEach { it.onOwnerDestroyed() }
+                    }
+                }
+            )
+        }
+    }
+
+    fun addOnOwnerDestroyedCallback(callback: OnOwnerDestroyedCallback) {
+        if (destroyed) callback.onOwnerDestroyed() else callbacks.add(callback)
+    }
+
+    /** Callback when the owner is destroyed */
+    interface OnOwnerDestroyedCallback {
+        fun onOwnerDestroyed()
+    }
+}
diff --git a/src/com/android/launcher3/views/ActivityContext.java b/src/com/android/launcher3/views/ActivityContext.java
index bcb9295..cbf5341 100644
--- a/src/com/android/launcher3/views/ActivityContext.java
+++ b/src/com/android/launcher3/views/ActivityContext.java
@@ -87,6 +87,7 @@
 import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.launcher3.util.SystemUiController;
 import com.android.launcher3.util.ViewCache;
+import com.android.launcher3.util.WeakCleanupSet;
 import com.android.launcher3.widget.picker.model.WidgetPickerDataProvider;
 
 import java.util.List;
@@ -520,6 +521,9 @@
         return new CellPosMapper(dp.isVerticalBarLayout(), dp.numShownHotseatIcons);
     }
 
+    /** Set to manage objects that can be cleaned up along with the context */
+    WeakCleanupSet getOwnerCleanupSet();
+
     /** Whether bubbles are enabled. */
     default boolean isBubbleBarEnabled() {
         return false;