From b39eb56d5293e1c758e3e26d8bf5189650ffe499 Mon Sep 17 00:00:00 2001
From: Shan Huang <shanh@google.com>
Date: Mon, 6 Feb 2023 12:14:54 -0800
Subject: [PATCH] Fix IME callback being unexpectedly garbage-collected and
 breaking back gesture.

Bug: 266528689
Test: Open GMail, Calendar, Launcher, find UI surfaces to bring up IME,
then swipe back. Repeat in different combinations of the steps and make
sure back always works.
Test: atest WindowOnBackInvokedDispatcherTest

Change-Id: I0cba6d6963c6b988f02a7baf56f2901d17b5c9b3
---
 .../window/ImeOnBackInvokedDispatcher.java    |  7 +++-
 .../window/WindowOnBackInvokedDispatcher.java | 33 ++++++++++++++++---
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/core/java/android/window/ImeOnBackInvokedDispatcher.java b/core/java/android/window/ImeOnBackInvokedDispatcher.java
index a0bd7f70ca585..a8c8760a0842d 100644
--- a/core/java/android/window/ImeOnBackInvokedDispatcher.java
+++ b/core/java/android/window/ImeOnBackInvokedDispatcher.java
@@ -82,8 +82,13 @@ public class ImeOnBackInvokedDispatcher implements OnBackInvokedDispatcher, Parc
             @NonNull OnBackInvokedCallback callback) {
         final Bundle bundle = new Bundle();
         // Always invoke back for ime without checking the window focus.
+        // We use strong reference in the binder wrapper to avoid accidentally GC the callback.
+        // This is necessary because the callback is sent to and registered from
+        // the app process, which may treat the IME callback as weakly referenced. This will not
+        // cause a memory leak because the app side already clears the reference correctly.
         final IOnBackInvokedCallback iCallback =
-                new WindowOnBackInvokedDispatcher.OnBackInvokedCallbackWrapper(callback);
+                new WindowOnBackInvokedDispatcher.OnBackInvokedCallbackWrapper(
+                        callback, false /* useWeakRef */);
         bundle.putBinder(RESULT_KEY_CALLBACK, iCallback.asBinder());
         bundle.putInt(RESULT_KEY_PRIORITY, priority);
         bundle.putInt(RESULT_KEY_ID, callback.hashCode());
diff --git a/core/java/android/window/WindowOnBackInvokedDispatcher.java b/core/java/android/window/WindowOnBackInvokedDispatcher.java
index 7a5510cc252c5..ed93ac184c666 100644
--- a/core/java/android/window/WindowOnBackInvokedDispatcher.java
+++ b/core/java/android/window/WindowOnBackInvokedDispatcher.java
@@ -233,10 +233,34 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
     }
 
     static class OnBackInvokedCallbackWrapper extends IOnBackInvokedCallback.Stub {
-        private final WeakReference<OnBackInvokedCallback> mCallback;
+        static class CallbackRef {
+            final WeakReference<OnBackInvokedCallback> mWeakRef;
+            final OnBackInvokedCallback mStrongRef;
+            CallbackRef(@NonNull OnBackInvokedCallback callback, boolean useWeakRef) {
+                if (useWeakRef) {
+                    mWeakRef = new WeakReference<>(callback);
+                    mStrongRef = null;
+                } else {
+                    mStrongRef = callback;
+                    mWeakRef = null;
+                }
+            }
+
+            OnBackInvokedCallback get() {
+                if (mStrongRef != null) {
+                    return mStrongRef;
+                }
+                return mWeakRef.get();
+            }
+        }
+        final CallbackRef mCallbackRef;
 
         OnBackInvokedCallbackWrapper(@NonNull OnBackInvokedCallback callback) {
-            mCallback = new WeakReference<>(callback);
+            mCallbackRef = new CallbackRef(callback, true /* useWeakRef */);
+        }
+
+        OnBackInvokedCallbackWrapper(@NonNull OnBackInvokedCallback callback, boolean useWeakRef) {
+            mCallbackRef = new CallbackRef(callback, useWeakRef);
         }
 
         @Override
@@ -279,8 +303,9 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
         public void onBackInvoked() throws RemoteException {
             Handler.getMain().post(() -> {
                 mProgressAnimator.reset();
-                final OnBackInvokedCallback callback = mCallback.get();
+                final OnBackInvokedCallback callback = mCallbackRef.get();
                 if (callback == null) {
+                    Log.d(TAG, "Trying to call onBackInvoked() on a null callback reference.");
                     return;
                 }
                 callback.onBackInvoked();
@@ -289,7 +314,7 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
 
         @Nullable
         private OnBackAnimationCallback getBackAnimationCallback() {
-            OnBackInvokedCallback callback = mCallback.get();
+            OnBackInvokedCallback callback = mCallbackRef.get();
             return callback instanceof OnBackAnimationCallback ? (OnBackAnimationCallback) callback
                     : null;
         }