Merge "Keep track of SystemUiContext before WMS is initialized" into main

5 files changed, 75 insertions, 30 deletions

diff --git a/core/java/android/window/WindowTokenClient.java b/core/java/android/window/WindowTokenClient.java
index 15adc80add01..6e76d8d345b2 100644
--- a/core/java/android/window/WindowTokenClient.java
+++ b/core/java/android/window/WindowTokenClient.java
@@ -19,6 +19,8 @@ import static android.window.ConfigurationHelper.freeTextLayoutCachesIfNeeded;
 import static android.window.ConfigurationHelper.isDifferentDisplay;
 import static android.window.ConfigurationHelper.shouldUpdateResources;
 
+import static com.android.internal.annotations.VisibleForTesting.Visibility.PACKAGE;
+
 import android.annotation.AnyThread;
 import android.annotation.MainThread;
 import android.annotation.NonNull;
@@ -104,7 +106,7 @@ public class WindowTokenClient extends Binder {
      * @param newConfig the updated {@link Configuration}
      * @param newDisplayId the updated {@link android.view.Display} ID
      */
-    @VisibleForTesting
+    @VisibleForTesting(visibility = PACKAGE)
     @MainThread
     public void onConfigurationChanged(Configuration newConfig, int newDisplayId) {
         onConfigurationChanged(newConfig, newDisplayId, true /* shouldReportConfigChange */);
@@ -113,7 +115,7 @@ public class WindowTokenClient extends Binder {
     /**
      * Posts an {@link #onConfigurationChanged} to the main thread.
      */
-    @VisibleForTesting
+    @VisibleForTesting(visibility = PACKAGE)
     public void postOnConfigurationChanged(@NonNull Configuration newConfig, int newDisplayId) {
         mHandler.post(PooledLambda.obtainRunnable(this::onConfigurationChanged, newConfig,
                 newDisplayId, true /* shouldReportConfigChange */).recycleOnUse());
@@ -232,7 +234,7 @@ public class WindowTokenClient extends Binder {
     /**
      * Called when the attached window is removed from the display.
      */
-    @VisibleForTesting
+    @VisibleForTesting(visibility = PACKAGE)
     @MainThread
     public void onWindowTokenRemoved() {
         final Context context = mContextRef.get();
diff --git a/core/java/android/window/WindowTokenClientController.java b/core/java/android/window/WindowTokenClientController.java
index abf7bb15b3d8..fa345956ec4d 100644
--- a/core/java/android/window/WindowTokenClientController.java
+++ b/core/java/android/window/WindowTokenClientController.java
@@ -28,7 +28,7 @@ import android.content.Context;
 import android.os.Bundle;
 import android.os.IBinder;
 import android.os.RemoteException;
-import android.util.ArrayMap;
+import android.util.ArraySet;
 import android.util.Log;
 import android.view.IWindowManager;
 import android.view.WindowManagerGlobal;
@@ -50,9 +50,9 @@ public class WindowTokenClientController {
     private final IApplicationThread mAppThread = ActivityThread.currentActivityThread()
             .getApplicationThread();
 
-    /** Mapping from a client defined token to the {@link WindowTokenClient} it represents. */
+    /** Attached {@link WindowTokenClient}. */
     @GuardedBy("mLock")
-    private final ArrayMap<IBinder, WindowTokenClient> mWindowTokenClientMap = new ArrayMap<>();
+    private final ArraySet<WindowTokenClient> mWindowTokenClients = new ArraySet<>();
 
     /** Gets the singleton controller. */
     @NonNull
@@ -85,11 +85,15 @@ public class WindowTokenClientController {
     /** Gets the {@link WindowContext} instance for the token. */
     @Nullable
     public Context getWindowContext(@NonNull IBinder clientToken) {
-        final WindowTokenClient windowTokenClient;
+        if (!(clientToken instanceof WindowTokenClient windowTokenClient)) {
+            return null;
+        }
         synchronized (mLock) {
-            windowTokenClient = mWindowTokenClientMap.get(clientToken);
+            if (!mWindowTokenClients.contains(windowTokenClient)) {
+                return null;
+            }
         }
-        return windowTokenClient != null ? windowTokenClient.getContext() : null;
+        return windowTokenClient.getContext();
     }
 
     /**
@@ -126,8 +130,14 @@ public class WindowTokenClientController {
      */
     public boolean attachToDisplayContent(@NonNull WindowTokenClient client, int displayId) {
         final IWindowManager wms = getWindowManagerService();
-        // #createSystemUiContext may call this method before WindowManagerService is initialized.
         if (wms == null) {
+            // #createSystemUiContext may call this method before WindowManagerService is
+            // initialized.
+            // Regardless of whether or not it is ready, keep track of the token so that when WMS
+            // is initialized later, the SystemUiContext will start reporting from
+            // DisplayContent#registerSystemUiContext, and WindowTokenClientController can report
+            // the Configuration to the correct client.
+            recordWindowContextToken(client);
             return false;
         }
         final WindowContextInfo info;
@@ -170,12 +180,18 @@ public class WindowTokenClientController {
     /** Detaches a {@link WindowTokenClient} from associated WindowContainer if there's one. */
     public void detachIfNeeded(@NonNull WindowTokenClient client) {
         synchronized (mLock) {
-            if (mWindowTokenClientMap.remove(client) == null) {
+            if (!mWindowTokenClients.remove(client)) {
                 return;
             }
         }
+        final IWindowManager wms = getWindowManagerService();
+        if (wms == null) {
+            // #createSystemUiContext may call this method before WindowManagerService is
+            // initialized. If it is GC'ed before WMS is initialized, skip calling into WMS.
+            return;
+        }
         try {
-            getWindowManagerService().detachWindowContext(client);
+            wms.detachWindowContext(client);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
@@ -183,9 +199,7 @@ public class WindowTokenClientController {
 
     private void onWindowContextTokenAttached(@NonNull WindowTokenClient client,
             @NonNull WindowContextInfo info, boolean shouldReportConfigChange) {
-        synchronized (mLock) {
-            mWindowTokenClientMap.put(client, client);
-        }
+        recordWindowContextToken(client);
         if (shouldReportConfigChange) {
             // Should trigger an #onConfigurationChanged callback to the WindowContext. Post the
             // dispatch in the next loop to prevent the callback from being dispatched before
@@ -199,10 +213,16 @@ public class WindowTokenClientController {
         }
     }
 
+    private void recordWindowContextToken(@NonNull WindowTokenClient client) {
+        synchronized (mLock) {
+            mWindowTokenClients.add(client);
+        }
+    }
+
     /** Called when receives {@link WindowContextInfoChangeItem}. */
     public void onWindowContextInfoChanged(@NonNull IBinder clientToken,
             @NonNull WindowContextInfo info) {
-        final WindowTokenClient windowTokenClient = getWindowTokenClient(clientToken);
+        final WindowTokenClient windowTokenClient = getWindowTokenClientIfAttached(clientToken);
         if (windowTokenClient != null) {
             windowTokenClient.onConfigurationChanged(info.getConfiguration(), info.getDisplayId());
         }
@@ -210,20 +230,23 @@ public class WindowTokenClientController {
 
     /** Called when receives {@link WindowContextWindowRemovalItem}. */
     public void onWindowContextWindowRemoved(@NonNull IBinder clientToken) {
-        final WindowTokenClient windowTokenClient = getWindowTokenClient(clientToken);
+        final WindowTokenClient windowTokenClient = getWindowTokenClientIfAttached(clientToken);
         if (windowTokenClient != null) {
             windowTokenClient.onWindowTokenRemoved();
         }
     }
 
     @Nullable
-    private WindowTokenClient getWindowTokenClient(@NonNull IBinder clientToken) {
-        final WindowTokenClient windowTokenClient;
-        synchronized (mLock) {
-            windowTokenClient = mWindowTokenClientMap.get(clientToken);
+    private WindowTokenClient getWindowTokenClientIfAttached(@NonNull IBinder clientToken) {
+        if (!(clientToken instanceof WindowTokenClient windowTokenClient)) {
+            Log.e(TAG, "getWindowTokenClient failed for non-window token " + clientToken);
+            return null;
         }
-        if (windowTokenClient == null) {
-            Log.w(TAG, "Can't find attached WindowTokenClient for " + clientToken);
+        synchronized (mLock) {
+            if (!mWindowTokenClients.contains(windowTokenClient)) {
+                Log.w(TAG, "Can't find attached WindowTokenClient for " + clientToken);
+                return null;
+            }
         }
         return windowTokenClient;
     }
diff --git a/core/tests/coretests/src/android/window/WindowTokenClientControllerTest.java b/core/tests/coretests/src/android/window/WindowTokenClientControllerTest.java
index a3725af04ba6..bb2fe1bcfc64 100644
--- a/core/tests/coretests/src/android/window/WindowTokenClientControllerTest.java
+++ b/core/tests/coretests/src/android/window/WindowTokenClientControllerTest.java
@@ -162,6 +162,22 @@ public class WindowTokenClientControllerTest {
     }
 
     @Test
+    public void testAttachToDisplayContent_keepTrackWithoutWMS() {
+        // WMS is not initialized
+        doReturn(null).when(mController).getWindowManagerService();
+
+        assertFalse(mController.attachToDisplayContent(mWindowTokenClient, DEFAULT_DISPLAY));
+
+        // Can report config change
+        mController.onWindowContextInfoChanged(mWindowTokenClient, mWindowContextInfo);
+
+        verify(mWindowTokenClient).onConfigurationChanged(mConfiguration, DEFAULT_DISPLAY);
+
+        // No crash to detach even if WMS is not initialized.
+        mController.detachIfNeeded(mWindowTokenClient);
+    }
+
+    @Test
     public void testAttachToWindowToken() throws RemoteException {
         doReturn(null).when(mWindowManagerService).attachWindowContextToWindowToken(
                 any(), any(), any());
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index e9e550e72a00..9a33df13bb6a 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -5502,14 +5502,18 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
             // Attach the SystemUiContext to this DisplayContent the get latest configuration.
             // Note that the SystemUiContext will be removed automatically if this DisplayContent
             // is detached.
-            final WindowProcessController wpc = mAtmService.getProcessController(
-                    getDisplayUiContext().getIApplicationThread());
-            mWmService.mWindowContextListenerController.registerWindowContainerListener(
-                    wpc, getDisplayUiContext().getWindowContextToken(), this,
-                    INVALID_WINDOW_TYPE, null /* options */);
+            registerSystemUiContext();
         }
     }
 
+    private void registerSystemUiContext() {
+        final WindowProcessController wpc = mAtmService.getProcessController(
+                getDisplayUiContext().getIApplicationThread());
+        mWmService.mWindowContextListenerController.registerWindowContainerListener(
+                wpc, getDisplayUiContext().getWindowContextToken(), this,
+                INVALID_WINDOW_TYPE, null /* options */);
+    }
+
     @Override
     void assignChildLayers(SurfaceControl.Transaction t) {
         assignRelativeLayerForIme(t, false /* forceUpdate */);
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 8268cae12e3d..63f29f639204 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -3030,8 +3030,8 @@ public class WindowManagerService extends IWindowManager.Stub
 
                 mWindowContextListenerController.unregisterWindowContainerListener(clientToken);
 
-                final WindowToken token = wc.asWindowToken();
-                if (token != null && token.isFromClient()) {
+                final WindowToken token = wc != null ? wc.asWindowToken() : null;
+                if (token != null && token.isFromClient() && token.getDisplayContent() != null) {
                     removeWindowToken(token.token, token.getDisplayContent().getDisplayId());
                 }
             }