Merge "Remove ImfLock dependencies from *Repository" into main
diff --git a/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java b/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java
index 8ca0458..99f4747 100644
--- a/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java
+++ b/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java
@@ -20,9 +20,9 @@
 import android.annotation.NonNull;
 import android.annotation.UserIdInt;
 import android.annotation.WorkerThread;
-import android.os.Handler;
 import android.os.Process;
 import android.util.IntArray;
+import android.util.Slog;
 import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
@@ -36,7 +36,10 @@
  * persistent storages.
  */
 final class AdditionalSubtypeMapRepository {
-    @GuardedBy("ImfLock.class")
+    private static final String TAG = "AdditionalSubtypeMapRepository";
+
+    // TODO(b/352594784): Should we user other lock primitives?
+    @GuardedBy("sPerUserMap")
     @NonNull
     private static final SparseArray<AdditionalSubtypeMap> sPerUserMap = new SparseArray<>();
 
@@ -192,29 +195,77 @@
     private AdditionalSubtypeMapRepository() {
     }
 
+    /**
+     * Returns {@link AdditionalSubtypeMap} for the given user.
+     *
+     * <p>This method is expected be called after {@link #ensureInitializedAndGet(int)}. Otherwise
+     * {@link AdditionalSubtypeMap#EMPTY_MAP} will be returned.</p>
+     *
+     * @param userId the user to be queried about
+     * @return {@link AdditionalSubtypeMap} for the given user
+     */
+    @AnyThread
     @NonNull
-    @GuardedBy("ImfLock.class")
     static AdditionalSubtypeMap get(@UserIdInt int userId) {
-        final AdditionalSubtypeMap map = sPerUserMap.get(userId);
-        if (map != null) {
-            return map;
+        final AdditionalSubtypeMap map;
+        synchronized (sPerUserMap) {
+            map = sPerUserMap.get(userId);
         }
-        final AdditionalSubtypeMap newMap = AdditionalSubtypeUtils.load(userId);
-        sPerUserMap.put(userId, newMap);
-        return newMap;
+        if (map == null) {
+            Slog.e(TAG, "get(userId=" + userId + ") is called before loadInitialDataAndGet()."
+                    + " Returning an empty map");
+            return AdditionalSubtypeMap.EMPTY_MAP;
+        }
+        return map;
     }
 
-    @GuardedBy("ImfLock.class")
+    /**
+     * Ensures that {@link AdditionalSubtypeMap} is initialized for the given user. Load it from
+     * the persistent storage if {@link #putAndSave(int, AdditionalSubtypeMap, InputMethodMap)} has
+     * not been called yet.
+     *
+     * @param userId the user to be initialized
+     * @return {@link AdditionalSubtypeMap} that is associated with the given user. If
+     *         {@link #putAndSave(int, AdditionalSubtypeMap, InputMethodMap)} is already called
+     *         then the given {@link AdditionalSubtypeMap}.
+     */
+    @AnyThread
+    @NonNull
+    static AdditionalSubtypeMap ensureInitializedAndGet(@UserIdInt int userId) {
+        final var map = AdditionalSubtypeUtils.load(userId);
+        synchronized (sPerUserMap) {
+            final AdditionalSubtypeMap previous = sPerUserMap.get(userId);
+            // If putAndSave() has already been called, then use it.
+            if (previous != null) {
+                return previous;
+            }
+            sPerUserMap.put(userId, map);
+        }
+        return map;
+    }
+
+    /**
+     * Puts {@link AdditionalSubtypeMap} for the given user then schedule an I/O task to save it
+     * to the storage.
+     *
+     * @param userId         the user for the given {@link AdditionalSubtypeMap} is to be saved
+     * @param map            {@link AdditionalSubtypeMap} to be saved
+     * @param inputMethodMap {@link InputMethodMap} to be used while saving the data
+     */
+    @AnyThread
     static void putAndSave(@UserIdInt int userId, @NonNull AdditionalSubtypeMap map,
             @NonNull InputMethodMap inputMethodMap) {
-        final AdditionalSubtypeMap previous = sPerUserMap.get(userId);
-        if (previous == map) {
-            return;
+        synchronized (sPerUserMap) {
+            final AdditionalSubtypeMap previous = sPerUserMap.get(userId);
+            if (previous == map) {
+                return;
+            }
+            sPerUserMap.put(userId, map);
+            sWriter.scheduleWriteTask(userId, map, inputMethodMap);
         }
-        sPerUserMap.put(userId, map);
-        sWriter.scheduleWriteTask(userId, map, inputMethodMap);
     }
 
+    @AnyThread
     static void startWriterThread() {
         sWriter.startThread();
     }
@@ -225,12 +276,10 @@
     }
 
     @AnyThread
-    static void remove(@UserIdInt int userId, @NonNull Handler ioHandler) {
-        sWriter.onUserRemoved(userId);
-        ioHandler.post(() -> {
-            synchronized (ImfLock.class) {
-                sPerUserMap.remove(userId);
-            }
-        });
+    static void remove(@UserIdInt int userId) {
+        synchronized (sPerUserMap) {
+            sWriter.onUserRemoved(userId);
+            sPerUserMap.remove(userId);
+        }
     }
 }
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index fbd9ac0..96ca570 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -219,6 +219,13 @@
     static final String TAG = "InputMethodManagerService";
     public static final String PROTO_ARG = "--proto";
 
+    /**
+     * Timeout in milliseconds in {@link #systemRunning()} to make sure that users are initialized
+     * in {@link Lifecycle#initializeUsersAsync(int[])}.
+     */
+    @DurationMillisLong
+    private static final long SYSTEM_READY_USER_INIT_TIMEOUT = 3000;
+
     @Retention(SOURCE)
     @IntDef({ShellCommandResult.SUCCESS, ShellCommandResult.FAILURE})
     private @interface ShellCommandResult {
@@ -1035,7 +1042,7 @@
             // Called directly from UserManagerService. Do not block the calling thread.
             final int userId = user.id;
             SecureSettingsWrapper.onUserRemoved(userId);
-            AdditionalSubtypeMapRepository.remove(userId, mService.mIoHandler);
+            AdditionalSubtypeMapRepository.remove(userId);
             InputMethodSettingsRepository.remove(userId);
             mService.mUserDataRepository.remove(userId);
         }
@@ -1066,39 +1073,31 @@
 
         @AnyThread
         private void initializeUsersAsync(@UserIdInt int[] userIds) {
+            Slog.d(TAG, "Schedule initialization for users=" + Arrays.toString(userIds));
             mService.mIoHandler.post(() -> {
                 final var service = mService;
                 final var context = service.mContext;
                 final var userManagerInternal = service.mUserManagerInternal;
 
-                // We first create InputMethodMap for each user without loading AdditionalSubtypes.
-                final int numUsers = userIds.length;
-                final InputMethodMap[] rawMethodMaps = new InputMethodMap[numUsers];
-                for (int i = 0; i < numUsers; ++i) {
-                    final int userId = userIds[i];
-                    rawMethodMaps[i] = InputMethodManagerService.queryInputMethodServicesInternal(
-                            context, userId, AdditionalSubtypeMap.EMPTY_MAP,
+                for (int userId : userIds) {
+                    Slog.d(TAG, "Start initialization for user=" + userId);
+                    final var additionalSubtypeMap =
+                            AdditionalSubtypeMapRepository.ensureInitializedAndGet(userId);
+                    final var settings = InputMethodManagerService.queryInputMethodServicesInternal(
+                            context, userId, additionalSubtypeMap,
                             DirectBootAwareness.AUTO).getMethodMap();
+                    InputMethodSettingsRepository.put(userId,
+                            InputMethodSettings.create(settings, userId));
+
                     final int profileParentId = userManagerInternal.getProfileParentId(userId);
                     final boolean value =
                             InputMethodDrawsNavBarResourceMonitor.evaluate(context,
                                     profileParentId);
                     final var userData = mService.getUserData(userId);
                     userData.mImeDrawsNavBar.set(value);
-                }
 
-                // Then create full InputMethodMap for each user. Note that
-                // AdditionalSubtypeMapRepository#get() and InputMethodSettingsRepository#put()
-                // need to be called with ImfLock held (b/352387655).
-                // TODO(b/343601565): Avoid ImfLock after fixing b/352387655.
-                synchronized (ImfLock.class) {
-                    for (int i = 0; i < numUsers; ++i) {
-                        final int userId = userIds[i];
-                        final var map = AdditionalSubtypeMapRepository.get(userId);
-                        final var methodMap = rawMethodMaps[i].applyAdditionalSubtypes(map);
-                        final var settings = InputMethodSettings.create(methodMap, userId);
-                        InputMethodSettingsRepository.put(userId, settings);
-                    }
+                    userData.mBackgroundLoadLatch.countDown();
+                    Slog.d(TAG, "Complete initialization for user=" + userId);
                 }
             });
         }
@@ -1345,10 +1344,42 @@
         }
     }
 
+    private void waitForUserInitialization() {
+        final int[] userIds = mUserManagerInternal.getUserIds();
+        final long deadlineNanos = SystemClock.elapsedRealtimeNanos()
+                + TimeUnit.MILLISECONDS.toNanos(SYSTEM_READY_USER_INIT_TIMEOUT);
+        boolean interrupted = false;
+        try {
+            for (int userId : userIds) {
+                final var latch = getUserData(userId).mBackgroundLoadLatch;
+                boolean awaitResult;
+                while (true) {
+                    try {
+                        final long remainingNanos =
+                                Math.max(deadlineNanos - SystemClock.elapsedRealtimeNanos(), 0);
+                        awaitResult = latch.await(remainingNanos, TimeUnit.NANOSECONDS);
+                        break;
+                    } catch (InterruptedException ignored) {
+                        interrupted = true;
+                    }
+                }
+                if (!awaitResult) {
+                    Slog.w(TAG, "Timed out for user#" + userId + " to be initialized");
+                }
+            }
+        } finally {
+            if (interrupted) {
+                Thread.currentThread().interrupt();
+            }
+        }
+    }
+
     /**
      * TODO(b/32343335): The entire systemRunning() method needs to be revisited.
      */
     public void systemRunning() {
+        waitForUserInitialization();
+
         synchronized (ImfLock.class) {
             if (DEBUG) {
                 Slog.d(TAG, "--- systemReady");
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java b/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java
index 50ba364..1b84036 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java
@@ -24,7 +24,8 @@
 import com.android.internal.annotations.GuardedBy;
 
 final class InputMethodSettingsRepository {
-    @GuardedBy("ImfLock.class")
+    // TODO(b/352594784): Should we user other lock primitives?
+    @GuardedBy("sPerUserMap")
     @NonNull
     private static final SparseArray<InputMethodSettings> sPerUserMap = new SparseArray<>();
 
@@ -35,23 +36,28 @@
     }
 
     @NonNull
-    @GuardedBy("ImfLock.class")
+    @AnyThread
     static InputMethodSettings get(@UserIdInt int userId) {
-        final InputMethodSettings obj = sPerUserMap.get(userId);
+        final InputMethodSettings obj;
+        synchronized (sPerUserMap) {
+            obj = sPerUserMap.get(userId);
+        }
         if (obj != null) {
             return obj;
         }
         return InputMethodSettings.createEmptyMap(userId);
     }
 
-    @GuardedBy("ImfLock.class")
+    @AnyThread
     static void put(@UserIdInt int userId, @NonNull InputMethodSettings obj) {
-        sPerUserMap.put(userId, obj);
+        synchronized (sPerUserMap) {
+            sPerUserMap.put(userId, obj);
+        }
     }
 
     @AnyThread
     static void remove(@UserIdInt int userId) {
-        synchronized (ImfLock.class) {
+        synchronized (sPerUserMap) {
             sPerUserMap.remove(userId);
         }
     }
diff --git a/services/core/java/com/android/server/inputmethod/UserData.java b/services/core/java/com/android/server/inputmethod/UserData.java
index ec5c9e6..be57321 100644
--- a/services/core/java/com/android/server/inputmethod/UserData.java
+++ b/services/core/java/com/android/server/inputmethod/UserData.java
@@ -28,6 +28,7 @@
 import com.android.internal.inputmethod.IRemoteAccessibilityInputConnection;
 import com.android.internal.inputmethod.IRemoteInputConnection;
 
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /** Placeholder for all IMMS user specific fields */
@@ -35,6 +36,13 @@
     @UserIdInt
     final int mUserId;
 
+    /**
+     * Tells whether {@link InputMethodManagerService.Lifecycle#initializeUsersAsync(int[])} is
+     * completed for this user or not.
+     */
+    @NonNull
+    final CountDownLatch mBackgroundLoadLatch = new CountDownLatch(1);
+
     @NonNull
     final InputMethodBindingController mBindingController;
 
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 c2a069d..267ce26 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java
@@ -161,6 +161,7 @@
                         .spyStatic(InputMethodUtils.class)
                         .mockStatic(ServiceManager.class)
                         .spyStatic(AdditionalSubtypeMapRepository.class)
+                        .spyStatic(AdditionalSubtypeUtils.class)
                         .startMocking();
 
         mContext = spy(InstrumentationRegistry.getInstrumentation().getContext());
@@ -235,6 +236,7 @@
 
         // The background writer thread in AdditionalSubtypeMapRepository should be stubbed out.
         doNothing().when(AdditionalSubtypeMapRepository::startWriterThread);
+        doReturn(AdditionalSubtypeMap.EMPTY_MAP).when(() -> AdditionalSubtypeUtils.load(anyInt()));
 
         mServiceThread =
                 new ServiceThread(
@@ -267,6 +269,10 @@
         LocalServices.removeServiceForTest(InputMethodManagerInternal.class);
         lifecycle.onStart();
 
+        // Emulate that the user initialization is done.
+        AdditionalSubtypeMapRepository.ensureInitializedAndGet(mCallingUserId);
+        mInputMethodManagerService.getUserData(mCallingUserId).mBackgroundLoadLatch.countDown();
+
         // After this boot phase, services can broadcast Intents.
         lifecycle.onBootPhase(SystemService.PHASE_ACTIVITY_MANAGER_READY);