Remove callback before WindowState removal

ViewRootImpl was removing the callback from the server after the
windowState was removed, which was consistently failing resulting in an
exception in the log.

Now the removal is done before the ViewRootImpl's window removal

Test: Manually check that the error is not appearing in the log
Bug: 226089354
Change-Id: Ie9cb9940791b16d6ae1b243079794583c87b1a5c

3 files changed, 12 insertions, 8 deletions

diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index 17e3914ab24b..5c85eb5e5551 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -8502,6 +8502,7 @@ public final class ViewRootImpl implements ViewParent,
                 return;
             }
             mRemoved = true;
+            mOnBackInvokedDispatcher.detachFromWindow();
             if (mAdded) {
                 dispatchDetachedFromWindow();
             }
@@ -8536,8 +8537,6 @@ public final class ViewRootImpl implements ViewParent,
 
             mAdded = false;
         }
-        unregisterCompatOnBackInvokedCallback();
-        mOnBackInvokedDispatcher.detachFromWindow();
         WindowManagerGlobal.getInstance().doRemoveView(this);
     }
 
diff --git a/core/java/android/window/WindowOnBackInvokedDispatcher.java b/core/java/android/window/WindowOnBackInvokedDispatcher.java
index 97573c291340..e0bee96002b8 100644
--- a/core/java/android/window/WindowOnBackInvokedDispatcher.java
+++ b/core/java/android/window/WindowOnBackInvokedDispatcher.java
@@ -20,6 +20,7 @@ import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.compat.CompatChanges;
 import android.content.Context;
+import android.os.Debug;
 import android.os.Handler;
 import android.os.RemoteException;
 import android.os.SystemProperties;
@@ -35,11 +36,11 @@ import java.util.TreeMap;
 
 /**
  * Provides window based implementation of {@link OnBackInvokedDispatcher}.
- *
+ * <p>
  * Callbacks with higher priorities receive back dispatching first.
  * Within the same priority, callbacks receive back dispatching in the reverse order
  * in which they are added.
- *
+ * <p>
  * When the top priority callback is updated, the new callback is propagated to the Window Manager
  * if the window the instance is associated with has been attached. It is allowed to register /
  * unregister {@link OnBackInvokedCallback}s before the window is attached, although
@@ -166,6 +167,10 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
                 mWindowSession.setOnBackInvokedCallback(
                         mWindow, new OnBackInvokedCallbackWrapper(callback), priority);
             }
+            if (DEBUG && callback == null) {
+                Log.d(TAG, TextUtils.formatSimple("setTopOnBackInvokedCallback(null) Callers:%s",
+                        Debug.getCallers(5, "  ")));
+            }
         } catch (RemoteException e) {
             Log.e(TAG, "Failed to set OnBackInvokedCallback to WM. Error: " + e);
         }
@@ -243,7 +248,7 @@ public class WindowOnBackInvokedDispatcher implements OnBackInvokedDispatcher {
 
     /**
      * Returns if the legacy back behavior should be used.
-     *
+     * <p>
      * Legacy back behavior dispatches KEYCODE_BACK instead of invoking the application registered
      * {@link OnBackInvokedCallback}.
      */
diff --git a/services/core/java/com/android/server/wm/Session.java b/services/core/java/com/android/server/wm/Session.java
index c4b7c56358b4..68ae0b44903e 100644
--- a/services/core/java/com/android/server/wm/Session.java
+++ b/services/core/java/com/android/server/wm/Session.java
@@ -915,12 +915,12 @@ class Session extends IWindowSession.Stub implements IBinder.DeathRecipient {
     public void setOnBackInvokedCallback(
             IWindow window,
             IOnBackInvokedCallback onBackInvokedCallback,
-            @OnBackInvokedDispatcher.Priority int priority) throws RemoteException {
+            @OnBackInvokedDispatcher.Priority int priority) {
         synchronized (mService.mGlobalLock) {
-            WindowState windowState = mService.windowForClientLocked(this, window, true);
+            WindowState windowState = mService.windowForClientLocked(this, window, false);
             if (windowState == null) {
                 Slog.e(TAG_WM,
-                        "setOnBackInvokedCallback(): Can't find window state for package:"
+                        "setOnBackInvokedCallback(): No window state for package:"
                                 + mPackageName);
             } else {
                 windowState.setOnBackInvokedCallback(onBackInvokedCallback, priority);