Merge cherrypicks of ['googleplex-android-review.googlesource.com/27929467'] into 24Q2-release.

Change-Id: I021b957565808d55675d2a853887aaaf38dc94be
diff --git a/services/core/java/com/android/server/am/UserController.java b/services/core/java/com/android/server/am/UserController.java
index 3abfe082..a7472d8 100644
--- a/services/core/java/com/android/server/am/UserController.java
+++ b/services/core/java/com/android/server/am/UserController.java
@@ -154,9 +154,6 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 
@@ -223,9 +220,18 @@
     private static final int USER_SWITCH_CALLBACKS_TIMEOUT_MS = 5 * 1000;
 
     /**
+     * Amount of time waited for
+     * {@link ActivityTaskManagerInternal.ScreenObserver#onKeyguardStateChanged} callbacks to be
+     * called after calling {@link WindowManagerService#lockDeviceNow}.
+     * Otherwise, we should throw a {@link RuntimeException} and never dismiss the
+     * {@link UserSwitchingDialog}.
+     */
+    static final int SHOW_KEYGUARD_TIMEOUT_MS = 20 * 1000;
+
+    /**
      * Amount of time waited for {@link WindowManagerService#dismissKeyguard} callbacks to be
      * called after dismissing the keyguard.
-     * Otherwise, we should move on to dismiss the dialog {@link #dismissUserSwitchDialog()}
+     * Otherwise, we should move on to dismiss the dialog {@link #dismissUserSwitchDialog}}
      * and report user switch is complete {@link #REPORT_USER_SWITCH_COMPLETE_MSG}.
      */
     private static final int DISMISS_KEYGUARD_TIMEOUT_MS = 2 * 1000;
@@ -1822,15 +1828,8 @@
                 updateProfileRelatedCaches();
                 mInjector.getWindowManager().setCurrentUser(userId);
                 mInjector.reportCurWakefulnessUsageEvent();
-                // Once the internal notion of the active user has switched, we lock the device
-                // with the option to show the user switcher on the keyguard.
                 if (userSwitchUiEnabled) {
                     mInjector.getWindowManager().setSwitchingUser(true);
-                    // Only lock if the user has a secure keyguard PIN/Pattern/Pwd
-                    if (mInjector.getKeyguardManager().isDeviceSecure(userId)) {
-                        // Make sure the device is locked before moving on with the user switch
-                        mInjector.lockDeviceNowAndWaitForKeyguardShown();
-                    }
                 }
 
             } else {
@@ -2341,32 +2340,54 @@
 
     @VisibleForTesting
     void completeUserSwitch(int oldUserId, int newUserId) {
-        final boolean isUserSwitchUiEnabled = isUserSwitchUiEnabled();
-        // serialize each conditional step
-        await(
-                // STEP 1 - If there is no challenge set, dismiss the keyguard right away
-                isUserSwitchUiEnabled && !mInjector.getKeyguardManager().isDeviceSecure(newUserId),
-                mInjector::dismissKeyguard,
-                () -> await(
-                        // STEP 2 - If user switch ui was enabled, dismiss user switch dialog
-                        isUserSwitchUiEnabled,
-                        this::dismissUserSwitchDialog,
-                        () -> {
-                            // STEP 3 - Send REPORT_USER_SWITCH_COMPLETE_MSG to broadcast
-                            // ACTION_USER_SWITCHED & call UserSwitchObservers.onUserSwitchComplete
-                            mHandler.removeMessages(REPORT_USER_SWITCH_COMPLETE_MSG);
-                            mHandler.sendMessage(mHandler.obtainMessage(
-                                    REPORT_USER_SWITCH_COMPLETE_MSG, oldUserId, newUserId));
-                        }
-                ));
+        final Runnable sendUserSwitchCompleteMessage = () -> {
+            mHandler.removeMessages(REPORT_USER_SWITCH_COMPLETE_MSG);
+            mHandler.sendMessage(mHandler.obtainMessage(
+                    REPORT_USER_SWITCH_COMPLETE_MSG, oldUserId, newUserId));
+        };
+        if (isUserSwitchUiEnabled()) {
+            if (mInjector.getKeyguardManager().isDeviceSecure(newUserId)) {
+                this.showKeyguard(() -> dismissUserSwitchDialog(sendUserSwitchCompleteMessage));
+            } else {
+                this.dismissKeyguard(() -> dismissUserSwitchDialog(sendUserSwitchCompleteMessage));
+            }
+        } else {
+            sendUserSwitchCompleteMessage.run();
+        }
     }
 
-    private void await(boolean condition, Consumer<Runnable> conditionalStep, Runnable nextStep) {
-        if (condition) {
-            conditionalStep.accept(nextStep);
-        } else {
-            nextStep.run();
-        }
+    protected void showKeyguard(Runnable runnable) {
+        runWithTimeout(mInjector::showKeyguard, SHOW_KEYGUARD_TIMEOUT_MS, runnable, () -> {
+            throw new RuntimeException(
+                    "Keyguard is not shown in " + SHOW_KEYGUARD_TIMEOUT_MS + " ms.");
+        }, "showKeyguard");
+    }
+
+    protected void dismissKeyguard(Runnable runnable) {
+        runWithTimeout(mInjector::dismissKeyguard, DISMISS_KEYGUARD_TIMEOUT_MS, runnable, runnable,
+                "dismissKeyguard");
+    }
+
+    private void runWithTimeout(Consumer<Runnable> task, int timeoutMs, Runnable onSuccess,
+            Runnable onTimeout, String traceMsg) {
+        final AtomicInteger state = new AtomicInteger(0); // state = 0 (RUNNING)
+
+        asyncTraceBegin(traceMsg, 0);
+
+        mHandler.postDelayed(() -> {
+            if (state.compareAndSet(0, 1)) { // state = 1 (TIMEOUT)
+                asyncTraceEnd(traceMsg, 0);
+                Slogf.w(TAG, "Timeout: %s did not finish in %d ms", traceMsg, timeoutMs);
+                onTimeout.run();
+            }
+        }, timeoutMs);
+
+        task.accept(() -> {
+            if (state.compareAndSet(0, 2)) { // state = 2 (SUCCESS)
+                asyncTraceEnd(traceMsg, 0);
+                onSuccess.run();
+            }
+        });
     }
 
     private void moveUserToForeground(UserState uss, int newUserId) {
@@ -3805,29 +3826,45 @@
             return IStorageManager.Stub.asInterface(ServiceManager.getService("mount"));
         }
 
-        protected void dismissKeyguard(Runnable runnable) {
-            final AtomicBoolean isFirst = new AtomicBoolean(true);
-            final Runnable runOnce = () -> {
-                if (isFirst.getAndSet(false)) {
-                    runnable.run();
-                }
-            };
+        protected void showKeyguard(Runnable runnable) {
+            if (getWindowManager().isKeyguardLocked()) {
+                runnable.run();
+                return;
+            }
+            getActivityTaskManagerInternal().registerScreenObserver(
+                    new ActivityTaskManagerInternal.ScreenObserver() {
+                        @Override
+                        public void onAwakeStateChanged(boolean isAwake) {
 
-            mHandler.postDelayed(runOnce, DISMISS_KEYGUARD_TIMEOUT_MS);
+                        }
+
+                        @Override
+                        public void onKeyguardStateChanged(boolean isShowing) {
+                            if (isShowing) {
+                                getActivityTaskManagerInternal().unregisterScreenObserver(this);
+                                runnable.run();
+                            }
+                        }
+                    }
+            );
+            getWindowManager().lockDeviceNow();
+        }
+
+        protected void dismissKeyguard(Runnable runnable) {
             getWindowManager().dismissKeyguard(new IKeyguardDismissCallback.Stub() {
                 @Override
                 public void onDismissError() throws RemoteException {
-                    mHandler.post(runOnce);
+                    runnable.run();
                 }
 
                 @Override
                 public void onDismissSucceeded() throws RemoteException {
-                    mHandler.post(runOnce);
+                    runnable.run();
                 }
 
                 @Override
                 public void onDismissCancelled() throws RemoteException {
-                    mHandler.post(runOnce);
+                    runnable.run();
                 }
             }, /* message= */ null);
         }
@@ -3853,43 +3890,5 @@
         void onSystemUserVisibilityChanged(boolean visible) {
             getUserManagerInternal().onSystemUserVisibilityChanged(visible);
         }
-
-        void lockDeviceNowAndWaitForKeyguardShown() {
-            if (getWindowManager().isKeyguardLocked()) {
-                return;
-            }
-
-            final TimingsTraceAndSlog t = new TimingsTraceAndSlog();
-            t.traceBegin("lockDeviceNowAndWaitForKeyguardShown");
-
-            final CountDownLatch latch = new CountDownLatch(1);
-            ActivityTaskManagerInternal.ScreenObserver screenObserver =
-                    new ActivityTaskManagerInternal.ScreenObserver() {
-                        @Override
-                        public void onAwakeStateChanged(boolean isAwake) {
-
-                        }
-
-                        @Override
-                        public void onKeyguardStateChanged(boolean isShowing) {
-                            if (isShowing) {
-                                latch.countDown();
-                            }
-                        }
-                    };
-
-            getActivityTaskManagerInternal().registerScreenObserver(screenObserver);
-            getWindowManager().lockDeviceNow();
-            try {
-                if (!latch.await(20, TimeUnit.SECONDS)) {
-                    throw new RuntimeException("Keyguard is not shown in 20 seconds");
-                }
-            } catch (InterruptedException e) {
-                throw new RuntimeException(e);
-            } finally {
-                getActivityTaskManagerInternal().unregisterScreenObserver(screenObserver);
-                t.traceEnd();
-            }
-        }
     }
 }
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 cea10ea..851c451 100644
--- a/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
@@ -57,6 +57,7 @@
 import static org.mockito.Matchers.anyBoolean;
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.doNothing;
@@ -89,6 +90,7 @@
 import android.os.Message;
 import android.os.PowerManagerInternal;
 import android.os.RemoteException;
+import android.os.SystemClock;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.os.storage.IStorageManager;
@@ -198,7 +200,10 @@
             doNothing().when(mInjector).activityManagerOnUserStopped(anyInt());
             doNothing().when(mInjector).clearBroadcastQueueForUser(anyInt());
             doNothing().when(mInjector).taskSupervisorRemoveUser(anyInt());
-            doNothing().when(mInjector).lockDeviceNowAndWaitForKeyguardShown();
+            doAnswer(invocation -> {
+                ((Runnable) invocation.getArgument(0)).run();
+                return null;
+            }).when(mInjector).showKeyguard(any());
             mockIsUsersOnSecondaryDisplaysEnabled(false);
             // All UserController params are set to default.
 
@@ -530,7 +535,6 @@
         expectedCodes.add(REPORT_USER_SWITCH_COMPLETE_MSG);
         if (backgroundUserStopping) {
             expectedCodes.add(CLEAR_USER_JOURNEY_SESSION_MSG);
-            expectedCodes.add(0); // this is for directly posting in stopping.
         }
         Set<Integer> actualCodes = mInjector.mHandler.getMessageCodes();
         assertEquals("Unexpected message sent", expectedCodes, actualCodes);
@@ -1047,21 +1051,13 @@
         // mock the device to be secure in order to expect the keyguard to be shown
         when(mInjector.mKeyguardManagerMock.isDeviceSecure(anyInt())).thenReturn(true);
 
-        // call real lockDeviceNowAndWaitForKeyguardShown method for this test
-        doCallRealMethod().when(mInjector).lockDeviceNowAndWaitForKeyguardShown();
+        // call real showKeyguard method for this test
+        doCallRealMethod().when(mInjector).showKeyguard(any());
 
-        // call startUser on a thread because we're expecting it to be blocked
-        Thread threadStartUser = new Thread(()-> {
-            mUserController.startUser(TEST_USER_ID, USER_START_MODE_FOREGROUND);
-        });
-        threadStartUser.start();
+        mUserController.completeUserSwitch(TEST_USER_ID1, TEST_USER_ID2);
 
-        // make sure the switch is stalled...
-        Thread.sleep(2000);
-        // by checking REPORT_USER_SWITCH_MSG is not sent yet
-        assertNull(mInjector.mHandler.getMessageForCode(REPORT_USER_SWITCH_MSG));
-        // and the thread is still alive
-        assertTrue(threadStartUser.isAlive());
+        // make sure the switch is stalled by checking the UserSwitchingDialog is not dismissed yet
+        verify(mInjector, never()).dismissUserSwitchingDialog(any());
 
         // mock send the keyguard shown event
         ArgumentCaptor<ActivityTaskManagerInternal.ScreenObserver> captor = ArgumentCaptor.forClass(
@@ -1069,12 +1065,41 @@
         verify(mInjector.mActivityTaskManagerInternal).registerScreenObserver(captor.capture());
         captor.getValue().onKeyguardStateChanged(true);
 
-        // verify the switch now moves on...
-        Thread.sleep(1000);
-        // by checking REPORT_USER_SWITCH_MSG is sent
-        assertNotNull(mInjector.mHandler.getMessageForCode(REPORT_USER_SWITCH_MSG));
-        // and the thread is finished
-        assertFalse(threadStartUser.isAlive());
+        // verify the switch now moves on by checking the UserSwitchingDialog is dismissed
+        verify(mInjector, atLeastOnce()).dismissUserSwitchingDialog(any());
+
+        // verify that SHOW_KEYGUARD_TIMEOUT is ignored and does not crash the system
+        try {
+            mInjector.mHandler.processPostDelayedCallbacksWithin(
+                    UserController.SHOW_KEYGUARD_TIMEOUT_MS);
+        } catch (RuntimeException e) {
+            throw new AssertionError(
+                    "SHOW_KEYGUARD_TIMEOUT is not ignored and crashed the system", e);
+        }
+    }
+
+    @Test
+    public void testRuntimeExceptionIsThrownIfTheKeyguardIsNotShown() throws Exception {
+        // enable user switch ui, because keyguard is only shown then
+        mUserController.setInitialConfig(/* userSwitchUiEnabled= */ true,
+                /* maxRunningUsers= */ 3, /* delayUserDataLocking= */ false);
+
+        // mock the device to be secure in order to expect the keyguard to be shown
+        when(mInjector.mKeyguardManagerMock.isDeviceSecure(anyInt())).thenReturn(true);
+
+        // suppress showKeyguard method for this test
+        doNothing().when(mInjector).showKeyguard(any());
+
+        mUserController.completeUserSwitch(TEST_USER_ID1, TEST_USER_ID2);
+
+        // verify that the system has crashed
+        assertThrows("Should have thrown RuntimeException", RuntimeException.class, () -> {
+            mInjector.mHandler.processPostDelayedCallbacksWithin(
+                    UserController.SHOW_KEYGUARD_TIMEOUT_MS);
+        });
+
+        // make sure the UserSwitchingDialog is not dismissed
+        verify(mInjector, never()).dismissUserSwitchingDialog(any());
     }
 
     private void setUpAndStartUserInBackground(int userId) throws Exception {
@@ -1410,7 +1435,9 @@
         Set<Integer> getMessageCodes() {
             Set<Integer> result = new LinkedHashSet<>();
             for (Message msg : mMessages) {
-                result.add(msg.what);
+                if (msg.what != 0) { // ignore mHandle.post and mHandler.postDelayed messages
+                    result.add(msg.what);
+                }
             }
             return result;
         }
@@ -1434,14 +1461,28 @@
 
         @Override
         public boolean sendMessageAtTime(Message msg, long uptimeMillis) {
+            final Runnable cb = msg.getCallback();
+            if (cb != null && uptimeMillis <= SystemClock.uptimeMillis()) {
+                // run mHandler.post calls immediately
+                cb.run();
+                return true;
+            }
             Message copy = new Message();
             copy.copyFrom(msg);
+            copy.setCallback(cb);
             mMessages.add(copy);
-            if (msg.getCallback() != null) {
-                msg.getCallback().run();
-                msg.setCallback(null);
-            }
             return super.sendMessageAtTime(msg, uptimeMillis);
         }
+
+        public void processPostDelayedCallbacksWithin(long millis) {
+            final long whenMax = SystemClock.uptimeMillis() + millis;
+            for (Message msg : mMessages) {
+                final Runnable cb = msg.getCallback();
+                if (cb != null && msg.getWhen() <= whenMax) {
+                    msg.setCallback(null);
+                    cb.run();
+                }
+            }
+        }
     }
 }