Merge "[RESTRICT AUTOMERGE] Associate SystemUiContext with DisplayContent" into sc-v2-dev

13 files changed, 244 insertions, 101 deletions

diff --git a/core/java/android/app/ActivityThread.java b/core/java/android/app/ActivityThread.java
index 17bb797bf8c5..19a1d7af5382 100644
--- a/core/java/android/app/ActivityThread.java
+++ b/core/java/android/app/ActivityThread.java
@@ -29,6 +29,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 android.window.ConfigurationHelper.diffPublicWithSizeBuckets;
 import static android.window.ConfigurationHelper.freeTextLayoutCachesIfNeeded;
@@ -320,7 +321,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;
@@ -2633,22 +2634,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) {
@@ -3767,7 +3772,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 b5ed1717496e..e76c1398b04a 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
@@ -3206,7 +3211,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 9e41e4d2906c..ae32a481691a 100644
--- a/core/java/android/view/IWindowManager.aidl
+++ b/core/java/android/view/IWindowManager.aidl
@@ -866,6 +866,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 5aa623388574..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,19 +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,
-                        false /* shouldReportConfigChange */);
-            }
-        }  catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
+        mAttachedToDisplayArea = mToken.attachToDisplayArea(type, displayId, options);
     }
 
     /**
@@ -120,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 f3e3859b4256..b331a9e81e27 100644
--- a/core/java/android/window/WindowTokenClient.java
+++ b/core/java/android/window/WindowTokenClient.java
@@ -21,6 +21,7 @@ import static android.window.ConfigurationHelper.isDifferentDisplay;
 import static android.window.ConfigurationHelper.shouldUpdateResources;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.app.ActivityThread;
 import android.app.IWindowToken;
 import android.app.ResourcesManager;
@@ -31,7 +32,11 @@ import android.os.Build;
 import android.os.Bundle;
 import android.os.Debug;
 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;
 
@@ -59,10 +64,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}.
@@ -84,6 +93,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, false /* shouldReportConfigChange */);
+            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, false /* shouldReportConfigChange */);
+            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 a6e351d9cee7..52cb9f318dd0 100644
--- a/core/tests/coretests/src/android/window/WindowContextControllerTest.java
+++ b/core/tests/coretests/src/android/window/WindowContextControllerTest.java
@@ -24,16 +24,13 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 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;
@@ -59,17 +56,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(), anyBoolean());
-        doReturn(new Configuration()).when(mMockWms).attachWindowContextToDisplayArea(any(),
-                anyInt(), anyInt(), any());
+        doReturn(true).when(mMockToken).attachToDisplayArea(anyInt(), anyInt(), any());
     }
 
     @Test(expected = IllegalStateException.class)
@@ -81,10 +75,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
@@ -93,8 +87,6 @@ public class WindowContextControllerTest {
                 null /* options */);
 
         assertThat(mController.mAttachedToDisplayArea).isTrue();
-        verify(mMockToken).onConfigurationChanged(any(), eq(DEFAULT_DISPLAY),
-                eq(false) /* shouldReportConfigChange */);
 
         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 07ab57145050..36579d96d6c5 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;
@@ -62,6 +63,7 @@ import static android.view.WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE;
 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_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;
@@ -4997,6 +4999,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 eb829573e45d..b2657e82eb6c 100644
--- a/services/core/java/com/android/server/wm/DisplayPolicy.java
+++ b/services/core/java/com/android/server/wm/DisplayPolicy.java
@@ -452,7 +452,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();
 
@@ -2260,7 +2260,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 b1eca9d5d4e4..3ffa62dbbe7e 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -2585,7 +2585,7 @@ 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);
             }
             // Drop any cached DisplayInfos associated with this display id - the values are now
             // out of date given this display added event.
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 943a3c12e5ed..cc7485c317c8 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -2741,6 +2741,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();
@@ -2831,6 +2834,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