Make InputMethodBindingController per user

Make InputMethodBindingController per user by moving
IMMS#mBindingController to `UserDataRepository.UserData`.

This is the final roll-forward of [1], which was reverted by [2]
due to an additional locking contention in IMMS constructor [3].

[1] Original CL: I870a76ac1d196436f1b2f172d65ac46d580650d6
[2] Revert   CL: I31bf05f1d5af17b0220a4ea0df9278a1155beacf
[3] Sync IMMS Constructor CL: I8e51b1ced4dc16cdca7e898885c64793665fafef

This CL represents an internal refactoring and shouldn't introduce any
observable breakage.

Bug: 325515685
Test: atest FrameworksInputMethodSystemServerTests
Test: atest CtsInputMethodTestCases
Test: atest FrameworksServicesTests
Test: atest --host FrameworksInputMethodSystemServerTestsRavenwood

Change-Id: Id6e9eca6caaa8edde8c4a405448c31c0ca986509

6 files changed, 138 insertions, 64 deletions

diff --git a/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java b/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
index 3e23f972bd45..b7091744c3b6 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
@@ -21,6 +21,7 @@ import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.UserIdInt;
 import android.app.ActivityOptions;
 import android.app.PendingIntent;
 import android.content.ComponentName;
@@ -63,6 +64,7 @@ final class InputMethodBindingController {
     /** Time in milliseconds that the IME service has to bind before it is reconnected. */
     static final long TIME_TO_RECONNECT = 3 * 1000;
 
+    @UserIdInt final int mUserId;
     @NonNull private final InputMethodManagerService mService;
     @NonNull private final Context mContext;
     @NonNull private final PackageManagerInternal mPackageManagerInternal;
@@ -107,12 +109,15 @@ final class InputMethodBindingController {
                     | Context.BIND_INCLUDE_CAPABILITIES
                     | Context.BIND_SHOWING_UI;
 
-    InputMethodBindingController(@NonNull InputMethodManagerService service) {
-        this(service, IME_CONNECTION_BIND_FLAGS, null /* latchForTesting */);
+    InputMethodBindingController(@UserIdInt int userId,
+            @NonNull InputMethodManagerService service) {
+        this(userId, service, IME_CONNECTION_BIND_FLAGS, null /* latchForTesting */);
     }
 
-    InputMethodBindingController(@NonNull InputMethodManagerService service,
-            int imeConnectionBindFlags, CountDownLatch latchForTesting) {
+    InputMethodBindingController(@UserIdInt int userId,
+            @NonNull InputMethodManagerService service, int imeConnectionBindFlags,
+            CountDownLatch latchForTesting) {
+        mUserId = userId;
         mService = service;
         mContext = mService.mContext;
         mPackageManagerInternal = mService.mPackageManagerInternal;
@@ -301,7 +306,8 @@ final class InputMethodBindingController {
                     }
                     if (DEBUG) Slog.v(TAG, "Initiating attach with token: " + mCurToken);
                     final InputMethodInfo info =
-                            mService.queryInputMethodForCurrentUserLocked(mSelectedMethodId);
+                            InputMethodSettingsRepository.get(mUserId).getMethodMap().get(
+                                    mSelectedMethodId);
                     boolean supportsStylusHwChanged =
                             mSupportsStylusHw != info.supportsStylusHandwriting();
                     mSupportsStylusHw = info.supportsStylusHandwriting();
@@ -339,7 +345,7 @@ final class InputMethodBindingController {
         private void updateCurrentMethodUid() {
             final String curMethodPackage = mCurIntent.getComponent().getPackageName();
             final int curMethodUid = mPackageManagerInternal.getPackageUid(
-                    curMethodPackage, 0 /* flags */, mService.getCurrentImeUserIdLocked());
+                    curMethodPackage, 0 /* flags */, mUserId);
             if (curMethodUid < 0) {
                 Slog.e(TAG, "Failed to get UID for package=" + curMethodPackage);
                 mCurMethodUid = Process.INVALID_UID;
@@ -425,7 +431,8 @@ final class InputMethodBindingController {
             return InputBindResult.NO_IME;
         }
 
-        InputMethodInfo info = mService.queryInputMethodForCurrentUserLocked(mSelectedMethodId);
+        InputMethodInfo info = InputMethodSettingsRepository.get(mUserId).getMethodMap().get(
+                mSelectedMethodId);
         if (info == null) {
             throw new IllegalArgumentException("Unknown id: " + mSelectedMethodId);
         }
@@ -497,8 +504,7 @@ final class InputMethodBindingController {
             Slog.e(TAG, "--- bind failed: service = " + mCurIntent + ", conn = " + conn);
             return false;
         }
-        return mContext.bindServiceAsUser(mCurIntent, conn, flags,
-                new UserHandle(mService.getCurrentImeUserIdLocked()));
+        return mContext.bindServiceAsUser(mCurIntent, conn, flags, new UserHandle(mUserId));
     }
 
     @GuardedBy("ImfLock.class")
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index 8985022ea1b8..14d04119c55f 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -205,6 +205,7 @@ import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import java.util.function.IntConsumer;
+import java.util.function.IntFunction;
 
 /**
  * This class provides a system service that manages input methods.
@@ -306,8 +307,6 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
     @MultiUserUnawareField
     private final InputMethodMenuController mMenuController;
     @MultiUserUnawareField
-    @NonNull private final InputMethodBindingController mBindingController;
-    @MultiUserUnawareField
     @NonNull private final AutofillSuggestionsController mAutofillController;
 
     @GuardedBy("ImfLock.class")
@@ -478,7 +477,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
     @GuardedBy("ImfLock.class")
     @Nullable
     String getSelectedMethodIdLocked() {
-        return mBindingController.getSelectedMethodId();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        return userData.mBindingController.getSelectedMethodId();
     }
 
     /**
@@ -487,7 +487,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
      */
     @GuardedBy("ImfLock.class")
     private int getSequenceNumberLocked() {
-        return mBindingController.getSequenceNumber();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        return userData.mBindingController.getSequenceNumber();
     }
 
     /**
@@ -496,7 +497,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
      */
     @GuardedBy("ImfLock.class")
     private void advanceSequenceNumberLocked() {
-        mBindingController.advanceSequenceNumber();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        userData.mBindingController.advanceSequenceNumber();
     }
 
     @GuardedBy("ImfLock.class")
@@ -556,7 +558,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
     @GuardedBy("ImfLock.class")
     @Nullable
     private String getCurIdLocked() {
-        return mBindingController.getCurId();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        return userData.mBindingController.getCurId();
     }
 
     /**
@@ -580,7 +583,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
      */
     @GuardedBy("ImfLock.class")
     private boolean hasConnectionLocked() {
-        return mBindingController.hasMainConnection();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        return userData.mBindingController.hasMainConnection();
     }
 
     /**
@@ -603,7 +607,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
     @GuardedBy("ImfLock.class")
     @Nullable
     private Intent getCurIntentLocked() {
-        return mBindingController.getCurIntent();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        return userData.mBindingController.getCurIntent();
     }
 
     /**
@@ -613,7 +618,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
     @GuardedBy("ImfLock.class")
     @Nullable
     IBinder getCurTokenLocked() {
-        return mBindingController.getCurToken();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        return userData.mBindingController.getCurToken();
     }
 
     /**
@@ -654,7 +660,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
     @GuardedBy("ImfLock.class")
     @Nullable
     IInputMethodInvoker getCurMethodLocked() {
-        return mBindingController.getCurMethod();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        return userData.mBindingController.getCurMethod();
     }
 
     /**
@@ -662,7 +669,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
      */
     @GuardedBy("ImfLock.class")
     private int getCurMethodUidLocked() {
-        return mBindingController.getCurMethodUid();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        return userData.mBindingController.getCurMethodUid();
     }
 
     /**
@@ -671,7 +679,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
      */
     @GuardedBy("ImfLock.class")
     private long getLastBindTimeLocked() {
-        return mBindingController.getLastBindTime();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        return userData.mBindingController.getLastBindTime();
     }
 
     /**
@@ -1353,7 +1362,7 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
     InputMethodManagerService(
             Context context,
             @Nullable ServiceThread serviceThreadForTesting,
-            @Nullable InputMethodBindingController bindingControllerForTesting) {
+            @Nullable IntFunction<InputMethodBindingController> bindingControllerForTesting) {
         synchronized (ImfLock.class) {
             mContext = context;
             mRes = context.getResources();
@@ -1392,7 +1401,12 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
             AdditionalSubtypeMapRepository.initialize(mHandler, mContext);
 
             mCurrentUserId = mActivityManagerInternal.getCurrentUserId();
-            mUserDataRepository = new UserDataRepository(mHandler, mUserManagerInternal);
+            @SuppressWarnings("GuardedBy") final IntFunction<InputMethodBindingController>
+                    bindingControllerFactory = userId -> new InputMethodBindingController(userId,
+                            InputMethodManagerService.this);
+            mUserDataRepository = new UserDataRepository(mHandler, mUserManagerInternal,
+                    bindingControllerForTesting != null ? bindingControllerForTesting
+                            : bindingControllerFactory);
             for (int id : mUserManagerInternal.getUserIds()) {
                 mUserDataRepository.getOrCreate(id);
             }
@@ -1406,12 +1420,7 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
                     new HardwareKeyboardShortcutController(settings.getMethodMap(),
                             settings.getUserId());
             mMenuController = new InputMethodMenuController(this);
-            mBindingController =
-                    bindingControllerForTesting != null
-                            ? bindingControllerForTesting
-                            : new InputMethodBindingController(this);
             mAutofillController = new AutofillSuggestionsController(this);
-
             mVisibilityStateComputer = new ImeVisibilityStateComputer(this);
             mVisibilityApplier = new DefaultImeVisibilityApplier(this);
 
@@ -1544,9 +1553,9 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
 
         // Note that in b/197848765 we want to see if we can keep the binding alive for better
         // profile switching.
-        mBindingController.unbindCurrentMethod();
-        // TODO(b/325515685): No need to do this once BindingController becomes per-user.
-        mBindingController.setSelectedMethodId(null);
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        userData.mBindingController.unbindCurrentMethod();
+
         unbindCurrentClientLocked(UnbindReason.SWITCH_USER);
 
         // Hereafter we start initializing things for "newUserId".
@@ -1763,9 +1772,10 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
 
             // Check if selected IME of current user supports handwriting.
             if (userId == mCurrentUserId) {
-                return mBindingController.supportsStylusHandwriting()
+                final var userData = mUserDataRepository.getOrCreate(userId);
+                return userData.mBindingController.supportsStylusHandwriting()
                         && (!connectionless
-                                || mBindingController.supportsConnectionlessStylusHandwriting());
+                        || userData.mBindingController.supportsConnectionlessStylusHandwriting());
             }
             final InputMethodSettings settings = InputMethodSettingsRepository.get(userId);
             final InputMethodInfo imi = settings.getMethodMap().get(
@@ -2095,7 +2105,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
                 curInputMethodInfo != null && curInputMethodInfo.suppressesSpellChecker();
         final SparseArray<IAccessibilityInputMethodSession> accessibilityInputMethodSessions =
                 createAccessibilityInputMethodSessions(mCurClient.mAccessibilitySessions);
-        if (mBindingController.supportsStylusHandwriting() && hasSupportedStylusLocked()) {
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        if (userData.mBindingController.supportsStylusHandwriting() && hasSupportedStylusLocked()) {
             mHwController.setInkWindowInitializer(new InkWindowInitializer());
         }
         return new InputBindResult(InputBindResult.ResultCode.SUCCESS_WITH_IME_SESSION,
@@ -2216,6 +2227,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
         if (connectionIsActive != connectionWasActive) {
             mInputManagerInternal.notifyInputMethodConnectionActive(connectionIsActive);
         }
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+
 
         // If configured, we want to avoid starting up the IME if it is not supposed to be showing
         if (shouldPreventImeStartupLocked(selectedMethodId, startInputFlags,
@@ -2224,7 +2237,7 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
                 Slog.d(TAG, "Avoiding IME startup and unbinding current input method.");
             }
             invalidateAutofillSessionLocked();
-            mBindingController.unbindCurrentMethod();
+            userData.mBindingController.unbindCurrentMethod();
             return InputBindResult.NO_EDITOR;
         }
 
@@ -2256,9 +2269,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
             }
         }
 
-        mBindingController.unbindCurrentMethod();
-
-        return mBindingController.bindCurrentMethod();
+        userData.mBindingController.unbindCurrentMethod();
+        return userData.mBindingController.bindCurrentMethod();
     }
 
     /**
@@ -2518,11 +2530,13 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
 
     @GuardedBy("ImfLock.class")
     void resetCurrentMethodAndClientLocked(@UnbindReason int unbindClientReason) {
-        mBindingController.setSelectedMethodId(null);
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        userData.mBindingController.setSelectedMethodId(null);
+
         // Callback before clean-up binding states.
         // TODO(b/338461930): Check if this is still necessary or not.
         onUnbindCurrentMethodByReset();
-        mBindingController.unbindCurrentMethod();
+        userData.mBindingController.unbindCurrentMethod();
         unbindCurrentClientLocked(unbindClientReason);
     }
 
@@ -3099,7 +3113,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
             // mCurMethodId should be updated after setSelectedInputMethodAndSubtypeLocked()
             // because mCurMethodId is stored as a history in
             // setSelectedInputMethodAndSubtypeLocked().
-            mBindingController.setSelectedMethodId(id);
+            final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+            userData.mBindingController.setSelectedMethodId(id);
 
             if (mActivityManagerInternal.isSystemReady()) {
                 Intent intent = new Intent(Intent.ACTION_INPUT_METHOD_CHANGED);
@@ -3154,7 +3169,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
             @Nullable String delegatorPackageName,
             @NonNull IConnectionlessHandwritingCallback callback) {
         synchronized (ImfLock.class) {
-            if (!mBindingController.supportsConnectionlessStylusHandwriting()) {
+            final var userData = mUserDataRepository.getOrCreate(userId);
+            if (!userData.mBindingController.supportsConnectionlessStylusHandwriting()) {
                 Slog.w(TAG, "Connectionless stylus handwriting mode unsupported by IME.");
                 try {
                     callback.onError(CONNECTIONLESS_HANDWRITING_ERROR_UNSUPPORTED);
@@ -3237,7 +3253,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
                 }
                 final long ident = Binder.clearCallingIdentity();
                 try {
-                    if (!mBindingController.supportsStylusHandwriting()) {
+                    final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+                    if (!userData.mBindingController.supportsStylusHandwriting()) {
                         Slog.w(TAG,
                                 "Stylus HW unsupported by IME. Ignoring startStylusHandwriting()");
                         return false;
@@ -3420,7 +3437,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
         mVisibilityStateComputer.requestImeVisibility(windowToken, true);
 
         // Ensure binding the connection when IME is going to show.
-        mBindingController.setCurrentMethodVisible();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        userData.mBindingController.setCurrentMethodVisible();
         final IInputMethodInvoker curMethod = getCurMethodLocked();
         ImeTracker.forLogging().onCancelled(mCurStatsToken, ImeTracker.PHASE_SERVER_WAIT_IME);
         final boolean readyToDispatchToIme;
@@ -3528,7 +3546,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
         } else {
             ImeTracker.forLogging().onCancelled(statsToken, ImeTracker.PHASE_SERVER_SHOULD_HIDE);
         }
-        mBindingController.setCurrentMethodNotVisible();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        userData.mBindingController.setCurrentMethodNotVisible();
         mVisibilityStateComputer.clearImeShowFlags();
         // Cancel existing statsToken for show IME as we got a hide request.
         ImeTracker.forLogging().onCancelled(mCurStatsToken, ImeTracker.PHASE_SERVER_WAIT_IME);
@@ -3810,7 +3829,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
                 // Note that we can trust client's display ID as long as it matches
                 // to the display ID obtained from the window.
                 if (cs.mSelfReportedDisplayId != mCurTokenDisplayId) {
-                    mBindingController.unbindCurrentMethod();
+                    final var userData = mUserDataRepository.getOrCreate(userId);
+                    userData.mBindingController.unbindCurrentMethod();
                 }
             }
         }
@@ -4271,8 +4291,9 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
         mStylusIds.add(deviceId);
         // a new Stylus is detected. If IME supports handwriting, and we don't have
         // handwriting initialized, lets do it now.
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
         if (!mHwController.getCurrentRequestId().isPresent()
-                && mBindingController.supportsStylusHandwriting()) {
+                && userData.mBindingController.supportsStylusHandwriting()) {
             scheduleResetStylusHandwriting();
         }
     }
@@ -4841,7 +4862,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
 
             case MSG_RESET_HANDWRITING: {
                 synchronized (ImfLock.class) {
-                    if (mBindingController.supportsStylusHandwriting()
+                    final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+                    if (userData.mBindingController.supportsStylusHandwriting()
                             && getCurMethodLocked() != null && hasSupportedStylusLocked()) {
                         Slog.d(TAG, "Initializing Handwriting Spy");
                         mHwController.initializeHandwritingSpy(mCurTokenDisplayId);
@@ -4866,11 +4888,12 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
                     if (curMethod == null || mImeBindingState.mFocusedWindow == null) {
                         return true;
                     }
+                    final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
                     final HandwritingModeController.HandwritingSession session =
                             mHwController.startHandwritingSession(
                                     msg.arg1 /*requestId*/,
                                     msg.arg2 /*pid*/,
-                                    mBindingController.getCurMethodUid(),
+                                    userData.mBindingController.getCurMethodUid(),
                                     mImeBindingState.mFocusedWindow);
                     if (session == null) {
                         Slog.e(TAG,
@@ -5164,7 +5187,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
 
     @GuardedBy("ImfLock.class")
     void sendOnNavButtonFlagsChangedLocked() {
-        final IInputMethodInvoker curMethod = mBindingController.getCurMethod();
+        final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
+        final IInputMethodInvoker curMethod = userData.mBindingController.getCurMethod();
         if (curMethod == null) {
             // No need to send the data if the IME is not yet bound.
             return;
@@ -5917,9 +5941,10 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
             p.println("  mCurClient=" + client + " mCurSeq=" + getSequenceNumberLocked());
             p.println("  mFocusedWindowPerceptible=" + mFocusedWindowPerceptible);
             mImeBindingState.dump("  ", p);
+            final var userData = mUserDataRepository.getOrCreate(mCurrentUserId);
             p.println("  mCurId=" + getCurIdLocked() + " mHaveConnection=" + hasConnectionLocked()
                     + " mBoundToMethod=" + mBoundToMethod + " mVisibleBound="
-                    + mBindingController.isVisibleBound());
+                    + userData.mBindingController.isVisibleBound());
             p.println("  mCurToken=" + getCurTokenLocked());
             p.println("  mCurTokenDisplayId=" + mCurTokenDisplayId);
             p.println("  mCurHostInputToken=" + mCurHostInputToken);
@@ -6413,7 +6438,8 @@ public final class InputMethodManagerService implements IInputMethodManagerImpl.
                     if (userId == mCurrentUserId) {
                         hideCurrentInputLocked(mImeBindingState.mFocusedWindow, 0 /* flags */,
                                 SoftInputShowHideReason.HIDE_RESET_SHELL_COMMAND);
-                        mBindingController.unbindCurrentMethod();
+                        final var userData = mUserDataRepository.getOrCreate(userId);
+                        userData.mBindingController.unbindCurrentMethod();
 
                         // Enable default IMEs, disable others
                         var toDisable = settings.getEnabledInputMethodList();
diff --git a/services/core/java/com/android/server/inputmethod/UserDataRepository.java b/services/core/java/com/android/server/inputmethod/UserDataRepository.java
index 7f00229395f8..825cfcbdf505 100644
--- a/services/core/java/com/android/server/inputmethod/UserDataRepository.java
+++ b/services/core/java/com/android/server/inputmethod/UserDataRepository.java
@@ -15,6 +15,7 @@
  */
 
 package com.android.server.inputmethod;
+
 import android.annotation.NonNull;
 import android.annotation.UserIdInt;
 import android.content.pm.UserInfo;
@@ -25,18 +26,21 @@ import com.android.internal.annotations.GuardedBy;
 import com.android.server.pm.UserManagerInternal;
 
 import java.util.function.Consumer;
+import java.util.function.IntFunction;
 
 final class UserDataRepository {
 
     @GuardedBy("ImfLock.class")
     private final SparseArray<UserData> mUserData = new SparseArray<>();
 
+    private final IntFunction<InputMethodBindingController> mBindingControllerFactory;
+
     @GuardedBy("ImfLock.class")
     @NonNull
     UserData getOrCreate(@UserIdInt int userId) {
         UserData userData = mUserData.get(userId);
         if (userData == null) {
-            userData = new UserData(userId);
+            userData = new UserData(userId, mBindingControllerFactory.apply(userId));
             mUserData.put(userId, userData);
         }
         return userData;
@@ -49,7 +53,9 @@ final class UserDataRepository {
         }
     }
 
-    UserDataRepository(@NonNull Handler handler, @NonNull UserManagerInternal userManagerInternal) {
+    UserDataRepository(@NonNull Handler handler, @NonNull UserManagerInternal userManagerInternal,
+            @NonNull IntFunction<InputMethodBindingController> bindingControllerFactory) {
+        mBindingControllerFactory = bindingControllerFactory;
         userManagerInternal.addUserLifecycleListener(
                 new UserManagerInternal.UserLifecycleListener() {
                     @Override
@@ -79,11 +85,16 @@ final class UserDataRepository {
         @UserIdInt
         final int mUserId;
 
-       /**
+        @NonNull
+        final InputMethodBindingController mBindingController;
+
+        /**
          * Intended to be instantiated only from this file.
          */
-        private UserData(@UserIdInt int userId) {
+        private UserData(@UserIdInt int userId,
+                @NonNull InputMethodBindingController bindingController) {
             mUserId = userId;
+            mBindingController = bindingController;
         }
     }
 }
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodBindingControllerTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodBindingControllerTest.java
index 1f0a37509989..70903cbc2b94 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodBindingControllerTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodBindingControllerTest.java
@@ -77,9 +77,13 @@ public class InputMethodBindingControllerTest extends InputMethodManagerServiceT
         mCountDownLatch = new CountDownLatch(1);
         // Remove flag Context.BIND_SCHEDULE_LIKE_TOP_APP because in tests we are not calling
         // from system.
-        mBindingController =
-                new InputMethodBindingController(
-                        mInputMethodManagerService, mImeConnectionBindFlags, mCountDownLatch);
+        synchronized (ImfLock.class) {
+            mBindingController =
+                    new InputMethodBindingController(
+                            mInputMethodManagerService.getCurrentImeUserIdLocked(),
+                            mInputMethodManagerService, mImeConnectionBindFlags,
+                            mCountDownLatch);
+        }
     }
 
     @Test
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java
index b4cf79941c33..cff22654e30c 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java
@@ -222,7 +222,7 @@ public class InputMethodManagerServiceTestBase {
                         Process.THREAD_PRIORITY_FOREGROUND, /* allowIo */
                         false);
         mInputMethodManagerService = new InputMethodManagerService(mContext, mServiceThread,
-                mMockInputMethodBindingController);
+                unusedUserId -> mMockInputMethodBindingController);
         spyOn(mInputMethodManagerService);
 
         // Start a InputMethodManagerService.Lifecycle to publish and manage the lifecycle of
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/UserDataRepositoryTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/UserDataRepositoryTest.java
index a15b17042174..c3a87dafe7ce 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/UserDataRepositoryTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/UserDataRepositoryTest.java
@@ -18,6 +18,7 @@ package com.android.server.inputmethod;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
@@ -38,6 +39,7 @@ import org.mockito.MockitoAnnotations;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.IntFunction;
 
 // This test is designed to run on both device and host (Ravenwood) side.
 public final class UserDataRepositoryTest {
@@ -51,19 +53,34 @@ public final class UserDataRepositoryTest {
     @Mock
     private UserManagerInternal mMockUserManagerInternal;
 
+    @Mock
+    private InputMethodManagerService mMockInputMethodManagerService;
+
     private Handler mHandler;
 
+    private IntFunction<InputMethodBindingController> mBindingControllerFactory;
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mHandler = new Handler(Looper.getMainLooper());
+        mBindingControllerFactory = new IntFunction<InputMethodBindingController>() {
+
+            @Override
+            public InputMethodBindingController apply(int userId) {
+                return new InputMethodBindingController(userId, mMockInputMethodManagerService);
+            }
+        };
     }
 
     @Test
     public void testUserDataRepository_addsNewUserInfoOnUserCreatedEvent() {
         // Create UserDataRepository and capture the user lifecycle listener
         final var captor = ArgumentCaptor.forClass(UserManagerInternal.UserLifecycleListener.class);
-        final var repository = new UserDataRepository(mHandler, mMockUserManagerInternal);
+        final var bindingControllerFactorySpy = spy(mBindingControllerFactory);
+        final var repository = new UserDataRepository(mHandler, mMockUserManagerInternal,
+                bindingControllerFactorySpy);
+
         verify(mMockUserManagerInternal, times(1)).addUserLifecycleListener(captor.capture());
         final var listener = captor.getValue();
 
@@ -77,14 +94,20 @@ public final class UserDataRepositoryTest {
         // Assert UserDataRepository contains the expected UserData
         final var allUserData = collectUserData(repository);
         assertThat(allUserData).hasSize(1);
-        assertThat(allUserData.get(0).mUserId).isEqualTo(userInfo.id);
+        assertThat(allUserData.get(0).mUserId).isEqualTo(ANY_USER_ID);
+
+        // Assert UserDataRepository called the InputMethodBindingController creator function.
+        verify(bindingControllerFactorySpy).apply(ANY_USER_ID);
+        assertThat(allUserData.get(0).mBindingController.mUserId).isEqualTo(ANY_USER_ID);
     }
 
     @Test
     public void testUserDataRepository_removesUserInfoOnUserRemovedEvent() {
         // Create UserDataRepository and capture the user lifecycle listener
         final var captor = ArgumentCaptor.forClass(UserManagerInternal.UserLifecycleListener.class);
-        final var repository = new UserDataRepository(mHandler, mMockUserManagerInternal);
+        final var repository = new UserDataRepository(mHandler, mMockUserManagerInternal,
+                userId -> new InputMethodBindingController(userId, mMockInputMethodManagerService));
+
         verify(mMockUserManagerInternal, times(1)).addUserLifecycleListener(captor.capture());
         final var listener = captor.getValue();
 
@@ -104,7 +127,8 @@ public final class UserDataRepositoryTest {
 
     @Test
     public void testGetOrCreate() {
-        final var repository = new UserDataRepository(mHandler, mMockUserManagerInternal);
+        final var repository = new UserDataRepository(mHandler, mMockUserManagerInternal,
+                mBindingControllerFactory);
 
         synchronized (ImfLock.class) {
             final var userData = repository.getOrCreate(ANY_USER_ID);
@@ -114,6 +138,9 @@ public final class UserDataRepositoryTest {
         final var allUserData = collectUserData(repository);
         assertThat(allUserData).hasSize(1);
         assertThat(allUserData.get(0).mUserId).isEqualTo(ANY_USER_ID);
+
+        // Assert UserDataRepository called the InputMethodBindingController creator function.
+        assertThat(allUserData.get(0).mBindingController.mUserId).isEqualTo(ANY_USER_ID);
     }
 
     private List<UserDataRepository.UserData> collectUserData(UserDataRepository repository) {