VirtualDisplayAdapter remove unnecessary owner uids from map and null check

Bug: 261791612
Bug: 383400124
Flag: com.android.server.display.feature.flags.virtual_display_limit
Test: atest VirtualDisplayAdapterTest
Test: atest CtsInputTestCases:android.input.cts.PointerIconTest
Change-Id: I416141a700aed2aa06df003b15356172566d1152

1 files changed, 8 insertions, 7 deletions

diff --git a/services/core/java/com/android/server/display/VirtualDisplayAdapter.java b/services/core/java/com/android/server/display/VirtualDisplayAdapter.java
index f14e452ab8d3..0bd67a826460 100644
--- a/services/core/java/com/android/server/display/VirtualDisplayAdapter.java
+++ b/services/core/java/com/android/server/display/VirtualDisplayAdapter.java
@@ -307,13 +307,14 @@ public class VirtualDisplayAdapter extends DisplayAdapter {
 
     private VirtualDisplayDevice removeVirtualDisplayDeviceLocked(IBinder appToken) {
         if (getFeatureFlags().isVirtualDisplayLimitEnabled()) {
-            int ownerUid = mOwnerUids.get(appToken);
-            int noOfDevices = mNoOfDevicesPerPackage.get(ownerUid, /* valueIfKeyNotFound= */ 0);
-            if (noOfDevices <= 1) {
-                mNoOfDevicesPerPackage.delete(ownerUid);
-                mOwnerUids.remove(appToken);
-            } else {
-                mNoOfDevicesPerPackage.put(ownerUid, noOfDevices - 1);
+            Integer ownerUid = mOwnerUids.remove(appToken);
+            if (ownerUid != null) {
+                int noOfDevices = mNoOfDevicesPerPackage.get(ownerUid, /* valueIfKeyNotFound= */ 0);
+                if (noOfDevices <= 1) {
+                    mNoOfDevicesPerPackage.delete(ownerUid);
+                } else {
+                    mNoOfDevicesPerPackage.put(ownerUid, noOfDevices - 1);
+                }
             }
         }
         return mVirtualDisplayDevices.remove(appToken);