Create ActivityDisplay if its registered in DisplayManager

Activity manager is notified asynchronously when new display is
created. Because of that someone may try to access or use it
before AM is aware of it, e.g. launch an activity to a just
create virtual display.
This CL guards the cases where we try to use a display by id in
AM and creates a new instance if needed.

Bug: 34262188
Test: android.server.cts.ActivityManagerDisplayTests
Test: #testImmediateLaunchOnNewDisplay
Test: go/wm-smoke
Change-Id: Iab0871708145aa80ff252a0c764beb6a3b5dc42d
diff --git a/services/core/java/com/android/server/am/ActivityStackSupervisor.java b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
index e180aef..440a67b 100644
--- a/services/core/java/com/android/server/am/ActivityStackSupervisor.java
+++ b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
@@ -470,7 +470,7 @@
     }
 
     Configuration getDisplayOverrideConfiguration(int displayId) {
-        final ActivityDisplay activityDisplay = mActivityDisplays.get(displayId);
+        final ActivityDisplay activityDisplay = getActivityDisplayOrCreateLocked(displayId);
         if (activityDisplay == null) {
             throw new IllegalArgumentException("No display found with id: " + displayId);
         }
@@ -479,7 +479,7 @@
     }
 
     void setDisplayOverrideConfiguration(Configuration overrideConfiguration, int displayId) {
-        final ActivityDisplay activityDisplay = mActivityDisplays.get(displayId);
+        final ActivityDisplay activityDisplay = getActivityDisplayOrCreateLocked(displayId);
         if (activityDisplay == null) {
             throw new IllegalArgumentException("No display found with id: " + displayId);
         }
@@ -507,7 +507,7 @@
         if (displayId == INVALID_DISPLAY) {
             return false;
         }
-        final ActivityDisplay targetDisplay = mActivityDisplays.get(displayId);
+        final ActivityDisplay targetDisplay = getActivityDisplayOrCreateLocked(displayId);
         if (targetDisplay == null) {
             throw new IllegalArgumentException("No display found with id: " + displayId);
         }
@@ -1672,7 +1672,7 @@
         if (DEBUG_TASKS) Slog.d(TAG, "Launch on display check: displayId=" + launchDisplayId
                 + " callingPid=" + callingPid + " callingUid=" + callingUid);
 
-        final ActivityDisplay activityDisplay = mActivityDisplays.get(launchDisplayId);
+        final ActivityDisplay activityDisplay = getActivityDisplayOrCreateLocked(launchDisplayId);
         if (activityDisplay == null) {
             Slog.w(TAG, "Launch on display check: display not found");
             return false;
@@ -2191,7 +2191,7 @@
      * @return Existing stack if there is a valid one, new dynamic stack if it is valid or null.
      */
     ActivityStack getValidLaunchStackOnDisplay(int displayId, @NonNull ActivityRecord r) {
-        final ActivityDisplay activityDisplay = mActivityDisplays.get(displayId);
+        final ActivityDisplay activityDisplay = getActivityDisplayOrCreateLocked(displayId);
         if (activityDisplay == null) {
             throw new IllegalArgumentException(
                     "Display with displayId=" + displayId + " not found.");
@@ -2242,10 +2242,9 @@
 
         for (int i = mTmpOrderedDisplayIds.size() - 1; i >= 0; --i) {
             final int displayId = mTmpOrderedDisplayIds.get(i);
-            final List<ActivityStack> stacks = mActivityDisplays.get(displayId).mStacks;
-            if (stacks == null) {
-                continue;
-            }
+            // If a display is registered in WM, it must also be available in AM.
+            @SuppressWarnings("ConstantConditions")
+            final List<ActivityStack> stacks = getActivityDisplayOrCreateLocked(displayId).mStacks;
             for (int j = stacks.size() - 1; j >= 0; --j) {
                 final ActivityStack stack = stacks.get(j);
                 if (stack != currentFocus && stack.isFocusable()
@@ -2576,7 +2575,7 @@
     }
 
     ActivityStack createStackOnDisplay(int stackId, int displayId, boolean onTop) {
-        final ActivityDisplay activityDisplay = mActivityDisplays.get(displayId);
+        final ActivityDisplay activityDisplay = getActivityDisplayOrCreateLocked(displayId);
         if (activityDisplay == null) {
             return null;
         }
@@ -2808,7 +2807,7 @@
      * @param onTop Indicates whether container should be place on top or on bottom.
      */
     void moveStackToDisplayLocked(int stackId, int displayId, boolean onTop) {
-        final ActivityDisplay activityDisplay = mActivityDisplays.get(displayId);
+        final ActivityDisplay activityDisplay = getActivityDisplayOrCreateLocked(displayId);
         if (activityDisplay == null) {
             throw new IllegalArgumentException("moveStackToDisplayLocked: Unknown displayId="
                     + displayId);
@@ -3915,25 +3914,44 @@
     }
 
     private void handleDisplayAdded(int displayId) {
-        boolean newDisplay;
         synchronized (mService) {
-            newDisplay = mActivityDisplays.get(displayId) == null;
-            if (newDisplay) {
-                ActivityDisplay activityDisplay = new ActivityDisplay(displayId);
-                if (activityDisplay.mDisplay == null) {
-                    Slog.w(TAG, "Display " + displayId + " gone before initialization complete");
-                    return;
-                }
-                mActivityDisplays.put(displayId, activityDisplay);
-                calculateDefaultMinimalSizeOfResizeableTasks(activityDisplay);
-                mWindowManager.onDisplayAdded(displayId);
-            }
+            getActivityDisplayOrCreateLocked(displayId);
         }
     }
 
     /** Check if display with specified id is added to the list. */
     boolean isDisplayAdded(int displayId) {
-        return mActivityDisplays.get(displayId) != null;
+        return getActivityDisplayOrCreateLocked(displayId) != null;
+    }
+
+    /**
+     * Get an existing instance of {@link ActivityDisplay} or create new if there is a
+     * corresponding record in display manager.
+     */
+    private ActivityDisplay getActivityDisplayOrCreateLocked(int displayId) {
+        ActivityDisplay activityDisplay = mActivityDisplays.get(displayId);
+        if (activityDisplay != null) {
+            return activityDisplay;
+        }
+        if (mDisplayManager == null) {
+            // The system isn't fully initialized yet.
+            return null;
+        }
+        final Display display = mDisplayManager.getDisplay(displayId);
+        if (display == null) {
+            // The display is not registered in DisplayManager.
+            return null;
+        }
+        // The display hasn't been added to ActivityManager yet, create a new record now.
+        activityDisplay = new ActivityDisplay(displayId);
+        if (activityDisplay.mDisplay == null) {
+            Slog.w(TAG, "Display " + displayId + " gone before initialization complete");
+            return null;
+        }
+        mActivityDisplays.put(displayId, activityDisplay);
+        calculateDefaultMinimalSizeOfResizeableTasks(activityDisplay);
+        mWindowManager.onDisplayAdded(displayId);
+        return activityDisplay;
     }
 
     private void calculateDefaultMinimalSizeOfResizeableTasks(ActivityDisplay display) {
@@ -3991,6 +4009,7 @@
         info.stackId = stack.mStackId;
         info.userId = stack.mCurrentUser;
         info.visible = stack.shouldBeVisible(null) == STACK_VISIBLE;
+        // A stack might be not attached to a display.
         info.position = display != null
                 ? display.mStacks.indexOf(stack)
                 : 0;
@@ -4618,7 +4637,7 @@
         @Override
         public void addToDisplay(int displayId) {
             synchronized (mService) {
-                ActivityDisplay activityDisplay = mActivityDisplays.get(displayId);
+                final ActivityDisplay activityDisplay = getActivityDisplayOrCreateLocked(displayId);
                 if (activityDisplay == null) {
                     return;
                 }