Introduce AMS.getLastResumedActivityUserId to fix work profile icon issue

Problem:
Work profile status bar icon feeature is relied on two callbacks
1. onForegroundProfileSwitch  (AMS.setResumedActivityLocked)
2. appTransitionStarting (WMS)

We assume callback 1 is always called before 2, but it is not the case.
These two callbacks are triggered by two handlers in two different threads,
and hence race condition happens.

Solution:
Not rely on onForegroundProfileSwitch to update mManagedProfileFocused
flag anymore. Query getLastResumedActivityUserId in appTransitionStarting.
Also, make sure mLastResumedActivity is updated before sending message
to WMS in setResumedActivityLocked.

Test: Start a work app, observe that the work icon is shown.
Test: Start a personal app, observe that work icon is gone.
Test: Dock the work app, tap on it (give it focus), observe that work
      icon is shown.
Test: Start a work app, switch user, can see the icon is gone. Switch back,
      icon is back.

Bug: 34159089

Change-Id: I2cee141d18e8b7d5607b26dd7a2fd5bc9cd0ebb3

3 files changed, 54 insertions, 69 deletions

diff --git a/core/java/android/app/IActivityManager.aidl b/core/java/android/app/IActivityManager.aidl
index 82229d5b1bc8..b9d1d91fabda 100644
--- a/core/java/android/app/IActivityManager.aidl
+++ b/core/java/android/app/IActivityManager.aidl
@@ -629,6 +629,11 @@ interface IActivityManager {
      */
     void setDisablePreviewScreenshots(IBinder token, boolean disable);
 
+    /**
+     * Return the user id of last resumed activity.
+     */
+    int getLastResumedActivityUserId();
+
     // WARNING: when these transactions are updated, check if they are any callers on the native
     // side. If so, make sure they are using the correct transaction ids and arguments.
     // If a transaction which will also be used on the native side is being inserted, add it
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
index bb6c8f2e1e73..53ec8c5048ad 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
@@ -127,7 +127,6 @@ public class PhoneStatusBarPolicy implements Callback, Callbacks,
     private boolean mVolumeVisible;
     private boolean mCurrentUserSetup;
 
-    private boolean mManagedProfileFocused = false;
     private boolean mManagedProfileIconVisible = false;
     private boolean mManagedProfileInQuietMode = false;
 
@@ -439,45 +438,30 @@ public class PhoneStatusBarPolicy implements Callback, Callbacks,
         }
     }
 
-    private void profileChanged(int userId) {
-        UserInfo user = null;
-        if (userId == UserHandle.USER_CURRENT) {
-            try {
-                user = ActivityManager.getService().getCurrentUser();
-            } catch (RemoteException e) {
-                // Ignore
-            }
-        } else {
-            user = mUserManager.getUserInfo(userId);
-        }
-
-        mManagedProfileFocused = user != null && user.isManagedProfile();
-        if (DEBUG) Log.v(TAG, "profileChanged: mManagedProfileFocused: " + mManagedProfileFocused);
-        // Actually update the icon later when transition starts.
-    }
-
     private void updateManagedProfile() {
-        if (DEBUG) {
-            Log.v(TAG, "updateManagedProfile: mManagedProfileFocused: "
-                    + mManagedProfileFocused);
-        }
-        final boolean showIcon;
-        if (mManagedProfileFocused && !mKeyguardMonitor.isShowing()) {
-            showIcon = true;
-            mIconController.setIcon(mSlotManagedProfile,
-                    R.drawable.stat_sys_managed_profile_status,
-                    mContext.getString(R.string.accessibility_managed_profile));
-        } else if (mManagedProfileInQuietMode) {
-            showIcon = true;
-            mIconController.setIcon(mSlotManagedProfile,
-                    R.drawable.stat_sys_managed_profile_status_off,
-                    mContext.getString(R.string.accessibility_managed_profile));
-        } else {
-            showIcon = false;
-        }
-        if (mManagedProfileIconVisible != showIcon) {
-            mIconController.setIconVisibility(mSlotManagedProfile, showIcon);
-            mManagedProfileIconVisible = showIcon;
+        try {
+            final boolean showIcon;
+            final int userId = ActivityManager.getService().getLastResumedActivityUserId();
+            if (mUserManager.isManagedProfile(userId) && !mKeyguardMonitor.isShowing()) {
+                showIcon = true;
+                mIconController.setIcon(mSlotManagedProfile,
+                        R.drawable.stat_sys_managed_profile_status,
+                        mContext.getString(R.string.accessibility_managed_profile));
+            } else if (mManagedProfileInQuietMode) {
+                showIcon = true;
+                mIconController.setIcon(mSlotManagedProfile,
+                        R.drawable.stat_sys_managed_profile_status_off,
+                        mContext.getString(R.string.accessibility_managed_profile));
+            } else {
+                showIcon = false;
+            }
+            if (mManagedProfileIconVisible != showIcon) {
+                mIconController.setIconVisibility(mSlotManagedProfile, showIcon);
+                mManagedProfileIconVisible = showIcon;
+            }
+        } catch (RemoteException ex) {
+            Log.w(TAG, "updateManagedProfile: ", ex);
+            // ignore
         }
     }
 
@@ -556,35 +540,16 @@ public class PhoneStatusBarPolicy implements Callback, Callbacks,
             new SynchronousUserSwitchObserver() {
                 @Override
                 public void onUserSwitching(int newUserId) throws RemoteException {
-                    mHandler.post(new Runnable() {
-                        @Override
-                        public void run() {
-                            mUserInfoController.reloadUserInfo();
-                        }
-                    });
+                    mHandler.post(() -> mUserInfoController.reloadUserInfo());
                 }
 
                 @Override
                 public void onUserSwitchComplete(int newUserId) throws RemoteException {
-                    mHandler.post(new Runnable() {
-                        @Override
-                        public void run() {
-                            updateAlarm();
-                            profileChanged(newUserId);
-                            updateQuietState();
-                            updateManagedProfile();
-                            updateForegroundInstantApps();
-                        }
-                    });
-                }
-
-                @Override
-                public void onForegroundProfileSwitch(int newProfileId) {
-                    mHandler.post(new Runnable() {
-                        @Override
-                        public void run() {
-                            profileChanged(newProfileId);
-                        }
+                    mHandler.post(() -> {
+                        updateAlarm();
+                        updateQuietState();
+                        updateManagedProfile();
+                        updateForegroundInstantApps();
                     });
                 }
             };
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index ead1b26a0a1f..8cb0eee07b3d 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -3212,18 +3212,18 @@ public class ActivityManagerService extends IActivityManager.Stub
             }
         }
 
-        mWindowManager.setFocusedApp(r.appToken, true);
-
-        applyUpdateLockStateLocked(r);
-        applyUpdateVrModeLocked(r);
         if (mLastResumedActivity != null && r.userId != mLastResumedActivity.userId) {
             mHandler.removeMessages(FOREGROUND_PROFILE_CHANGED_MSG);
             mHandler.obtainMessage(
                     FOREGROUND_PROFILE_CHANGED_MSG, r.userId, 0).sendToTarget();
         }
-
         mLastResumedActivity = r;
 
+        mWindowManager.setFocusedApp(r.appToken, true);
+
+        applyUpdateLockStateLocked(r);
+        applyUpdateVrModeLocked(r);
+
         EventLogTags.writeAmSetResumedActivity(
                 r == null ? -1 : r.userId,
                 r == null ? "NULL" : r.shortComponentName,
@@ -23599,6 +23599,21 @@ public class ActivityManagerService extends IActivityManager.Stub
         }
     }
 
+    /**
+     * Return the user id of the last resumed activity.
+     */
+    @Override
+    public @UserIdInt int getLastResumedActivityUserId() {
+        enforceCallingPermission(
+                permission.INTERACT_ACROSS_USERS_FULL, "getLastResumedActivityUserId()");
+        synchronized (this) {
+            if (mLastResumedActivity == null) {
+                return mUserController.getCurrentUserIdLocked();
+            }
+            return mLastResumedActivity.userId;
+        }
+    }
+
     private final class SleepTokenImpl extends SleepToken {
         private final String mTag;
         private final long mAcquireTime;