From 459d5bd503c11478637a0f61f7c952eda1127e74 Mon Sep 17 00:00:00 2001
From: Bryce Lee <brycelee@google.com>
Date: Thu, 3 Mar 2022 14:00:43 -0800
Subject: [PATCH] Make OnBackInvoked callback a WeakReference.

OnBackInvokedCallbackWrapper instances are sent to the WindowState over
AIDL. This leads to Activity leak issues even if the reference is
cleared in the remote side. This change makes the held reference weak
So it is not retained.

Test: verified leak no longer present through hprof
Bug: 221285764
Change-Id: I1c529d27769e2426898776e47f37bdb7abe1f681
---
 .../window/WindowOnBackInvokedDispatcher.java | 46 ++++++++++++++-----
 1 file changed, 35 insertions(+), 11 deletions(-)

diff --git a/core/java/android/window/WindowOnBackInvokedDispatcher.java b/core/java/android/window/WindowOnBackInvokedDispatcher.java
index 5f82fb00ed752..5dbb551da681c 100644
--- a/core/java/android/window/WindowOnBackInvokedDispatcher.java
+++ b/core/java/android/window/WindowOnBackInvokedDispatcher.java
@@ -28,6 +28,7 @@ import android.util.Log;
 import android.view.IWindow;
 import android.view.IWindowSession;
 
+import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.TreeMap;
@@ -185,35 +186,58 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
     }
 
     private static class OnBackInvokedCallbackWrapper extends IOnBackInvokedCallback.Stub {
-        private final OnBackInvokedCallback mCallback;
+        private final WeakReference<OnBackInvokedCallback> mCallback;
 
         OnBackInvokedCallbackWrapper(@NonNull OnBackInvokedCallback callback) {
-            mCallback = callback;
-        }
-
-        @NonNull
-        public OnBackInvokedCallback getCallback() {
-            return mCallback;
+            mCallback = new WeakReference<>(callback);
         }
 
         @Override
         public void onBackStarted() {
-            Handler.getMain().post(() -> mCallback.onBackStarted());
+            Handler.getMain().post(() -> {
+                final OnBackInvokedCallback callback = mCallback.get();
+                if (callback == null) {
+                    return;
+                }
+
+                callback.onBackStarted();
+            });
         }
 
         @Override
         public void onBackProgressed(BackEvent backEvent) {
-            Handler.getMain().post(() -> mCallback.onBackProgressed(backEvent));
+            Handler.getMain().post(() -> {
+                final OnBackInvokedCallback callback = mCallback.get();
+                if (callback == null) {
+                    return;
+                }
+
+                callback.onBackProgressed(backEvent);
+            });
         }
 
         @Override
         public void onBackCancelled() {
-            Handler.getMain().post(() -> mCallback.onBackCancelled());
+            Handler.getMain().post(() -> {
+                final OnBackInvokedCallback callback = mCallback.get();
+                if (callback == null) {
+                    return;
+                }
+
+                callback.onBackCancelled();
+            });
         }
 
         @Override
         public void onBackInvoked() throws RemoteException {
-            Handler.getMain().post(() -> mCallback.onBackInvoked());
+            Handler.getMain().post(() -> {
+                final OnBackInvokedCallback callback = mCallback.get();
+                if (callback == null) {
+                    return;
+                }
+
+                callback.onBackInvoked();
+            });
         }
     }