Allow unregistring callbacks in DeviceStateController

Add unregisterDeviceStateCallback() to DeviceStateController to avoid
memory leak due to unnecessary references.

Bug: 271071505
Test: atest com.android.server.wm.DeviceStateControllerTests
Change-Id: Id1e13f9e5c163cb4bfcf3467ba510f1227016942

services/core/java/com/android/server/wm/DeviceStateController.java

@@ -24,6 +24,7 @@ import android.os.HandlerExecutor;
 
 import com.android.internal.R;
 import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.ArrayUtils;
 
 import java.util.ArrayList;
@@ -51,7 +52,8 @@ final class DeviceStateController implements DeviceStateManager.DeviceStateCallb
     private final int[] mReverseRotationAroundZAxisStates;
     @GuardedBy("this")
     @NonNull
-    private final List<Consumer<DeviceState>> mDeviceStateCallbacks = new ArrayList<>();
+    @VisibleForTesting
+    final List<Consumer<DeviceState>> mDeviceStateCallbacks = new ArrayList<>();
 
     private final boolean mMatchBuiltInDisplayOrientationToDefaultDisplay;
 
@@ -98,6 +100,12 @@ final class DeviceStateController implements DeviceStateManager.DeviceStateCallb
         }
     }
 
+    void unregisterDeviceStateCallback(@NonNull Consumer<DeviceState> callback) {
+        synchronized (this) {
+            mDeviceStateCallbacks.remove(callback);
+        }
+    }
+
     /**
      * @return true if the rotation direction on the Z axis should be reversed.
      */

services/core/java/com/android/server/wm/DisplayContent.java

@@ -601,6 +601,7 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
 
     @VisibleForTesting
     final DeviceStateController mDeviceStateController;
+    final Consumer<DeviceStateController.DeviceState> mDeviceStateConsumer;
     private final PhysicalDisplaySwitchTransitionLauncher mDisplaySwitchTransitionLauncher;
     final RemoteDisplayChangeController mRemoteDisplayChangeController;
 
@@ -1166,12 +1167,12 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         mDisplayRotation = new DisplayRotation(mWmService, this, mDisplayInfo.address,
                 mDeviceStateController, root.getDisplayRotationCoordinator());
 
-        final Consumer<DeviceStateController.DeviceState> deviceStateConsumer =
+        mDeviceStateConsumer =
                 (@NonNull DeviceStateController.DeviceState newFoldState) -> {
                     mDisplaySwitchTransitionLauncher.foldStateChanged(newFoldState);
                     mDisplayRotation.foldStateChanged(newFoldState);
                 };
-        mDeviceStateController.registerDeviceStateCallback(deviceStateConsumer);
+        mDeviceStateController.registerDeviceStateCallback(mDeviceStateConsumer);
 
         mCloseToSquareMaxAspectRatio = mWmService.mContext.getResources().getFloat(
                 R.dimen.config_closeToSquareDisplayMaxAspectRatio);
@@ -3283,6 +3284,7 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
             handleAnimatingStoppedAndTransition();
             mWmService.stopFreezingDisplayLocked();
             mDisplayRotation.removeDefaultDisplayRotationChangedCallback();
+            mDeviceStateController.unregisterDeviceStateCallback(mDeviceStateConsumer);
             super.removeImmediately();
             if (DEBUG_DISPLAY) Slog.v(TAG_WM, "Removing display=" + this);
             mPointerEventDispatcher.dispose();

services/tests/wmtests/src/com/android/server/wm/DeviceStateControllerTests.java

@@ -21,6 +21,7 @@ import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.when;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 
 import android.content.Context;
@@ -55,6 +56,7 @@ public class DeviceStateControllerTests {
     private DeviceStateManager mMockDeviceStateManager;
     private DeviceStateController.DeviceState mCurrentState =
             DeviceStateController.DeviceState.UNKNOWN;
+    private Consumer<DeviceStateController.DeviceState> mDelegate;
 
     @Before
     public void setUp() {
@@ -64,10 +66,10 @@ public class DeviceStateControllerTests {
 
     private void initialize(boolean supportFold, boolean supportHalfFold) {
         mBuilder.setSupportFold(supportFold, supportHalfFold);
-        Consumer<DeviceStateController.DeviceState> delegate = (newFoldState) -> {
+        mDelegate = (newFoldState) -> {
             mCurrentState = newFoldState;
         };
-        mBuilder.setDelegate(delegate);
+        mBuilder.setDelegate(mDelegate);
         mBuilder.build();
         verify(mMockDeviceStateManager).registerCallback(any(), any());
     }
@@ -111,6 +113,24 @@ public class DeviceStateControllerTests {
         assertEquals(DeviceStateController.DeviceState.CONCURRENT, mCurrentState);
     }
 
+    @Test
+    public void testUnregisterDeviceStateCallback() {
+        initialize(true /* supportFold */, true /* supportHalfFolded */);
+        assertEquals(1, mTarget.mDeviceStateCallbacks.size());
+        assertEquals(mDelegate, mTarget.mDeviceStateCallbacks.get(0));
+
+        mTarget.onStateChanged(mOpenDeviceStates[0]);
+        assertEquals(DeviceStateController.DeviceState.OPEN, mCurrentState);
+        mTarget.onStateChanged(mFoldedStates[0]);
+        assertEquals(DeviceStateController.DeviceState.FOLDED, mCurrentState);
+
+        // The callback should not receive state change when the it is unregistered.
+        mTarget.unregisterDeviceStateCallback(mDelegate);
+        assertTrue(mTarget.mDeviceStateCallbacks.isEmpty());
+        mTarget.onStateChanged(mOpenDeviceStates[0]);
+        assertEquals(DeviceStateController.DeviceState.FOLDED /* unchanged */, mCurrentState);
+    }
+
     private final int[] mFoldedStates = {0};
     private final int[] mOpenDeviceStates = {1};
     private final int[] mHalfFoldedStates = {2};