Compare binder instead of the proxy for DisplayAreaOrganizer

Before, we directly compared the da.mOrganizer when unregister, which
failed to remove it. Now, we compare the binder instead.

Bug: 169224778
Test: manual: test with OneHanded mode
Test: atest WmTests:DisplayAreaOrganizerTest
Change-Id: I907615c111b10787dfb53b6fe2f683d651c21b23

2 files changed, 39 insertions, 11 deletions

diff --git a/services/core/java/com/android/server/wm/DisplayAreaOrganizerController.java b/services/core/java/com/android/server/wm/DisplayAreaOrganizerController.java
index 01c007e381b1..dd92f507a33d 100644
--- a/services/core/java/com/android/server/wm/DisplayAreaOrganizerController.java
+++ b/services/core/java/com/android/server/wm/DisplayAreaOrganizerController.java
@@ -52,10 +52,7 @@ public class DisplayAreaOrganizerController extends IDisplayAreaOrganizerControl
         public void binderDied() {
             synchronized (mGlobalLock) {
                 mOrganizersByFeatureIds.remove(mFeature);
-                mService.mRootWindowContainer.forAllDisplayAreas((da) -> {
-                    if (da.mOrganizer != mOrganizer) return;
-                    da.setOrganizer(null);
-                });
+                removeOrganizer(mOrganizer);
             }
         }
     }
@@ -112,11 +109,7 @@ public class DisplayAreaOrganizerController extends IDisplayAreaOrganizerControl
                         organizer.asBinder(), uid);
                 mOrganizersByFeatureIds.entrySet().removeIf(
                         entry -> entry.getValue().asBinder() == organizer.asBinder());
-
-                mService.mRootWindowContainer.forAllDisplayAreas((da) -> {
-                    if (da.mOrganizer != organizer) return;
-                    da.setOrganizer(null);
-                });
+                removeOrganizer(organizer);
             }
         } finally {
             Binder.restoreCallingIdentity(origId);
@@ -151,4 +144,13 @@ public class DisplayAreaOrganizerController extends IDisplayAreaOrganizerControl
             // Oh well...
         }
     }
+
+    private void removeOrganizer(IDisplayAreaOrganizer organizer) {
+        IBinder organizerBinder = organizer.asBinder();
+        mService.mRootWindowContainer.forAllDisplayAreas((da) -> {
+            if (da.mOrganizer != null && da.mOrganizer.asBinder().equals(organizerBinder)) {
+                da.setOrganizer(null);
+            }
+        });
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayAreaOrganizerTest.java b/services/tests/wmtests/src/com/android/server/wm/DisplayAreaOrganizerTest.java
index f4b50dc6b553..54b2b3b4a009 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayAreaOrganizerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayAreaOrganizerTest.java
@@ -18,6 +18,8 @@ package com.android.server.wm;
 
 import static android.window.DisplayAreaOrganizer.FEATURE_VENDOR_FIRST;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -39,6 +41,10 @@ import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+/**
+ * Build/Install/Run:
+ *  atest WmTests:DisplayAreaOrganizerTest
+ */
 @SmallTest
 @Presubmit
 @RunWith(WindowTestRunner.class)
@@ -61,14 +67,22 @@ public class DisplayAreaOrganizerTest extends WindowTestsBase {
     }
 
     private IDisplayAreaOrganizer registerMockOrganizer(int feature) {
-        final IDisplayAreaOrganizer organizer = mock(IDisplayAreaOrganizer.class);
-        when(organizer.asBinder()).thenReturn(new Binder());
+        return registerMockOrganizer(feature, new Binder());
+    }
 
+    private IDisplayAreaOrganizer registerMockOrganizer(int feature, Binder binder) {
+        final IDisplayAreaOrganizer organizer = createMockOrganizer(binder);
         mWm.mAtmService.mWindowOrganizerController.mDisplayAreaOrganizerController
                 .registerOrganizer(organizer, feature);
         return organizer;
     }
 
+    private IDisplayAreaOrganizer createMockOrganizer(Binder binder) {
+        final IDisplayAreaOrganizer organizer = mock(IDisplayAreaOrganizer.class);
+        when(organizer.asBinder()).thenReturn(binder);
+        return organizer;
+    }
+
     private void unregisterMockOrganizer(IDisplayAreaOrganizer organizer) {
         mWm.mAtmService.mWindowOrganizerController.mDisplayAreaOrganizerController
                 .unregisterOrganizer(organizer);
@@ -99,4 +113,16 @@ public class DisplayAreaOrganizerTest extends WindowTestsBase {
         // Ensure it was still only called once if the bounds didn't change
         verify(organizer).onDisplayAreaInfoChanged(any());
     }
+
+    @Test
+    public void testUnregisterOrganizer() {
+        final Binder binder = new Binder();
+        registerMockOrganizer(FEATURE_VENDOR_FIRST, binder);
+
+        assertThat(mTestDisplayArea.mOrganizer).isNotNull();
+
+        unregisterMockOrganizer(createMockOrganizer(binder));
+
+        assertThat(mTestDisplayArea.mOrganizer).isNull();
+    }
 }