Add UserSwitchObserver.onBeforeUserSwitching and make SysUI use it.

During a user switch, allow SystemUI to update its source of truth
for current user before WallpaperManagerService, to put wallpaper
colors changed events in order. Otherwise, ThemeOverlayController
receives onColorsChanged call before UserTrackerImpl updates current
user id and thinks colors are changed for non-current user.

Bug: 274925897
Test: atest UserTrackerImplTest
Test: atest UserControllerTest
Change-Id: Ib5e467a18d75d633b2ebaecbaf00fb65f8be6c77

6 files changed, 24 insertions, 3 deletions

diff --git a/core/java/android/app/IUserSwitchObserver.aidl b/core/java/android/app/IUserSwitchObserver.aidl
index 234da8f36e96..cfdb426d6026 100644
--- a/core/java/android/app/IUserSwitchObserver.aidl
+++ b/core/java/android/app/IUserSwitchObserver.aidl
@@ -20,6 +20,7 @@ import android.os.IRemoteCallback;
 
 /** {@hide} */
 oneway interface IUserSwitchObserver {
+    void onBeforeUserSwitching(int newUserId);
     void onUserSwitching(int newUserId, IRemoteCallback reply);
     void onUserSwitchComplete(int newUserId);
     void onForegroundProfileSwitch(int newProfileId);
diff --git a/core/java/android/app/UserSwitchObserver.java b/core/java/android/app/UserSwitchObserver.java
index 6abc4f09ba38..727799a1f948 100644
--- a/core/java/android/app/UserSwitchObserver.java
+++ b/core/java/android/app/UserSwitchObserver.java
@@ -30,6 +30,9 @@ public class UserSwitchObserver extends IUserSwitchObserver.Stub {
     }
 
     @Override
+    public void onBeforeUserSwitching(int newUserId) throws RemoteException {}
+
+    @Override
     public void onUserSwitching(int newUserId, IRemoteCallback reply) throws RemoteException {
         if (reply != null) {
             reply.sendResult(null);
diff --git a/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt b/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt
index 0b2ae05b7c9b..72286f175671 100644
--- a/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt
@@ -161,6 +161,10 @@ open class UserTrackerImpl internal constructor(
 
     private fun registerUserSwitchObserver() {
         iActivityManager.registerUserSwitchObserver(object : UserSwitchObserver() {
+            override fun onBeforeUserSwitching(newUserId: Int) {
+                setUserIdInternal(newUserId)
+            }
+
             override fun onUserSwitching(newUserId: Int, reply: IRemoteCallback?) {
                 backgroundHandler.run {
                     handleUserSwitching(newUserId)
@@ -181,8 +185,6 @@ open class UserTrackerImpl internal constructor(
         Assert.isNotMainThread()
         Log.i(TAG, "Switching to user $newUserId")
 
-        setUserIdInternal(newUserId)
-
         val list = synchronized(callbacks) {
             callbacks.toList()
         }
@@ -205,7 +207,6 @@ open class UserTrackerImpl internal constructor(
         Assert.isNotMainThread()
         Log.i(TAG, "Switched to user $newUserId")
 
-        setUserIdInternal(newUserId)
         notifySubscribers {
             onUserChanged(newUserId, userContext)
             onProfilesChanged(userProfiles)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt
index 71ba21538a8e..aa98f08e9015 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt
@@ -167,6 +167,7 @@ class UserTrackerImplTest : SysuiTestCase() {
 
         val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
         verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
+        captor.value.onBeforeUserSwitching(newID)
         captor.value.onUserSwitching(newID, userSwitchingReply)
         verify(userSwitchingReply).sendResult(any())
 
@@ -290,6 +291,7 @@ class UserTrackerImplTest : SysuiTestCase() {
 
         val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
         verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
+        captor.value.onBeforeUserSwitching(newID)
         captor.value.onUserSwitching(newID, userSwitchingReply)
         verify(userSwitchingReply).sendResult(any())
 
@@ -308,6 +310,7 @@ class UserTrackerImplTest : SysuiTestCase() {
 
         val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
         verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
+        captor.value.onBeforeUserSwitching(newID)
         captor.value.onUserSwitchComplete(newID)
 
         assertThat(callback.calledOnUserChanged).isEqualTo(1)
diff --git a/services/core/java/com/android/server/am/UserController.java b/services/core/java/com/android/server/am/UserController.java
index 1b378837e558..d926c2c7c7a8 100644
--- a/services/core/java/com/android/server/am/UserController.java
+++ b/services/core/java/com/android/server/am/UserController.java
@@ -2141,6 +2141,17 @@ class UserController implements Handler.Callback {
 
         final int observerCount = mUserSwitchObservers.beginBroadcast();
         if (observerCount > 0) {
+            for (int i = 0; i < observerCount; i++) {
+                final String name = "#" + i + " " + mUserSwitchObservers.getBroadcastCookie(i);
+                t.traceBegin("onBeforeUserSwitching-" + name);
+                try {
+                    mUserSwitchObservers.getBroadcastItem(i).onBeforeUserSwitching(newUserId);
+                } catch (RemoteException e) {
+                    // Ignore
+                } finally {
+                    t.traceEnd();
+                }
+            }
             final ArraySet<String> curWaitingUserSwitchCallbacks = new ArraySet<>();
             synchronized (mLock) {
                 uss.switching = true;
diff --git a/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java b/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
index 8994a488bd56..ab8f3f2279fe 100644
--- a/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
@@ -383,6 +383,7 @@ public class UserControllerTest {
         // Call dispatchUserSwitch and verify that observer was called only once
         mInjector.mHandler.clearAllRecordedMessages();
         mUserController.dispatchUserSwitch(userState, oldUserId, newUserId);
+        verify(observer, times(1)).onBeforeUserSwitching(eq(TEST_USER_ID));
         verify(observer, times(1)).onUserSwitching(eq(TEST_USER_ID), any());
         Set<Integer> expectedCodes = Collections.singleton(CONTINUE_USER_SWITCH_MSG);
         Set<Integer> actualCodes = mInjector.mHandler.getMessageCodes();
@@ -413,6 +414,7 @@ public class UserControllerTest {
         // Call dispatchUserSwitch and verify that observer was called only once
         mInjector.mHandler.clearAllRecordedMessages();
         mUserController.dispatchUserSwitch(userState, oldUserId, newUserId);
+        verify(observer, times(1)).onBeforeUserSwitching(eq(TEST_USER_ID));
         verify(observer, times(1)).onUserSwitching(eq(TEST_USER_ID), any());
         // Verify that CONTINUE_USER_SWITCH_MSG is not sent (triggers timeout)
         Set<Integer> actualCodes = mInjector.mHandler.getMessageCodes();