Associate SystemUiContext with DisplayContent

Before this CL, SystemUiContexts were fixed to the Display metrics
when the SystemUiContexts were created. It is the same mechanism
as DisplayContext but applies to SystemUiContexts unexpectedly.
This CL makes SystemUiContexts associate with DisplayContent with
the corresponding display ID. In this way, SystemUiContext would
receive updates when there is a Display property change.

Bug: 220022170
Bug: 194262507
Bug: 191064581
Bug: 205859784
Test: atest InputMethodMenuControllerTest WindowContextControllerTest
Test: atest NexusLauncherTests

Change-Id: I64a1614f32d097785915f6105b1813a929e0fe32
(cherry picked from commit 46ffd51d0723ad18b82d1093814ed7995996c8b9)

13 files changed, 249 insertions, 99 deletions

diff --git a/core/java/android/app/ActivityThread.java b/core/java/android/app/ActivityThread.java
index 8f6f552ad1ff..5c942d217fc4 100644
--- a/core/java/android/app/ActivityThread.java
+++ b/core/java/android/app/ActivityThread.java
@@ -30,6 +30,7 @@ import static android.app.servertransaction.ActivityLifecycleItem.ON_STOP;
 import static android.app.servertransaction.ActivityLifecycleItem.PRE_ON_CREATE;
 import static android.content.ContentResolver.DEPRECATE_DATA_COLUMNS;
 import static android.content.ContentResolver.DEPRECATE_DATA_PREFIX;
+import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.Display.INVALID_DISPLAY;
 
 import static com.android.internal.annotations.VisibleForTesting.Visibility.PACKAGE;
@@ -313,7 +314,7 @@ public final class ActivityThread extends ClientTransactionHandler
 
     @UnsupportedAppUsage
     private ContextImpl mSystemContext;
-    private ContextImpl mSystemUiContext;
+    private final SparseArray<ContextImpl> mDisplaySystemUiContexts = new SparseArray<>();
 
     @UnsupportedAppUsage
     static volatile IPackageManager sPackageManager;
@@ -2610,22 +2611,26 @@ public final class ActivityThread extends ClientTransactionHandler
     }
 
     @Override
+    @NonNull
     public ContextImpl getSystemUiContext() {
-        synchronized (this) {
-            if (mSystemUiContext == null) {
-                mSystemUiContext = ContextImpl.createSystemUiContext(getSystemContext());
-            }
-            return mSystemUiContext;
-        }
+        return getSystemUiContext(DEFAULT_DISPLAY);
     }
 
     /**
-     * Create the context instance base on system resources & display information which used for UI.
+     * Gets the context instance base on system resources & display information which used for UI.
      * @param displayId The ID of the display where the UI is shown.
      * @see ContextImpl#createSystemUiContext(ContextImpl, int)
      */
-    public ContextImpl createSystemUiContext(int displayId) {
-        return ContextImpl.createSystemUiContext(getSystemUiContext(), displayId);
+    @NonNull
+    public ContextImpl getSystemUiContext(int displayId) {
+        synchronized (this) {
+            ContextImpl systemUiContext = mDisplaySystemUiContexts.get(displayId);
+            if (systemUiContext == null) {
+                systemUiContext = ContextImpl.createSystemUiContext(getSystemContext(), displayId);
+                mDisplaySystemUiContexts.put(displayId, systemUiContext);
+            }
+            return systemUiContext;
+        }
     }
 
     public void installSystemApplicationInfo(ApplicationInfo info, ClassLoader classLoader) {
@@ -3740,7 +3745,7 @@ public final class ActivityThread extends ClientTransactionHandler
         if (pkgName != null && !pkgName.isEmpty()
                 && r.packageInfo.mPackageName.contains(pkgName)) {
             for (int id : dm.getDisplayIds()) {
-                if (id != Display.DEFAULT_DISPLAY) {
+                if (id != DEFAULT_DISPLAY) {
                     Display display =
                             dm.getCompatibleDisplay(id, appContext.getResources());
                     appContext = (ContextImpl) appContext.createDisplayContext(display);
diff --git a/core/java/android/app/ContextImpl.java b/core/java/android/app/ContextImpl.java
index 1397f5ea10dc..1093baade42c 100644
--- a/core/java/android/app/ContextImpl.java
+++ b/core/java/android/app/ContextImpl.java
@@ -2618,7 +2618,10 @@ class ContextImpl extends Context {
                 overrideConfig, display.getDisplayAdjustments().getCompatibilityInfo(),
                 mResources.getLoaders()));
         context.mDisplay = display;
-        context.mContextType = CONTEXT_TYPE_DISPLAY_CONTEXT;
+        // Inherit context type if the container is from System or System UI context to bypass
+        // UI context check.
+        context.mContextType = mContextType == CONTEXT_TYPE_SYSTEM_OR_SYSTEM_UI
+                ? CONTEXT_TYPE_SYSTEM_OR_SYSTEM_UI : CONTEXT_TYPE_DISPLAY_CONTEXT;
         // Display contexts and any context derived from a display context should always override
         // the display that would otherwise be inherited from mToken (or the global configuration if
         // mToken is null).
@@ -2671,7 +2674,8 @@ class ContextImpl extends Context {
 
         // Step 2. Create the base context of the window context, it will also create a Resources
         //         associated with the WindowTokenClient and set the token to the base context.
-        final ContextImpl windowContextBase = createWindowContextBase(windowTokenClient, display);
+        final ContextImpl windowContextBase = createWindowContextBase(windowTokenClient,
+                display.getDisplayId());
 
         // Step 3. Create a WindowContext instance and set it as the outer context of the base
         //         context to make the service obtained by #getSystemService(String) able to query
@@ -2696,9 +2700,7 @@ class ContextImpl extends Context {
         if (display == null) {
             throw new IllegalArgumentException("Display must not be null");
         }
-        final ContextImpl tokenContext = createWindowContextBase(token, display);
-        tokenContext.setResources(createWindowContextResources(tokenContext));
-        return tokenContext;
+        return createWindowContextBase(token, display.getDisplayId());
     }
 
     /**
@@ -2706,13 +2708,13 @@ class ContextImpl extends Context {
      * window.
      *
      * @param token The token to associate with {@link Resources}
-     * @param display The {@link Display} to associate with.
+     * @param displayId The ID of {@link Display} to associate with.
      *
      * @see #createWindowContext(Display, int, Bundle)
      * @see #createTokenContext(IBinder, Display)
      */
     @UiContext
-    ContextImpl createWindowContextBase(@NonNull IBinder token, @NonNull Display display) {
+    ContextImpl createWindowContextBase(@NonNull IBinder token, int displayId) {
         ContextImpl baseContext = new ContextImpl(this, mMainThread, mPackageInfo, mParams,
                 mAttributionSource.getAttributionTag(),
                 mAttributionSource.getNext(),
@@ -2726,8 +2728,8 @@ class ContextImpl extends Context {
         baseContext.setResources(windowContextResources);
         // Associate the display with window context resources so that configuration update from
         // the server side will also apply to the display's metrics.
-        baseContext.mDisplay = ResourcesManager.getInstance()
-                .getAdjustedDisplay(display.getDisplayId(), windowContextResources);
+        baseContext.mDisplay = ResourcesManager.getInstance().getAdjustedDisplay(displayId,
+                windowContextResources);
 
         return baseContext;
     }
@@ -2963,6 +2965,18 @@ class ContextImpl extends Context {
         mContentCaptureOptions = options;
     }
 
+    @Override
+    protected void finalize() throws Throwable {
+        // If mToken is a WindowTokenClient, the Context is usually associated with a
+        // WindowContainer. We should detach from WindowContainer when the Context is finalized
+        // if this Context is not a WindowContext. WindowContext finalization is handled in
+        // WindowContext class.
+        if (mToken instanceof WindowTokenClient && mContextType != CONTEXT_TYPE_WINDOW_CONTEXT) {
+            ((WindowTokenClient) mToken).detachFromWindowContainerIfNeeded();
+        }
+        super.finalize();
+    }
+
     @UnsupportedAppUsage
     static ContextImpl createSystemContext(ActivityThread mainThread) {
         LoadedApk packageInfo = new LoadedApk(mainThread);
@@ -2983,22 +2997,13 @@ class ContextImpl extends Context {
      * @param displayId The ID of the display where the UI is shown.
      */
     static ContextImpl createSystemUiContext(ContextImpl systemContext, int displayId) {
-        final LoadedApk packageInfo = systemContext.mPackageInfo;
-        ContextImpl context = new ContextImpl(null, systemContext.mMainThread, packageInfo,
-                ContextParams.EMPTY, null, null, null, null, null, 0, null, null);
-        context.setResources(createResources(null, packageInfo, null, displayId, null,
-                packageInfo.getCompatibilityInfo(), null));
-        context.updateDisplay(displayId);
+        final WindowTokenClient token = new WindowTokenClient();
+        final ContextImpl context = systemContext.createWindowContextBase(token, displayId);
+        token.attachContext(context);
+        token.attachToDisplayContent(displayId);
         context.mContextType = CONTEXT_TYPE_SYSTEM_OR_SYSTEM_UI;
-        return context;
-    }
 
-    /**
-     * The overloaded method of {@link #createSystemUiContext(ContextImpl, int)}.
-     * Uses {@Code Display.DEFAULT_DISPLAY} as the target display.
-     */
-    static ContextImpl createSystemUiContext(ContextImpl systemContext) {
-        return createSystemUiContext(systemContext, Display.DEFAULT_DISPLAY);
+        return context;
     }
 
     @UnsupportedAppUsage
@@ -3212,7 +3217,13 @@ class ContextImpl extends Context {
 
     @Override
     public IBinder getWindowContextToken() {
-        return mContextType == CONTEXT_TYPE_WINDOW_CONTEXT ? mToken : null;
+        switch (mContextType) {
+            case CONTEXT_TYPE_WINDOW_CONTEXT:
+            case CONTEXT_TYPE_SYSTEM_OR_SYSTEM_UI:
+                return mToken;
+            default:
+                return null;
+        }
     }
 
     private void checkMode(int mode) {
diff --git a/core/java/android/view/IWindowManager.aidl b/core/java/android/view/IWindowManager.aidl
index 19343acc438a..9000fd00e851 100644
--- a/core/java/android/view/IWindowManager.aidl
+++ b/core/java/android/view/IWindowManager.aidl
@@ -843,6 +843,23 @@ interface IWindowManager
     void attachWindowContextToWindowToken(IBinder clientToken, IBinder token);
 
     /**
+     * Attaches a {@code clientToken} to associate with DisplayContent.
+     * <p>
+     * Note that this API should be invoked after calling
+     * {@link android.window.WindowTokenClient#attachContext(Context)}
+     * </p>
+     *
+     * @param clientToken {@link android.window.WindowContext#getWindowContextToken()
+     * the WindowContext's token}
+     * @param displayId The display associated with the window context
+     *
+     * @return the DisplayContent's {@link android.app.res.Configuration} if the Context is
+     * attached to the DisplayContent successfully. {@code null}, otherwise.
+     * @throws android.view.WindowManager.InvalidDisplayException if the display ID is invalid
+     */
+    Configuration attachToDisplayContent(IBinder clientToken, int displayId);
+
+    /**
      * Detaches {@link android.window.WindowContext} from the window manager node it's currently
      * attached to. It is no-op if the WindowContext is not attached to a window manager node.
      *
diff --git a/core/java/android/window/WindowContextController.java b/core/java/android/window/WindowContextController.java
index 505b45008663..17b675f93f86 100644
--- a/core/java/android/window/WindowContextController.java
+++ b/core/java/android/window/WindowContextController.java
@@ -19,13 +19,10 @@ package android.window;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
-import android.content.res.Configuration;
 import android.os.Bundle;
 import android.os.IBinder;
-import android.os.RemoteException;
 import android.view.IWindowManager;
 import android.view.WindowManager.LayoutParams.WindowType;
-import android.view.WindowManagerGlobal;
 
 import com.android.internal.annotations.VisibleForTesting;
 
@@ -38,7 +35,6 @@ import com.android.internal.annotations.VisibleForTesting;
  * @hide
  */
 public class WindowContextController {
-    private final IWindowManager mWms;
     /**
      * {@code true} to indicate that the {@code mToken} is associated with a
      * {@link com.android.server.wm.DisplayArea}. Note that {@code mToken} is able to attach a
@@ -56,14 +52,7 @@ public class WindowContextController {
      *              {@link Context#getWindowContextToken()}.
      */
     public WindowContextController(@NonNull WindowTokenClient token) {
-        this(token, WindowManagerGlobal.getWindowManagerService());
-    }
-
-    /** Used for test only. DO NOT USE it in production code. */
-    @VisibleForTesting
-    public WindowContextController(@NonNull WindowTokenClient token, IWindowManager mockWms) {
         mToken = token;
-        mWms = mockWms;
     }
 
     /**
@@ -80,18 +69,7 @@ public class WindowContextController {
             throw new IllegalStateException("A Window Context can be only attached to "
                     + "a DisplayArea once.");
         }
-        try {
-            final Configuration configuration = mWms.attachWindowContextToDisplayArea(mToken, type,
-                    displayId, options);
-            if (configuration != null) {
-                mAttachedToDisplayArea = true;
-                // Send the DisplayArea's configuration to WindowContext directly instead of
-                // waiting for dispatching from WMS.
-                mToken.onConfigurationChanged(configuration, displayId);
-            }
-        }  catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
+        mAttachedToDisplayArea = mToken.attachToDisplayArea(type, displayId, options);
     }
 
     /**
@@ -119,22 +97,14 @@ public class WindowContextController {
             throw new IllegalStateException("The Window Context should have been attached"
                     + " to a DisplayArea.");
         }
-        try {
-            mWms.attachWindowContextToWindowToken(mToken, windowToken);
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
+        mToken.attachToWindowToken(windowToken);
     }
 
     /** Detaches the window context from the node it's currently associated with. */
     public void detachIfNeeded() {
         if (mAttachedToDisplayArea) {
-            try {
-                mWms.detachWindowContextFromWindowContainer(mToken);
-                mAttachedToDisplayArea = false;
-            } catch (RemoteException e) {
-                throw e.rethrowFromSystemServer();
-            }
+            mToken.detachFromWindowContainerIfNeeded();
+            mAttachedToDisplayArea = false;
         }
     }
 }
diff --git a/core/java/android/window/WindowTokenClient.java b/core/java/android/window/WindowTokenClient.java
index 4dcd2e74a53f..876272b2b87e 100644
--- a/core/java/android/window/WindowTokenClient.java
+++ b/core/java/android/window/WindowTokenClient.java
@@ -16,6 +16,7 @@
 package android.window;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.app.ActivityThread;
 import android.app.IWindowToken;
 import android.app.ResourcesManager;
@@ -23,6 +24,11 @@ import android.content.Context;
 import android.content.res.Configuration;
 import android.os.Bundle;
 import android.os.IBinder;
+import android.os.RemoteException;
+import android.util.Log;
+import android.view.IWindowManager;
+import android.view.WindowManager.LayoutParams.WindowType;
+import android.view.WindowManagerGlobal;
 
 import com.android.internal.annotations.VisibleForTesting;
 
@@ -48,6 +54,14 @@ public class WindowTokenClient extends IWindowToken.Stub {
 
     private final ResourcesManager mResourcesManager = ResourcesManager.getInstance();
 
+    private IWindowManager mWms;
+
+    private final Configuration mConfiguration = new Configuration();
+
+    private boolean mShouldDumpConfigForIme;
+
+    private boolean mAttachToWindowContainer;
+
     /**
      * Attaches {@code context} to this {@link WindowTokenClient}. Each {@link WindowTokenClient}
      * can only attach one {@link Context}.
@@ -66,6 +80,88 @@ public class WindowTokenClient extends IWindowToken.Stub {
     }
 
     /**
+     * Attaches this {@link WindowTokenClient} to a {@link com.android.server.wm.DisplayArea}.
+     *
+     * @param type The window type of the {@link WindowContext}
+     * @param displayId The {@link Context#getDisplayId() ID of display} to associate with
+     * @param options The window context launched option
+     * @return {@code true} if attaching successfully.
+     */
+    public boolean attachToDisplayArea(@WindowType int type, int displayId,
+            @Nullable Bundle options) {
+        try {
+            final Configuration configuration = getWindowManagerService()
+                    .attachWindowContextToDisplayArea(this, type, displayId, options);
+            if (configuration == null) {
+                return false;
+            }
+            onConfigurationChanged(configuration, displayId);
+            mAttachToWindowContainer = true;
+            return true;
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Attaches this {@link WindowTokenClient} to a {@code DisplayContent}.
+     *
+     * @param displayId The {@link Context#getDisplayId() ID of display} to associate with
+     * @return {@code true} if attaching successfully.
+     */
+    public boolean attachToDisplayContent(int displayId) {
+        final IWindowManager wms = getWindowManagerService();
+        // #createSystemUiContext may call this method before WindowManagerService is initialized.
+        if (wms == null) {
+            return false;
+        }
+        try {
+            final Configuration configuration = wms.attachToDisplayContent(this, displayId);
+            if (configuration == null) {
+                return false;
+            }
+            onConfigurationChanged(configuration, displayId);
+            mAttachToWindowContainer = true;
+            return true;
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Attaches this {@link WindowTokenClient} to a {@code windowToken}.
+     *
+     * @param windowToken the window token to associated with
+     */
+    public void attachToWindowToken(IBinder windowToken) {
+        try {
+            getWindowManagerService().attachWindowContextToWindowToken(this, windowToken);
+            mAttachToWindowContainer = true;
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /** Detaches this {@link WindowTokenClient} from associated WindowContainer if there's one. */
+    public void detachFromWindowContainerIfNeeded() {
+        if (!mAttachToWindowContainer) {
+            return;
+        }
+        try {
+            getWindowManagerService().detachWindowContextFromWindowContainer(this);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    private IWindowManager getWindowManagerService() {
+        if (mWms == null) {
+            mWms = WindowManagerGlobal.getWindowManagerService();
+        }
+        return mWms;
+    }
+
+    /**
      * Called when {@link Configuration} updates from the server side receive.
      *
      * @param newConfig the updated {@link Configuration}
diff --git a/core/tests/coretests/src/android/window/WindowContextControllerTest.java b/core/tests/coretests/src/android/window/WindowContextControllerTest.java
index 073e46827bbb..d7b36a8d7067 100644
--- a/core/tests/coretests/src/android/window/WindowContextControllerTest.java
+++ b/core/tests/coretests/src/android/window/WindowContextControllerTest.java
@@ -23,16 +23,13 @@ import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
-import android.content.res.Configuration;
 import android.os.Binder;
 import android.platform.test.annotations.Presubmit;
-import android.view.IWindowManager;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
@@ -58,17 +55,14 @@ import org.mockito.MockitoAnnotations;
 public class WindowContextControllerTest {
     private WindowContextController mController;
     @Mock
-    private IWindowManager mMockWms;
-    @Mock
     private WindowTokenClient mMockToken;
 
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
-        mController = new WindowContextController(mMockToken, mMockWms);
+        mController = new WindowContextController(mMockToken);
         doNothing().when(mMockToken).onConfigurationChanged(any(), anyInt());
-        doReturn(new Configuration()).when(mMockWms).attachWindowContextToDisplayArea(any(),
-                anyInt(), anyInt(), any());
+        doReturn(true).when(mMockToken).attachToDisplayArea(anyInt(), anyInt(), any());
     }
 
     @Test(expected = IllegalStateException.class)
@@ -80,10 +74,10 @@ public class WindowContextControllerTest {
     }
 
     @Test
-    public void testDetachIfNeeded_NotAttachedYet_DoNothing() throws Exception {
+    public void testDetachIfNeeded_NotAttachedYet_DoNothing() {
         mController.detachIfNeeded();
 
-        verify(mMockWms, never()).detachWindowContextFromWindowContainer(any());
+        verify(mMockToken, never()).detachFromWindowContainerIfNeeded();
     }
 
     @Test
@@ -92,7 +86,6 @@ public class WindowContextControllerTest {
                 null /* options */);
 
         assertThat(mController.mAttachedToDisplayArea).isTrue();
-        verify(mMockToken).onConfigurationChanged(any(), eq(DEFAULT_DISPLAY));
 
         mController.detachIfNeeded();
 
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodMenuController.java b/services/core/java/com/android/server/inputmethod/InputMethodMenuController.java
index 73baf79ea4b1..82b34c35cfd2 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodMenuController.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodMenuController.java
@@ -223,7 +223,7 @@ public class InputMethodMenuController {
     public Context getSettingsContext(int displayId) {
         if (mSettingsContext == null || mSettingsContext.getDisplayId() != displayId) {
             final Context systemUiContext = ActivityThread.currentActivityThread()
-                    .createSystemUiContext(displayId);
+                    .getSystemUiContext(displayId);
             final Context windowContext = systemUiContext.createWindowContext(
                     WindowManager.LayoutParams.TYPE_INPUT_METHOD_DIALOG, null /* options */);
             mSettingsContext = new ContextThemeWrapper(
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 9335846e7805..8a9d5d438a47 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -34,6 +34,7 @@ import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
 import static android.content.res.Configuration.ORIENTATION_PORTRAIT;
 import static android.content.res.Configuration.ORIENTATION_UNDEFINED;
 import static android.os.Build.VERSION_CODES.N;
+import static android.os.Process.SYSTEM_UID;
 import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 import static android.util.DisplayMetrics.DENSITY_DEFAULT;
 import static android.util.RotationUtils.deltaRotation;
@@ -63,6 +64,7 @@ import static android.view.WindowManager.LayoutParams.FLAG_NOT_TOUCHABLE;
 import static android.view.WindowManager.LayoutParams.FLAG_NOT_TOUCH_MODAL;
 import static android.view.WindowManager.LayoutParams.FLAG_SECURE;
 import static android.view.WindowManager.LayoutParams.FLAG_SPLIT_TOUCH;
+import static android.view.WindowManager.LayoutParams.INVALID_WINDOW_TYPE;
 import static android.view.WindowManager.LayoutParams.LAST_APPLICATION_WINDOW;
 import static android.view.WindowManager.LayoutParams.SOFT_INPUT_STATE_ALWAYS_HIDDEN;
 import static android.view.WindowManager.LayoutParams.SOFT_INPUT_STATE_HIDDEN;
@@ -4813,6 +4815,12 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
             reconfigureDisplayLocked();
             onRequestedOverrideConfigurationChanged(getRequestedOverrideConfiguration());
             mWmService.mDisplayNotificationController.dispatchDisplayAdded(this);
+            // Attach the SystemUiContext to this DisplayContent the get latest configuration.
+            // Note that the SystemUiContext will be removed automatically if this DisplayContent
+            // is detached.
+            mWmService.mWindowContextListenerController.registerWindowContainerListener(
+                    getDisplayUiContext().getWindowContextToken(), this, SYSTEM_UID,
+                    INVALID_WINDOW_TYPE, null /* options */);
         }
     }
 
diff --git a/services/core/java/com/android/server/wm/DisplayPolicy.java b/services/core/java/com/android/server/wm/DisplayPolicy.java
index 05356f54b61f..03ea15e2875e 100644
--- a/services/core/java/com/android/server/wm/DisplayPolicy.java
+++ b/services/core/java/com/android/server/wm/DisplayPolicy.java
@@ -475,7 +475,7 @@ public class DisplayPolicy {
                 : service.mContext.createDisplayContext(displayContent.getDisplay());
         mUiContext = displayContent.isDefaultDisplay ? service.mAtmService.mUiContext
                 : service.mAtmService.mSystemThread
-                        .createSystemUiContext(displayContent.getDisplayId());
+                        .getSystemUiContext(displayContent.getDisplayId());
         mDisplayContent = displayContent;
         mLock = service.getWindowManagerLock();
 
@@ -2254,7 +2254,8 @@ public class DisplayPolicy {
         // user's package info (see ContextImpl.createDisplayContext)
         final LoadedApk pi = ActivityThread.currentActivityThread().getPackageInfo(
                 uiContext.getPackageName(), null, 0, userId);
-        mCurrentUserResources = ResourcesManager.getInstance().getResources(null,
+        mCurrentUserResources = ResourcesManager.getInstance().getResources(
+                uiContext.getWindowContextToken(),
                 pi.getResDir(),
                 null /* splitResDirs */,
                 pi.getOverlayDirs(),
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index e159bfbf1811..6844656b6e09 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -2659,7 +2659,7 @@ public class RootWindowContainer extends WindowContainer<DisplayContent>
             // starts. Instead, we expect home activities to be launched when the system is ready
             // (ActivityManagerService#systemReady).
             if (mService.isBooted() || mService.isBooting()) {
-                startSystemDecorations(display.mDisplayContent);
+                startSystemDecorations(display);
             }
         }
     }
diff --git a/services/core/java/com/android/server/wm/WindowContextListenerController.java b/services/core/java/com/android/server/wm/WindowContextListenerController.java
index bc530416c8cd..86e356a876b5 100644
--- a/services/core/java/com/android/server/wm/WindowContextListenerController.java
+++ b/services/core/java/com/android/server/wm/WindowContextListenerController.java
@@ -45,7 +45,7 @@ import java.util.Objects;
  *
  * <ul>
  *   <li>When a {@link WindowContext} is created, it registers the listener via
- *     {@link WindowManagerService#registerWindowContextListener(IBinder, int, int, Bundle)}
+ *     {@link WindowManagerService#attachWindowContextToDisplayArea(IBinder, int, int, Bundle)}
  *     automatically.</li>
  *   <li>When the {@link WindowContext} adds the first window to the screen via
  *     {@link android.view.WindowManager#addView(View, android.view.ViewGroup.LayoutParams)},
@@ -53,7 +53,7 @@ import java.util.Objects;
  *     to corresponding {@link WindowToken} via this controller.</li>
  *   <li>When the {@link WindowContext} is GCed, it unregisters the previously
  *     registered listener via
- *     {@link WindowManagerService#unregisterWindowContextListener(IBinder)}.
+ *     {@link WindowManagerService#detachWindowContextFromWindowContainer(IBinder)}.
  *     {@link WindowManagerService} is also responsible for removing the
  *     {@link WindowContext} created {@link WindowToken}.</li>
  * </ul>
@@ -68,7 +68,7 @@ class WindowContextListenerController {
 
     /**
      * Registers the listener to a {@code container} which is associated with
-     * a {@code clientToken}, which is a {@link android.app.WindowContext} representation. If the
+     * a {@code clientToken}, which is a {@link android.window.WindowContext} representation. If the
      * listener associated with {@code clientToken} hasn't been initialized yet, create one
      * {@link WindowContextListenerImpl}. Otherwise, the listener associated with
      * {@code clientToken} switches to listen to the {@code container}.
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 1289ca3804e6..45e833e84789 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -2723,6 +2723,9 @@ public class WindowManagerService extends IWindowManager.Stub
     @Override
     public Configuration attachWindowContextToDisplayArea(IBinder clientToken, int
             type, int displayId, Bundle options) {
+        if (clientToken == null) {
+            throw new IllegalArgumentException("clientToken must not be null!");
+        }
         final boolean callerCanManageAppTokens = checkCallingPermission(MANAGE_APP_TOKENS,
                 "attachWindowContextToDisplayArea", false /* printLog */);
         final int callingUid = Binder.getCallingUid();
@@ -2813,6 +2816,39 @@ public class WindowManagerService extends IWindowManager.Stub
         }
     }
 
+    @Override
+    public Configuration attachToDisplayContent(IBinder clientToken, int displayId) {
+        if (clientToken == null) {
+            throw new IllegalArgumentException("clientToken must not be null!");
+        }
+        final int callingUid = Binder.getCallingUid();
+        final long origId = Binder.clearCallingIdentity();
+        try {
+            synchronized (mGlobalLock) {
+                // We use "getDisplayContent" instead of "getDisplayContentOrCreate" because
+                // this method may be called in DisplayPolicy's constructor and may cause
+                // infinite loop. In this scenario, we early return here and switch to do the
+                // registration in DisplayContent#onParentChanged at DisplayContent initialization.
+                final DisplayContent dc = mRoot.getDisplayContent(displayId);
+                if (dc == null) {
+                    if (Binder.getCallingPid() != myPid()) {
+                        throw new WindowManager.InvalidDisplayException("attachToDisplayContent: "
+                                + "trying to attach to a non-existing display:" + displayId);
+                    }
+                    // Early return if this method is invoked from system process.
+                    // See above comments for more detail.
+                    return null;
+                }
+
+                mWindowContextListenerController.registerWindowContainerListener(clientToken, dc,
+                        callingUid, INVALID_WINDOW_TYPE, null /* options */);
+                return dc.getConfiguration();
+            }
+        } finally {
+            Binder.restoreCallingIdentity(origId);
+        }
+    }
+
     /** Returns {@code true} if this binder is a registered window token. */
     @Override
     public boolean isWindowToken(IBinder binder) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/InputMethodMenuControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/InputMethodMenuControllerTest.java
index a8ede13e5de6..d7daa57cc9da 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InputMethodMenuControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InputMethodMenuControllerTest.java
@@ -29,7 +29,9 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
 
+import android.app.ActivityThread;
 import android.content.Context;
 import android.content.res.Resources;
 import android.graphics.Rect;
@@ -44,6 +46,7 @@ import android.view.WindowManagerGlobal;
 import com.android.server.inputmethod.InputMethodManagerService;
 import com.android.server.inputmethod.InputMethodMenuController;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -62,6 +65,9 @@ public class InputMethodMenuControllerTest extends WindowTestsBase {
     private InputMethodMenuController mController;
     private DualDisplayAreaGroupPolicyTest.DualDisplayContent mSecondaryDisplay;
 
+    private IWindowManager mIWindowManager;
+    private DisplayManagerGlobal mDisplayManagerGlobal;
+
     @Before
     public void setUp() throws Exception {
         // Let the Display to be created with the DualDisplay policy.
@@ -70,10 +76,12 @@ public class InputMethodMenuControllerTest extends WindowTestsBase {
         Mockito.doReturn(policyProvider).when(mWm).getDisplayAreaPolicyProvider();
 
         mController = new InputMethodMenuController(mock(InputMethodManagerService.class));
+        mSecondaryDisplay = new DualDisplayAreaGroupPolicyTest.DualDisplayContent
+                .Builder(mAtm, 1000, 1000).build();
 
         // Mock addWindowTokenWithOptions to create a test window token.
-        IWindowManager wms = WindowManagerGlobal.getWindowManagerService();
-        spyOn(wms);
+        mIWindowManager = WindowManagerGlobal.getWindowManagerService();
+        spyOn(mIWindowManager);
         doAnswer(invocation -> {
             Object[] args = invocation.getArguments();
             IBinder clientToken = (IBinder) args[0];
@@ -83,19 +91,24 @@ public class InputMethodMenuControllerTest extends WindowTestsBase {
                     dc.getImeContainer(), 1000 /* ownerUid */, TYPE_INPUT_METHOD_DIALOG,
                     null /* options */);
             return dc.getImeContainer().getConfiguration();
-        }).when(wms).attachWindowContextToDisplayArea(any(), eq(TYPE_INPUT_METHOD_DIALOG),
-                anyInt(), any());
-
-        mSecondaryDisplay = new DualDisplayAreaGroupPolicyTest.DualDisplayContent
-                .Builder(mAtm, 1000, 1000).build();
-
-        // Mock DisplayManagerGlobal to return test display when obtaining Display instance.
+        }).when(mIWindowManager).attachWindowContextToDisplayArea(any(),
+                eq(TYPE_INPUT_METHOD_DIALOG), anyInt(), any());
+        mDisplayManagerGlobal = DisplayManagerGlobal.getInstance();
+        spyOn(mDisplayManagerGlobal);
         final int displayId = mSecondaryDisplay.getDisplayId();
         final Display display = mSecondaryDisplay.getDisplay();
-        DisplayManagerGlobal displayManagerGlobal = DisplayManagerGlobal.getInstance();
-        spyOn(displayManagerGlobal);
-        doReturn(display).when(displayManagerGlobal).getCompatibleDisplay(eq(displayId),
+        doReturn(display).when(mDisplayManagerGlobal).getCompatibleDisplay(eq(displayId),
                 (Resources) any());
+        Context systemUiContext = ActivityThread.currentActivityThread()
+                .getSystemUiContext(displayId);
+        spyOn(systemUiContext);
+        doReturn(display).when(systemUiContext).getDisplay();
+    }
+
+    @After
+    public void tearDown() {
+        reset(mIWindowManager);
+        reset(mDisplayManagerGlobal);
     }
 
     @Test