Fix Notification animation controller memory leak

This fixes leaking of NotificationLaunchAnimatorController instances (together with their referenced Notification entries) by breaking a dependency cycle through IPC.
Also removes completion callback - it's a very easy avenue to leak even more objects accidentally and is currently unused in SysUI code.

Bug: 299069405
Bug: 298028699
Bug: 293929945
Test: unit tests for attached class. Reproduced issue on cheetah,
confirmed fix of issue with a profiler.

Change-Id: Ib17577cba17ed35b24c4ab32ce0ed55fc7a24605

3 files changed, 151 insertions, 5 deletions

diff --git a/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityLaunchAnimator.kt b/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityLaunchAnimator.kt
index 0f2e4bace46d..4aac27932924 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityLaunchAnimator.kt
+++ b/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityLaunchAnimator.kt
@@ -39,6 +39,7 @@ import android.view.ViewGroup
 import android.view.WindowManager
 import android.view.animation.Interpolator
 import android.view.animation.PathInterpolator
+import androidx.annotation.AnyThread
 import androidx.annotation.BinderThread
 import androidx.annotation.UiThread
 import com.android.app.animation.Interpolators
@@ -149,6 +150,10 @@ class ActivityLaunchAnimator(
             override fun onLaunchAnimationProgress(linearProgress: Float) {
                 listeners.forEach { it.onLaunchAnimationProgress(linearProgress) }
             }
+
+            override fun onLaunchAnimationCancelled() {
+                listeners.forEach { it.onLaunchAnimationCancelled() }
+            }
         }
 
     /**
@@ -191,6 +196,7 @@ class ActivityLaunchAnimator(
                     "ActivityLaunchAnimator.callback must be set before using this animator"
                 )
         val runner = createRunner(controller)
+        val runnerDelegate = runner.delegate!!
         val hideKeyguardWithAnimation = callback.isOnKeyguard() && !showOverLockscreen
 
         // Pass the RemoteAnimationAdapter to the intent starter only if we are not hiding the
@@ -241,12 +247,15 @@ class ActivityLaunchAnimator(
         // If we expect an animation, post a timeout to cancel it in case the remote animation is
         // never started.
         if (willAnimate) {
-            runner.delegate.postTimeout()
+            runnerDelegate.postTimeout()
 
             // Hide the keyguard using the launch animation instead of the default unlock animation.
             if (hideKeyguardWithAnimation) {
                 callback.hideKeyguardWithAnimation(runner)
             }
+        } else {
+            // We need to make sure delegate references are dropped to avoid memory leaks.
+            runner.dispose()
         }
     }
 
@@ -344,6 +353,13 @@ class ActivityLaunchAnimator(
          */
         fun onLaunchAnimationEnd() {}
 
+        /**
+         * The animation was cancelled. Note that [onLaunchAnimationEnd] will still be called after
+         * this if the animation was already started, i.e. if [onLaunchAnimationStart] was called
+         * before the cancellation.
+         */
+        fun onLaunchAnimationCancelled() {}
+
         /** Called when an activity launch animation made progress. */
         fun onLaunchAnimationProgress(linearProgress: Float) {}
     }
@@ -426,6 +442,39 @@ class ActivityLaunchAnimator(
         fun onLaunchAnimationCancelled(newKeyguardOccludedState: Boolean? = null) {}
     }
 
+    /**
+     * Invokes [onAnimationComplete] when animation is either cancelled or completed. Delegates all
+     * events to the passed [delegate].
+     */
+    @VisibleForTesting
+    inner class DelegatingAnimationCompletionListener(
+        private val delegate: Listener?,
+        private val onAnimationComplete: () -> Unit
+    ) : Listener {
+        var cancelled = false
+
+        override fun onLaunchAnimationStart() {
+            delegate?.onLaunchAnimationStart()
+        }
+
+        override fun onLaunchAnimationProgress(linearProgress: Float) {
+            delegate?.onLaunchAnimationProgress(linearProgress)
+        }
+
+        override fun onLaunchAnimationEnd() {
+            delegate?.onLaunchAnimationEnd()
+            if (!cancelled) {
+                onAnimationComplete.invoke()
+            }
+        }
+
+        override fun onLaunchAnimationCancelled() {
+            cancelled = true
+            delegate?.onLaunchAnimationCancelled()
+            onAnimationComplete.invoke()
+        }
+    }
+
     @VisibleForTesting
     inner class Runner(
         controller: Controller,
@@ -436,11 +485,21 @@ class ActivityLaunchAnimator(
         listener: Listener? = null
     ) : IRemoteAnimationRunner.Stub() {
         private val context = controller.launchContainer.context
-        internal val delegate: AnimationDelegate
+
+        // This is being passed across IPC boundaries and cycles (through PendingIntentRecords,
+        // etc.) are possible. So we need to make sure we drop any references that might
+        // transitively cause leaks when we're done with animation.
+        @VisibleForTesting var delegate: AnimationDelegate?
 
         init {
             delegate =
-                AnimationDelegate(controller, callback, listener, launchAnimator, disableWmTimeout)
+                AnimationDelegate(
+                    controller,
+                    callback,
+                    DelegatingAnimationCompletionListener(listener, this::dispose),
+                    launchAnimator,
+                    disableWmTimeout
+                )
         }
 
         @BinderThread
@@ -451,14 +510,33 @@ class ActivityLaunchAnimator(
             nonApps: Array<out RemoteAnimationTarget>?,
             finishedCallback: IRemoteAnimationFinishedCallback?
         ) {
+            val delegate = delegate
             context.mainExecutor.execute {
-                delegate.onAnimationStart(transit, apps, wallpapers, nonApps, finishedCallback)
+                if (delegate == null) {
+                    Log.i(TAG, "onAnimationStart called after completion")
+                    // Animation started too late and timed out already. We need to still
+                    // signal back that we're done with it.
+                    finishedCallback?.onAnimationFinished()
+                } else {
+                    delegate.onAnimationStart(transit, apps, wallpapers, nonApps, finishedCallback)
+                }
             }
         }
 
         @BinderThread
         override fun onAnimationCancelled() {
-            context.mainExecutor.execute { delegate.onAnimationCancelled() }
+            val delegate = delegate
+            context.mainExecutor.execute {
+                delegate ?: Log.wtf(TAG, "onAnimationCancelled called after completion")
+                delegate?.onAnimationCancelled()
+            }
+        }
+
+        @AnyThread
+        fun dispose() {
+            // Drop references to animation controller once we're done with the animation
+            // to avoid leaking.
+            context.mainExecutor.execute { delegate = null }
         }
     }
 
@@ -584,6 +662,7 @@ class ActivityLaunchAnimator(
                     )
                 }
                 controller.onLaunchAnimationCancelled()
+                listener?.onLaunchAnimationCancelled()
                 return
             }
 
@@ -821,6 +900,7 @@ class ActivityLaunchAnimator(
                 Log.d(TAG, "Calling controller.onLaunchAnimationCancelled() [animation timed out]")
             }
             controller.onLaunchAnimationCancelled()
+            listener?.onLaunchAnimationCancelled()
         }
 
         @UiThread
@@ -842,6 +922,7 @@ class ActivityLaunchAnimator(
                 )
             }
             controller.onLaunchAnimationCancelled()
+            listener?.onLaunchAnimationCancelled()
         }
 
         private fun IRemoteAnimationFinishedCallback.invoke() {
diff --git a/packages/SystemUI/src/com/android/systemui/util/ReferenceExt.kt b/packages/SystemUI/src/com/android/systemui/util/ReferenceExt.kt
new file mode 100644
index 000000000000..ac04d31041b6
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/util/ReferenceExt.kt
@@ -0,0 +1,50 @@
+package com.android.systemui.util
+
+import java.lang.ref.SoftReference
+import java.lang.ref.WeakReference
+import kotlin.properties.ReadWriteProperty
+import kotlin.reflect.KProperty
+
+/**
+ * Creates a Kotlin idiomatic weak reference.
+ *
+ * Usage:
+ * ```
+ * var weakReferenceObj: Object? by weakReference(null)
+ * weakReferenceObj = Object()
+ * ```
+ */
+fun <T> weakReference(obj: T? = null): ReadWriteProperty<Any?, T?> {
+    return object : ReadWriteProperty<Any?, T?> {
+        var weakRef = WeakReference<T?>(obj)
+        override fun getValue(thisRef: Any?, property: KProperty<*>): T? {
+            return weakRef.get()
+        }
+
+        override fun setValue(thisRef: Any?, property: KProperty<*>, value: T?) {
+            weakRef = WeakReference(value)
+        }
+    }
+}
+
+/**
+ * Creates a Kotlin idiomatic soft reference.
+ *
+ * Usage:
+ * ```
+ * var softReferenceObj: Object? by softReference(null)
+ * softReferenceObj = Object()
+ * ```
+ */
+fun <T> softReference(obj: T? = null): ReadWriteProperty<Any?, T?> {
+    return object : ReadWriteProperty<Any?, T?> {
+        var softRef = SoftReference<T?>(obj)
+        override fun getValue(thisRef: Any?, property: KProperty<*>): T? {
+            return softRef.get()
+        }
+
+        override fun setValue(thisRef: Any?, property: KProperty<*>, value: T?) {
+            softRef = SoftReference(value)
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/animation/ActivityLaunchAnimatorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/animation/ActivityLaunchAnimatorTest.kt
index 59c7e7669b63..8faf715521f8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/animation/ActivityLaunchAnimatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/animation/ActivityLaunchAnimatorTest.kt
@@ -166,6 +166,9 @@ class ActivityLaunchAnimatorTest : SysuiTestCase() {
         waitForIdleSync()
         verify(controller).onLaunchAnimationCancelled()
         verify(controller, never()).onLaunchAnimationStart(anyBoolean())
+        verify(listener).onLaunchAnimationCancelled()
+        verify(listener, never()).onLaunchAnimationStart()
+        assertNull(runner.delegate)
     }
 
     @Test
@@ -176,6 +179,9 @@ class ActivityLaunchAnimatorTest : SysuiTestCase() {
         waitForIdleSync()
         verify(controller).onLaunchAnimationCancelled()
         verify(controller, never()).onLaunchAnimationStart(anyBoolean())
+        verify(listener).onLaunchAnimationCancelled()
+        verify(listener, never()).onLaunchAnimationStart()
+        assertNull(runner.delegate)
     }
 
     @Test
@@ -194,6 +200,15 @@ class ActivityLaunchAnimatorTest : SysuiTestCase() {
         }
     }
 
+    @Test
+    fun disposeRunner_delegateDereferenced() {
+        val runner = activityLaunchAnimator.createRunner(controller)
+        assertNotNull(runner.delegate)
+        runner.dispose()
+        waitForIdleSync()
+        assertNull(runner.delegate)
+    }
+
     private fun fakeWindow(): RemoteAnimationTarget {
         val bounds = Rect(10 /* left */, 20 /* top */, 30 /* right */, 40 /* bottom */)
         val taskInfo = ActivityManager.RunningTaskInfo()