Make UserSwitchObserver.onBeforeUserSwitching oneway but still blocking.

Bug: 371536480
Test: atest UserControllerTest
Test: atest UserTrackerImplTest
Flag: EXEMPT bugfix

Change-Id: I9b71ff4d4ee726c40e37e67ee2d2bba1f2262555

7 files changed, 207 insertions, 112 deletions

diff --git a/core/java/android/app/IUserSwitchObserver.aidl b/core/java/android/app/IUserSwitchObserver.aidl
index 1ff7a17e578f..d71ee7c712e7 100644
--- a/core/java/android/app/IUserSwitchObserver.aidl
+++ b/core/java/android/app/IUserSwitchObserver.aidl
@@ -19,10 +19,10 @@ package android.app;
 import android.os.IRemoteCallback;
 
 /** {@hide} */
-interface IUserSwitchObserver {
-    void onBeforeUserSwitching(int newUserId);
-    oneway void onUserSwitching(int newUserId, IRemoteCallback reply);
-    oneway void onUserSwitchComplete(int newUserId);
-    oneway void onForegroundProfileSwitch(int newProfileId);
-    oneway void onLockedBootComplete(int newUserId);
+oneway interface IUserSwitchObserver {
+    void onBeforeUserSwitching(int newUserId, IRemoteCallback reply);
+    void onUserSwitching(int newUserId, IRemoteCallback reply);
+    void onUserSwitchComplete(int newUserId);
+    void onForegroundProfileSwitch(int newProfileId);
+    void onLockedBootComplete(int newUserId);
 }
diff --git a/core/java/android/app/UserSwitchObserver.java b/core/java/android/app/UserSwitchObserver.java
index 727799a1f948..1664cfb6f7a8 100644
--- a/core/java/android/app/UserSwitchObserver.java
+++ b/core/java/android/app/UserSwitchObserver.java
@@ -30,7 +30,11 @@ public class UserSwitchObserver extends IUserSwitchObserver.Stub {
     }
 
     @Override
-    public void onBeforeUserSwitching(int newUserId) throws RemoteException {}
+    public void onBeforeUserSwitching(int newUserId, IRemoteCallback reply) throws RemoteException {
+        if (reply != null) {
+            reply.sendResult(null);
+        }
+    }
 
     @Override
     public void onUserSwitching(int newUserId, IRemoteCallback reply) throws RemoteException {
diff --git a/packages/SystemUI/src/com/android/systemui/settings/UserTracker.kt b/packages/SystemUI/src/com/android/systemui/settings/UserTracker.kt
index e1631ccdcb06..bbb13d5c1dfe 100644
--- a/packages/SystemUI/src/com/android/systemui/settings/UserTracker.kt
+++ b/packages/SystemUI/src/com/android/systemui/settings/UserTracker.kt
@@ -61,9 +61,18 @@ interface UserTracker : UserContentResolverProvider, UserContextProvider {
     /** Callback for notifying of changes. */
     @WeaklyReferencedCallback
     interface Callback {
-        /** Notifies that the current user will be changed. */
+        /**
+         * Same as {@link onBeforeUserSwitching(Int, Runnable)} but the callback will be called
+         * automatically after the completion of this method.
+         */
         fun onBeforeUserSwitching(newUser: Int) {}
 
+        /** Notifies that the current user will be changed. */
+        fun onBeforeUserSwitching(newUser: Int, resultCallback: Runnable) {
+            onBeforeUserSwitching(newUser)
+            resultCallback.run()
+        }
+
         /**
          * Same as {@link onUserChanging(Int, Context, Runnable)} but the callback will be called
          * automatically after the completion of this method.
diff --git a/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt b/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt
index b7a3aedc565e..42d83637ec1a 100644
--- a/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/settings/UserTrackerImpl.kt
@@ -196,8 +196,9 @@ internal constructor(
     private fun registerUserSwitchObserver() {
         iActivityManager.registerUserSwitchObserver(
             object : UserSwitchObserver() {
-                override fun onBeforeUserSwitching(newUserId: Int) {
+                override fun onBeforeUserSwitching(newUserId: Int, reply: IRemoteCallback?) {
                     handleBeforeUserSwitching(newUserId)
+                    reply?.sendResult(null)
                 }
 
                 override fun onUserSwitching(newUserId: Int, reply: IRemoteCallback?) {
@@ -236,8 +237,7 @@ internal constructor(
         setUserIdInternal(newUserId)
 
         notifySubscribers { callback, resultCallback ->
-                callback.onBeforeUserSwitching(newUserId)
-                resultCallback.run()
+                callback.onBeforeUserSwitching(newUserId, resultCallback)
             }
             .await()
     }
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 a0ecb802dd61..f695c13a9e62 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt
@@ -76,6 +76,8 @@ class UserTrackerImplTest : SysuiTestCase() {
 
     @Mock private lateinit var iActivityManager: IActivityManager
 
+    @Mock private lateinit var beforeUserSwitchingReply: IRemoteCallback
+
     @Mock private lateinit var userSwitchingReply: IRemoteCallback
 
     @Mock(stubOnly = true) private lateinit var dumpManager: DumpManager
@@ -199,9 +201,10 @@ class UserTrackerImplTest : SysuiTestCase() {
 
             val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
             verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
-            captor.value.onBeforeUserSwitching(newID)
+            captor.value.onBeforeUserSwitching(newID, beforeUserSwitchingReply)
             captor.value.onUserSwitching(newID, userSwitchingReply)
             runCurrent()
+            verify(beforeUserSwitchingReply).sendResult(any())
             verify(userSwitchingReply).sendResult(any())
 
             verify(userManager).getProfiles(newID)
@@ -341,10 +344,11 @@ class UserTrackerImplTest : SysuiTestCase() {
 
             val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
             verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
-            captor.value.onBeforeUserSwitching(newID)
+            captor.value.onBeforeUserSwitching(newID, beforeUserSwitchingReply)
             captor.value.onUserSwitching(newID, userSwitchingReply)
             runCurrent()
 
+            verify(beforeUserSwitchingReply).sendResult(any())
             verify(userSwitchingReply).sendResult(any())
             assertThat(callback.calledOnUserChanging).isEqualTo(1)
             assertThat(callback.lastUser).isEqualTo(newID)
@@ -395,7 +399,7 @@ class UserTrackerImplTest : SysuiTestCase() {
 
             val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
             verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
-            captor.value.onBeforeUserSwitching(newID)
+            captor.value.onBeforeUserSwitching(newID, any())
             captor.value.onUserSwitchComplete(newID)
             runCurrent()
 
@@ -453,8 +457,10 @@ class UserTrackerImplTest : SysuiTestCase() {
 
             val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
             verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
+            captor.value.onBeforeUserSwitching(newID, beforeUserSwitchingReply)
             captor.value.onUserSwitching(newID, userSwitchingReply)
             runCurrent()
+            verify(beforeUserSwitchingReply).sendResult(any())
             verify(userSwitchingReply).sendResult(any())
             captor.value.onUserSwitchComplete(newID)
 
@@ -488,6 +494,7 @@ class UserTrackerImplTest : SysuiTestCase() {
         }
 
     private class TestCallback : UserTracker.Callback {
+        var calledOnBeforeUserChanging = 0
         var calledOnUserChanging = 0
         var calledOnUserChanged = 0
         var calledOnProfilesChanged = 0
@@ -495,6 +502,11 @@ class UserTrackerImplTest : SysuiTestCase() {
         var lastUserContext: Context? = null
         var lastUserProfiles = emptyList<UserInfo>()
 
+        override fun onBeforeUserSwitching(newUser: Int) {
+            calledOnBeforeUserChanging++
+            lastUser = newUser
+        }
+
         override fun onUserChanging(newUser: Int, userContext: Context) {
             calledOnUserChanging++
             lastUser = newUser
diff --git a/services/core/java/com/android/server/am/UserController.java b/services/core/java/com/android/server/am/UserController.java
index c31b9ef60bd2..70f2a8e1dd1b 100644
--- a/services/core/java/com/android/server/am/UserController.java
+++ b/services/core/java/com/android/server/am/UserController.java
@@ -160,6 +160,7 @@ 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.BiConsumer;
 import java.util.function.Consumer;
 
 /**
@@ -176,6 +177,9 @@ import java.util.function.Consumer;
 class UserController implements Handler.Callback {
     private static final String TAG = TAG_WITH_CLASS_NAME ? "UserController" : TAG_AM;
 
+    // Amount of time we wait for observers to handle onBeforeUserSwitching, before crashing system.
+    static final int DEFAULT_BEFORE_USER_SWITCH_TIMEOUT_MS = 20 * 1000;
+
     // Amount of time we wait for observers to handle a user switch before
     // giving up on them and dismissing the user switching dialog.
     static final int DEFAULT_USER_SWITCH_TIMEOUT_MS = 3 * 1000;
@@ -1920,8 +1924,14 @@ class UserController implements Handler.Callback {
                 return false;
             }
 
-            mHandler.post(() -> startUserInternalOnHandler(userId, oldUserId, userStartMode,
-                    unlockListener, callingUid, callingPid));
+            final Runnable continueStartUserInternal = () -> continueStartUserInternal(userInfo,
+                    oldUserId, userStartMode, unlockListener, callingUid, callingPid);
+            if (foreground) {
+                mHandler.post(() -> dispatchOnBeforeUserSwitching(userId, () ->
+                        mHandler.post(continueStartUserInternal)));
+            } else {
+                continueStartUserInternal.run();
+            }
         } finally {
             Binder.restoreCallingIdentity(ident);
         }
@@ -1929,11 +1939,11 @@ class UserController implements Handler.Callback {
         return true;
     }
 
-    private void startUserInternalOnHandler(int userId, int oldUserId, int userStartMode,
+    private void continueStartUserInternal(UserInfo userInfo, int oldUserId, int userStartMode,
             IProgressListener unlockListener, int callingUid, int callingPid) {
         final TimingsTraceAndSlog t = new TimingsTraceAndSlog();
         final boolean foreground = userStartMode == USER_START_MODE_FOREGROUND;
-        final UserInfo userInfo = getUserInfo(userId);
+        final int userId = userInfo.id;
 
         boolean needStart = false;
         boolean updateUmState = false;
@@ -1995,7 +2005,6 @@ class UserController implements Handler.Callback {
             // it should be moved outside, but for now it's not as there are many calls to
             // external components here afterwards
             updateProfileRelatedCaches();
-            dispatchOnBeforeUserSwitching(userId);
             mInjector.getWindowManager().setCurrentUser(userId);
             mInjector.reportCurWakefulnessUsageEvent();
             // Once the internal notion of the active user has switched, we lock the device
@@ -2296,25 +2305,40 @@ class UserController implements Handler.Callback {
         mUserSwitchObservers.finishBroadcast();
     }
 
-    private void dispatchOnBeforeUserSwitching(@UserIdInt int newUserId) {
+    private void dispatchOnBeforeUserSwitching(@UserIdInt int newUserId, Runnable onComplete) {
         final TimingsTraceAndSlog t = new TimingsTraceAndSlog();
         t.traceBegin("dispatchOnBeforeUserSwitching-" + newUserId);
-        final int observerCount = mUserSwitchObservers.beginBroadcast();
-        for (int i = 0; i < observerCount; i++) {
-            final String name = "#" + i + " " + mUserSwitchObservers.getBroadcastCookie(i);
-            t.traceBegin("onBeforeUserSwitching-" + name);
+        final AtomicBoolean isSuccessful = new AtomicBoolean(false);
+        startTimeoutForOnBeforeUserSwitching(isSuccessful);
+        informUserSwitchObservers((observer, callback) -> {
             try {
-                mUserSwitchObservers.getBroadcastItem(i).onBeforeUserSwitching(newUserId);
+                observer.onBeforeUserSwitching(newUserId, callback);
             } catch (RemoteException e) {
-                // Ignore
-            } finally {
-                t.traceEnd();
+                // ignore
             }
-        }
-        mUserSwitchObservers.finishBroadcast();
+        }, () -> {
+            isSuccessful.set(true);
+            onComplete.run();
+        }, "onBeforeUserSwitching");
         t.traceEnd();
     }
 
+    private void startTimeoutForOnBeforeUserSwitching(AtomicBoolean isSuccessful) {
+        mHandler.postDelayed(() -> {
+            if (isSuccessful.get()) {
+                return;
+            }
+            String unresponsiveObservers;
+            synchronized (mLock) {
+                unresponsiveObservers = String.join(", ", mCurWaitingUserSwitchCallbacks);
+            }
+            throw new RuntimeException("Timeout on dispatchOnBeforeUserSwitching. "
+                    + "These UserSwitchObservers did not respond in "
+                    + DEFAULT_BEFORE_USER_SWITCH_TIMEOUT_MS + "ms: " + unresponsiveObservers + ".");
+        }, DEFAULT_BEFORE_USER_SWITCH_TIMEOUT_MS);
+    }
+
+
     /** Called on handler thread */
     @VisibleForTesting
     void dispatchUserSwitchComplete(@UserIdInt int oldUserId, @UserIdInt int newUserId) {
@@ -2527,70 +2551,76 @@ class UserController implements Handler.Callback {
         t.traceBegin("dispatchUserSwitch-" + oldUserId + "-to-" + newUserId);
 
         EventLog.writeEvent(EventLogTags.UC_DISPATCH_USER_SWITCH, oldUserId, newUserId);
+        uss.switching = true;
+        informUserSwitchObservers((observer, callback) -> {
+            try {
+                observer.onUserSwitching(newUserId, callback);
+            } catch (RemoteException e) {
+                // ignore
+            }
+        }, () -> {
+            synchronized (mLock) {
+                sendContinueUserSwitchLU(uss, oldUserId, newUserId);
+            }
+        }, "onUserSwitching");
+        t.traceEnd();
+    }
 
+    void informUserSwitchObservers(BiConsumer<IUserSwitchObserver, IRemoteCallback> consumer,
+            final Runnable onComplete, String trace) {
         final int observerCount = mUserSwitchObservers.beginBroadcast();
-        if (observerCount > 0) {
-            final ArraySet<String> curWaitingUserSwitchCallbacks = new ArraySet<>();
+        if (observerCount == 0) {
+            onComplete.run();
+            mUserSwitchObservers.finishBroadcast();
+            return;
+        }
+        final ArraySet<String> curWaitingUserSwitchCallbacks = new ArraySet<>();
+        synchronized (mLock) {
+            mCurWaitingUserSwitchCallbacks = curWaitingUserSwitchCallbacks;
+        }
+        final AtomicInteger waitingCallbacksCount = new AtomicInteger(observerCount);
+        final long userSwitchTimeoutMs = getUserSwitchTimeoutMs();
+        final long dispatchStartedTime = SystemClock.elapsedRealtime();
+        for (int i = 0; i < observerCount; i++) {
+            final long dispatchStartedTimeForObserver = SystemClock.elapsedRealtime();
+            // Prepend with unique prefix to guarantee that keys are unique
+            final String name = "#" + i + " " + mUserSwitchObservers.getBroadcastCookie(i);
             synchronized (mLock) {
-                uss.switching = true;
-                mCurWaitingUserSwitchCallbacks = curWaitingUserSwitchCallbacks;
-            }
-            final AtomicInteger waitingCallbacksCount = new AtomicInteger(observerCount);
-            final long userSwitchTimeoutMs = getUserSwitchTimeoutMs();
-            final long dispatchStartedTime = SystemClock.elapsedRealtime();
-            for (int i = 0; i < observerCount; i++) {
-                final long dispatchStartedTimeForObserver = SystemClock.elapsedRealtime();
-                try {
-                    // Prepend with unique prefix to guarantee that keys are unique
-                    final String name = "#" + i + " " + mUserSwitchObservers.getBroadcastCookie(i);
+                curWaitingUserSwitchCallbacks.add(name);
+            }
+            final IRemoteCallback callback = new IRemoteCallback.Stub() {
+                @Override
+                public void sendResult(Bundle data) throws RemoteException {
+                    asyncTraceEnd(trace + "-" + name, 0);
                     synchronized (mLock) {
-                        curWaitingUserSwitchCallbacks.add(name);
-                    }
-                    final IRemoteCallback callback = new IRemoteCallback.Stub() {
-                        @Override
-                        public void sendResult(Bundle data) throws RemoteException {
-                            asyncTraceEnd("onUserSwitching-" + name, newUserId);
-                            synchronized (mLock) {
-                                long delayForObserver = SystemClock.elapsedRealtime()
-                                        - dispatchStartedTimeForObserver;
-                                if (delayForObserver > LONG_USER_SWITCH_OBSERVER_WARNING_TIME_MS) {
-                                    Slogf.w(TAG, "User switch slowed down by observer " + name
-                                            + ": result took " + delayForObserver
-                                            + " ms to process.");
-                                }
-
-                                long totalDelay = SystemClock.elapsedRealtime()
-                                        - dispatchStartedTime;
-                                if (totalDelay > userSwitchTimeoutMs) {
-                                    Slogf.e(TAG, "User switch timeout: observer " + name
-                                            + "'s result was received " + totalDelay
-                                            + " ms after dispatchUserSwitch.");
-                                }
-
-                                curWaitingUserSwitchCallbacks.remove(name);
-                                // Continue switching if all callbacks have been notified and
-                                // user switching session is still valid
-                                if (waitingCallbacksCount.decrementAndGet() == 0
-                                        && (curWaitingUserSwitchCallbacks
-                                        == mCurWaitingUserSwitchCallbacks)) {
-                                    sendContinueUserSwitchLU(uss, oldUserId, newUserId);
-                                }
-                            }
+                        long delayForObserver = SystemClock.elapsedRealtime()
+                                - dispatchStartedTimeForObserver;
+                        if (delayForObserver > LONG_USER_SWITCH_OBSERVER_WARNING_TIME_MS) {
+                            Slogf.w(TAG, "User switch slowed down by observer " + name
+                                    + ": result took " + delayForObserver
+                                    + " ms to process. " + trace);
                         }
-                    };
-                    asyncTraceBegin("onUserSwitching-" + name, newUserId);
-                    mUserSwitchObservers.getBroadcastItem(i).onUserSwitching(newUserId, callback);
-                } catch (RemoteException e) {
-                    // Ignore
+                        long totalDelay = SystemClock.elapsedRealtime() - dispatchStartedTime;
+                        if (totalDelay > userSwitchTimeoutMs) {
+                            Slogf.e(TAG, "User switch timeout: observer " + name
+                                    + "'s result was received " + totalDelay
+                                    + " ms after dispatchUserSwitch. " + trace);
+                        }
+                        curWaitingUserSwitchCallbacks.remove(name);
+                        // Continue switching if all callbacks have been notified and
+                        // user switching session is still valid
+                        if (waitingCallbacksCount.decrementAndGet() == 0
+                                && (curWaitingUserSwitchCallbacks
+                                == mCurWaitingUserSwitchCallbacks)) {
+                            onComplete.run();
+                        }
+                    }
                 }
-            }
-        } else {
-            synchronized (mLock) {
-                sendContinueUserSwitchLU(uss, oldUserId, newUserId);
-            }
+            };
+            asyncTraceBegin(trace + "-" + name, 0);
+            consumer.accept(mUserSwitchObservers.getBroadcastItem(i), callback);
         }
         mUserSwitchObservers.finishBroadcast();
-        t.traceEnd(); // end dispatchUserSwitch-
     }
 
     @GuardedBy("mLock")
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 2fe6918630f6..7dbbff22a537 100644
--- a/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
@@ -33,6 +33,7 @@ import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentat
 import static com.android.server.am.UserController.CLEAR_USER_JOURNEY_SESSION_MSG;
 import static com.android.server.am.UserController.COMPLETE_USER_SWITCH_MSG;
 import static com.android.server.am.UserController.CONTINUE_USER_SWITCH_MSG;
+import static com.android.server.am.UserController.DEFAULT_BEFORE_USER_SWITCH_TIMEOUT_MS;
 import static com.android.server.am.UserController.REPORT_LOCKED_BOOT_COMPLETE_MSG;
 import static com.android.server.am.UserController.REPORT_USER_SWITCH_COMPLETE_MSG;
 import static com.android.server.am.UserController.REPORT_USER_SWITCH_MSG;
@@ -94,6 +95,7 @@ import android.os.Looper;
 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;
@@ -181,14 +183,12 @@ public class UserControllerTest {
             Intent.ACTION_USER_STARTING);
 
     private static final Set<Integer> START_FOREGROUND_USER_MESSAGE_CODES = newHashSet(
-            0, // for startUserInternalOnHandler
             REPORT_USER_SWITCH_MSG,
             USER_SWITCH_TIMEOUT_MSG,
             USER_START_MSG,
             USER_CURRENT_MSG);
 
     private static final Set<Integer> START_BACKGROUND_USER_MESSAGE_CODES = newHashSet(
-            0, // for startUserInternalOnHandler
             USER_START_MSG,
             REPORT_LOCKED_BOOT_COMPLETE_MSG);
 
@@ -376,7 +376,7 @@ public class UserControllerTest {
         // and the cascade effect goes on...). In fact, a better approach would to not assert the
         // binder calls, but their side effects (in this case, that the user is stopped right away)
         assertWithMessage("wrong binder message calls").that(mInjector.mHandler.getMessageCodes())
-                .containsExactly(/* for startUserInternalOnHandler */ 0, USER_START_MSG);
+                .containsExactly(USER_START_MSG);
     }
 
     private void startUserAssertions(
@@ -419,17 +419,12 @@ public class UserControllerTest {
     @Test
     public void testDispatchUserSwitch() throws RemoteException {
         // Prepare mock observer and register it
-        IUserSwitchObserver observer = mock(IUserSwitchObserver.class);
-        when(observer.asBinder()).thenReturn(new Binder());
-        doAnswer(invocation -> {
-            IRemoteCallback callback = (IRemoteCallback) invocation.getArguments()[1];
-            callback.sendResult(null);
-            return null;
-        }).when(observer).onUserSwitching(anyInt(), any());
-        mUserController.registerUserSwitchObserver(observer, "mock");
+        IUserSwitchObserver observer = registerUserSwitchObserver(
+                /* replyToOnBeforeUserSwitchingCallback= */ true,
+                /* replyToOnUserSwitchingCallback= */ true);
         // Start user -- this will update state of mUserController
         mUserController.startUser(TEST_USER_ID, USER_START_MODE_FOREGROUND);
-        verify(observer, times(1)).onBeforeUserSwitching(eq(TEST_USER_ID));
+        verify(observer, times(1)).onBeforeUserSwitching(eq(TEST_USER_ID), any());
         Message reportMsg = mInjector.mHandler.getMessageForCode(REPORT_USER_SWITCH_MSG);
         assertNotNull(reportMsg);
         UserState userState = (UserState) reportMsg.obj;
@@ -453,14 +448,26 @@ public class UserControllerTest {
     }
 
     @Test
+    public void testShouldCrashWhenOnBeforeUserSwitchingTimeouts() throws RemoteException {
+        IUserSwitchObserver observer = registerUserSwitchObserver(
+                /* replyToOnBeforeUserSwitchingCallback= */ false,
+                /* replyToOnUserSwitchingCallback= */ true);
+        mUserController.startUser(TEST_USER_ID, USER_START_MODE_FOREGROUND);
+        verify(observer, times(1)).onBeforeUserSwitching(eq(TEST_USER_ID), any());
+        assertThrows("Should have crashed when observers don't reply to onBeforeUserSwitching in "
+                        + DEFAULT_BEFORE_USER_SWITCH_TIMEOUT_MS + " ms", RuntimeException.class,
+                mInjector.mHandler::runPendingCallbacks);
+    }
+
+    @Test
     public void testDispatchUserSwitchBadReceiver() throws RemoteException {
-        // Prepare mock observer which doesn't notify the callback and register it
-        IUserSwitchObserver observer = mock(IUserSwitchObserver.class);
-        when(observer.asBinder()).thenReturn(new Binder());
-        mUserController.registerUserSwitchObserver(observer, "mock");
+        // Prepare mock observer which doesn't notify the onUserSwitching callback and register it
+        IUserSwitchObserver observer = registerUserSwitchObserver(
+                /* replyToOnBeforeUserSwitchingCallback= */ true,
+                /* replyToOnUserSwitchingCallback= */ false);
         // Start user -- this will update state of mUserController
         mUserController.startUser(TEST_USER_ID, USER_START_MODE_FOREGROUND);
-        verify(observer, times(1)).onBeforeUserSwitching(eq(TEST_USER_ID));
+        verify(observer, times(1)).onBeforeUserSwitching(eq(TEST_USER_ID), any());
         Message reportMsg = mInjector.mHandler.getMessageForCode(REPORT_USER_SWITCH_MSG);
         assertNotNull(reportMsg);
         UserState userState = (UserState) reportMsg.obj;
@@ -551,7 +558,6 @@ public class UserControllerTest {
         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.
         }
         if (expectScheduleBackgroundUserStopping) {
             expectedCodes.add(SCHEDULED_STOP_BACKGROUND_USER_MSG);
@@ -567,9 +573,9 @@ public class UserControllerTest {
     @Test
     public void testDispatchUserSwitchComplete() throws RemoteException {
         // Prepare mock observer and register it
-        IUserSwitchObserver observer = mock(IUserSwitchObserver.class);
-        when(observer.asBinder()).thenReturn(new Binder());
-        mUserController.registerUserSwitchObserver(observer, "mock");
+        IUserSwitchObserver observer = registerUserSwitchObserver(
+                /* replyToOnBeforeUserSwitchingCallback= */ true,
+                /* replyToOnUserSwitchingCallback= */ true);
         // Start user -- this will update state of mUserController
         mUserController.startUser(TEST_USER_ID, USER_START_MODE_FOREGROUND);
         Message reportMsg = mInjector.mHandler.getMessageForCode(REPORT_USER_SWITCH_MSG);
@@ -1752,6 +1758,29 @@ public class UserControllerTest {
         verify(mInjector, never()).onSystemUserVisibilityChanged(anyBoolean());
     }
 
+    private IUserSwitchObserver registerUserSwitchObserver(
+            boolean replyToOnBeforeUserSwitchingCallback, boolean replyToOnUserSwitchingCallback)
+            throws RemoteException {
+        IUserSwitchObserver observer = mock(IUserSwitchObserver.class);
+        when(observer.asBinder()).thenReturn(new Binder());
+        if (replyToOnBeforeUserSwitchingCallback) {
+            doAnswer(invocation -> {
+                IRemoteCallback callback = (IRemoteCallback) invocation.getArguments()[1];
+                callback.sendResult(null);
+                return null;
+            }).when(observer).onBeforeUserSwitching(anyInt(), any());
+        }
+        if (replyToOnUserSwitchingCallback) {
+            doAnswer(invocation -> {
+                IRemoteCallback callback = (IRemoteCallback) invocation.getArguments()[1];
+                callback.sendResult(null);
+                return null;
+            }).when(observer).onUserSwitching(anyInt(), any());
+        }
+        mUserController.registerUserSwitchObserver(observer, "mock");
+        return observer;
+    }
+
     // Should be public to allow mocking
     private static class TestInjector extends UserController.Injector {
         public final TestHandler mHandler;
@@ -1957,6 +1986,7 @@ public class UserControllerTest {
          * fix this, but in the meantime, this is your warning.
          */
         private final List<Message> mMessages = new ArrayList<>();
+        private final List<Runnable> mPendingCallbacks = new ArrayList<>();
 
         TestHandler(Looper looper) {
             super(looper);
@@ -1989,14 +2019,24 @@ public class UserControllerTest {
 
         @Override
         public boolean sendMessageAtTime(Message msg, long uptimeMillis) {
-            Message copy = new Message();
-            copy.copyFrom(msg);
-            mMessages.add(copy);
-            if (msg.getCallback() != null) {
-                msg.getCallback().run();
+            if (msg.getCallback() == null) {
+                Message copy = new Message();
+                copy.copyFrom(msg);
+                mMessages.add(copy);
+            } else {
+                if (SystemClock.uptimeMillis() >= uptimeMillis) {
+                    msg.getCallback().run();
+                } else {
+                    mPendingCallbacks.add(msg.getCallback());
+                }
                 msg.setCallback(null);
             }
             return super.sendMessageAtTime(msg, uptimeMillis);
         }
+
+        private void runPendingCallbacks() {
+            mPendingCallbacks.forEach(Runnable::run);
+            mPendingCallbacks.clear();
+        }
     }
 }