Merge "[Panlingual] App locales restore Stage the per-app locales in the storage instead of memory" into main
diff --git a/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java b/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java
index 563f93e..b9e0960 100644
--- a/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java
+++ b/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java
@@ -84,9 +84,16 @@
      * from the delegate selector.
      */
     private static final String LOCALES_FROM_DELEGATE_PREFS = "LocalesFromDelegatePrefs.xml";
+    private static final String LOCALES_STAGED_DATA_PREFS = "LocalesStagedDataPrefs.xml";
+    private static final String ARCHIVED_PACKAGES_PREFS = "ArchivedPackagesPrefs.xml";
     // Stage data would be deleted on reboot since it's stored in memory. So it's retained until
     // retention period OR next reboot, whichever happens earlier.
     private static final Duration STAGE_DATA_RETENTION_PERIOD = Duration.ofDays(3);
+    // Store the locales staged data for the specified package in the SharedPreferences. The format
+    // is locales s:setFromDelegate
+    // For example: en-US s:true
+    private static final String STRING_SPLIT = " s:";
+    private static final String KEY_STAGED_DATA_TIME = "staged_data_time";
 
     private final LocaleManagerService mLocaleManagerService;
     private final PackageManager mPackageManager;
@@ -94,39 +101,34 @@
     private final Context mContext;
     private final Object mStagedDataLock = new Object();
 
-    // Staged data map keyed by user-id to handle multi-user scenario / work profiles. We are using
-    // SparseArray because it is more memory-efficient than a HashMap.
-    private final SparseArray<StagedData> mStagedData;
-
     // SharedPreferences to store packages whose app-locale was set by a delegate, as opposed to
     // the application setting the app-locale itself.
     private final SharedPreferences mDelegateAppLocalePackages;
+    // For unit tests
+    private final SparseArray<File> mStagedDataFiles;
+    private final File mArchivedPackagesFile;
+
     private final BroadcastReceiver mUserMonitor;
-    // To determine whether an app is pre-archived, check for Intent.EXTRA_ARCHIVAL upon receiving
-    // the initial PACKAGE_ADDED broadcast. If it is indeed pre-archived, perform the data
-    // restoration during the second PACKAGE_ADDED broadcast, which is sent subsequently when the
-    // app is installed.
-    private final Set<String> mPkgsToRestore;
 
     LocaleManagerBackupHelper(LocaleManagerService localeManagerService,
             PackageManager packageManager, HandlerThread broadcastHandlerThread) {
         this(localeManagerService.mContext, localeManagerService, packageManager, Clock.systemUTC(),
-                new SparseArray<>(), broadcastHandlerThread, null);
+                broadcastHandlerThread, null, null, null);
     }
 
-    @VisibleForTesting LocaleManagerBackupHelper(Context context,
-            LocaleManagerService localeManagerService,
-            PackageManager packageManager, Clock clock, SparseArray<StagedData> stagedData,
-            HandlerThread broadcastHandlerThread, SharedPreferences delegateAppLocalePackages) {
+    @VisibleForTesting
+    LocaleManagerBackupHelper(Context context, LocaleManagerService localeManagerService,
+            PackageManager packageManager, Clock clock, HandlerThread broadcastHandlerThread,
+            SparseArray<File> stagedDataFiles, File archivedPackagesFile,
+            SharedPreferences delegateAppLocalePackages) {
         mContext = context;
         mLocaleManagerService = localeManagerService;
         mPackageManager = packageManager;
         mClock = clock;
-        mStagedData = stagedData;
         mDelegateAppLocalePackages = delegateAppLocalePackages != null ? delegateAppLocalePackages
-                : createPersistedInfo();
-        mPkgsToRestore = new ArraySet<>();
-
+            : createPersistedInfo();
+        mArchivedPackagesFile = archivedPackagesFile;
+        mStagedDataFiles = stagedDataFiles;
         mUserMonitor = new UserMonitor();
         IntentFilter filter = new IntentFilter();
         filter.addAction(Intent.ACTION_USER_REMOVED);
@@ -148,7 +150,7 @@
         }
 
         synchronized (mStagedDataLock) {
-            cleanStagedDataForOldEntriesLocked();
+            cleanStagedDataForOldEntriesLocked(userId);
         }
 
         HashMap<String, LocalesInfo> pkgStates = new HashMap<>();
@@ -207,14 +209,11 @@
         return out.toByteArray();
     }
 
-    private void cleanStagedDataForOldEntriesLocked() {
-        for (int i = 0; i < mStagedData.size(); i++) {
-            int userId = mStagedData.keyAt(i);
-            StagedData stagedData = mStagedData.get(userId);
-            if (stagedData.mCreationTimeMillis
-                    < mClock.millis() - STAGE_DATA_RETENTION_PERIOD.toMillis()) {
-                deleteStagedDataLocked(userId);
-            }
+    private void cleanStagedDataForOldEntriesLocked(@UserIdInt int userId) {
+        Long created_time = getStagedDataSp(userId).getLong(KEY_STAGED_DATA_TIME, -1);
+        if (created_time != -1
+                && created_time < mClock.millis() - STAGE_DATA_RETENTION_PERIOD.toMillis()) {
+            deleteStagedDataLocked(userId);
         }
     }
 
@@ -252,20 +251,16 @@
         // performed simultaneously.
         synchronized (mStagedDataLock) {
             // Backups for apps which are yet to be installed.
-            StagedData stagedData = new StagedData(mClock.millis(), new HashMap<>());
-
             for (String pkgName : pkgStates.keySet()) {
                 LocalesInfo localesInfo = pkgStates.get(pkgName);
                 // Check if the application is already installed for the concerned user.
                 if (isPackageInstalledForUser(pkgName, userId)) {
-                    if (mPkgsToRestore != null) {
-                        mPkgsToRestore.remove(pkgName);
-                    }
+                    removeFromArchivedPackagesInfo(userId, pkgName);
                     // Don't apply the restore if the locales have already been set for the app.
                     checkExistingLocalesAndApplyRestore(pkgName, localesInfo, userId);
                 } else {
                     // Stage the data if the app isn't installed.
-                    stagedData.mPackageStates.put(pkgName, localesInfo);
+                    storeStagedDataInfo(userId, pkgName, localesInfo);
                     if (DEBUG) {
                         Slog.d(TAG, "Add locales=" + localesInfo.mLocales
                                 + " fromDelegate=" + localesInfo.mSetFromDelegate
@@ -274,8 +269,9 @@
                 }
             }
 
-            if (!stagedData.mPackageStates.isEmpty()) {
-                mStagedData.put(userId, stagedData);
+            // Create the time if the data is being staged.
+            if (!getStagedDataSp(userId).getAll().isEmpty()) {
+                storeStagedDataCreatedTime(userId);
             }
         }
     }
@@ -293,14 +289,23 @@
      * added on device.
      */
     void onPackageAddedWithExtras(String packageName, int uid, Bundle extras) {
-        boolean archived = false;
+        int userId = UserHandle.getUserId(uid);
         if (extras != null) {
-            archived = extras.getBoolean(Intent.EXTRA_ARCHIVAL, false);
-            if (archived && mPkgsToRestore != null) {
-                mPkgsToRestore.add(packageName);
+            // To determine whether an app is pre-archived, check for Intent.EXTRA_ARCHIVAL upon
+            // receiving the initial PACKAGE_ADDED broadcast. If it is indeed pre-archived, perform
+            // the data restoration during the second PACKAGE_ADDED broadcast, which is sent
+            // subsequently when the app is installed.
+            boolean archived = extras.getBoolean(Intent.EXTRA_ARCHIVAL, false);
+            if (DEBUG) {
+                Slog.d(TAG,
+                        "onPackageAddedWithExtras packageName: " + packageName + ", userId: "
+                                + userId + ", archived: " + archived);
+            }
+            if (archived) {
+                addInArchivedPackagesInfo(userId, packageName);
             }
         }
-        checkStageDataAndApplyRestore(packageName, uid);
+        checkStageDataAndApplyRestore(packageName, userId);
     }
 
     /**
@@ -310,9 +315,32 @@
      */
     void onPackageUpdateFinished(String packageName, int uid) {
         int userId = UserHandle.getUserId(uid);
-        if (mPkgsToRestore != null && mPkgsToRestore.contains(packageName)) {
-            mPkgsToRestore.remove(packageName);
-            checkStageDataAndApplyRestore(packageName, uid);
+        if (DEBUG) {
+            Slog.d(TAG,
+                    "onPackageUpdateFinished userId: " + userId + ", packageName: " + packageName);
+        }
+        String user = Integer.toString(userId);
+        File file = getArchivedPackagesFile();
+        if (file.exists()) {
+            SharedPreferences sp = getArchivedPackagesSp(file);
+            Set<String> packageNames = new ArraySet<>(sp.getStringSet(user, new ArraySet<>()));
+            if (packageNames.remove(packageName)) {
+                SharedPreferences.Editor editor = sp.edit();
+                if (packageNames.isEmpty()) {
+                    if (!editor.remove(user).commit()) {
+                        Slog.e(TAG, "Failed to remove the user");
+                    }
+                    if (sp.getAll().isEmpty()) {
+                        file.delete();
+                    }
+                } else {
+                    // commit and log the result.
+                    if (!editor.putStringSet(user, packageNames).commit()) {
+                        Slog.e(TAG, "failed to remove the package");
+                    }
+                }
+                checkStageDataAndApplyRestore(packageName, userId);
+            }
         }
         cleanApplicationLocalesIfNeeded(packageName, userId);
     }
@@ -347,16 +375,16 @@
         }
     }
 
-    private void checkStageDataAndApplyRestore(String packageName, int uid) {
+    private void checkStageDataAndApplyRestore(String packageName, int userId) {
         try {
             synchronized (mStagedDataLock) {
-                cleanStagedDataForOldEntriesLocked();
-
-                int userId = UserHandle.getUserId(uid);
-                if (mStagedData.contains(userId)) {
-                    if (mPkgsToRestore != null) {
-                        mPkgsToRestore.remove(packageName);
+                cleanStagedDataForOldEntriesLocked(userId);
+                if (!getStagedDataSp(userId).getString(packageName, "").isEmpty()) {
+                    if (DEBUG) {
+                        Slog.d(TAG,
+                                "checkStageDataAndApplyRestore, remove package and restore data");
                     }
+                    removeFromArchivedPackagesInfo(userId, packageName);
                     // Perform lazy restore only if the staged data exists.
                     doLazyRestoreLocked(packageName, userId);
                 }
@@ -417,8 +445,17 @@
         }
     }
 
-    private void deleteStagedDataLocked(@UserIdInt int userId) {
-        mStagedData.remove(userId);
+    void deleteStagedDataLocked(@UserIdInt int userId) {
+        File stagedFile = getStagedDataFile(userId);
+        SharedPreferences sp = getStagedDataSp(stagedFile);
+        // commit and log the result.
+        if (!sp.edit().clear().commit()) {
+            Slog.e(TAG, "Failed to commit data!");
+        }
+
+        if (stagedFile.exists()) {
+            stagedFile.delete();
+        }
     }
 
     /**
@@ -473,16 +510,6 @@
         out.endDocument();
     }
 
-    static class StagedData {
-        final long mCreationTimeMillis;
-        final HashMap<String, LocalesInfo> mPackageStates;
-
-        StagedData(long creationTimeMillis, HashMap<String, LocalesInfo> pkgStates) {
-            mCreationTimeMillis = creationTimeMillis;
-            mPackageStates = pkgStates;
-        }
-    }
-
     static class LocalesInfo {
         final String mLocales;
         final boolean mSetFromDelegate;
@@ -508,6 +535,7 @@
                     synchronized (mStagedDataLock) {
                         deleteStagedDataLocked(userId);
                         removeProfileFromPersistedInfo(userId);
+                        removeArchivedPackagesForUser(userId);
                     }
                 }
             } catch (Exception e) {
@@ -533,26 +561,159 @@
             return;
         }
 
-        StagedData stagedData = mStagedData.get(userId);
-        for (String pkgName : stagedData.mPackageStates.keySet()) {
-            LocalesInfo localesInfo = stagedData.mPackageStates.get(pkgName);
-
-            if (pkgName.equals(packageName)) {
-
-                checkExistingLocalesAndApplyRestore(pkgName, localesInfo, userId);
-
-                // Remove the restored entry from the staged data list.
-                stagedData.mPackageStates.remove(pkgName);
-
-                // Remove the stage data entry for user if there are no more packages to restore.
-                if (stagedData.mPackageStates.isEmpty()) {
-                    mStagedData.remove(userId);
-                }
-
-                // No need to loop further after restoring locales because the staged data will
-                // contain at most one entry for the newly added package.
-                break;
+        SharedPreferences sp = getStagedDataSp(userId);
+        String value = sp.getString(packageName, "");
+        if (!value.isEmpty()) {
+            String[] info = value.split(STRING_SPLIT);
+            if (info == null || info.length != 2) {
+                Slog.e(TAG, "Failed to restore data");
+                return;
             }
+            LocalesInfo localesInfo = new LocalesInfo(info[0], Boolean.parseBoolean(info[1]));
+            checkExistingLocalesAndApplyRestore(packageName, localesInfo, userId);
+
+            // Remove the restored entry from the staged data list.
+            if (!sp.edit().remove(packageName).commit()) {
+                Slog.e(TAG, "Failed to commit data!");
+            }
+        }
+
+        // Remove the stage data entry for user if there are no more packages to restore.
+        if (sp.getAll().size() == 1 && sp.getLong(KEY_STAGED_DATA_TIME, -1) != -1) {
+            deleteStagedDataLocked(userId);
+        }
+    }
+
+    private File getStagedDataFile(@UserIdInt int userId) {
+        return mStagedDataFiles == null ? new File(Environment.getDataSystemDeDirectory(userId),
+            LOCALES_STAGED_DATA_PREFS) : mStagedDataFiles.get(userId);
+    }
+
+    private SharedPreferences getStagedDataSp(File file) {
+        return mStagedDataFiles == null ? mContext.createDeviceProtectedStorageContext()
+            .getSharedPreferences(file, Context.MODE_PRIVATE)
+            : mContext.getSharedPreferences(file, Context.MODE_PRIVATE);
+    }
+
+    private SharedPreferences getStagedDataSp(@UserIdInt int userId) {
+        return mStagedDataFiles == null ? mContext.createDeviceProtectedStorageContext()
+            .getSharedPreferences(getStagedDataFile(userId), Context.MODE_PRIVATE)
+            : mContext.getSharedPreferences(mStagedDataFiles.get(userId), Context.MODE_PRIVATE);
+    }
+
+    /**
+     * Store the staged locales info.
+     */
+    private void storeStagedDataInfo(@UserIdInt int userId, @NonNull String packageName,
+            @NonNull LocalesInfo localesInfo) {
+        if (DEBUG) {
+            Slog.d(TAG, "storeStagedDataInfo, userId: " + userId + ", packageName: " + packageName
+                    + ", localesInfo.mLocales: " + localesInfo.mLocales
+                    + ", localesInfo.mSetFromDelegate: " + localesInfo.mSetFromDelegate);
+        }
+        String info =
+                localesInfo.mLocales + STRING_SPLIT + String.valueOf(localesInfo.mSetFromDelegate);
+        SharedPreferences sp = getStagedDataSp(userId);
+        // commit and log the result.
+        if (!sp.edit().putString(packageName, info).commit()) {
+            Slog.e(TAG, "Failed to commit data!");
+        }
+    }
+
+    /**
+     * Store the time of creation for staged locales info.
+     */
+    private void storeStagedDataCreatedTime(@UserIdInt int userId) {
+        SharedPreferences sp = getStagedDataSp(userId);
+        // commit and log the result.
+        if (!sp.edit().putLong(KEY_STAGED_DATA_TIME, mClock.millis()).commit()) {
+            Slog.e(TAG, "Failed to commit data!");
+        }
+    }
+
+    private File getArchivedPackagesFile() {
+        return mArchivedPackagesFile == null ? new File(
+            Environment.getDataSystemDeDirectory(UserHandle.USER_SYSTEM),
+            ARCHIVED_PACKAGES_PREFS) : mArchivedPackagesFile;
+    }
+
+    private SharedPreferences getArchivedPackagesSp(File file) {
+        return mArchivedPackagesFile == null ? mContext.createDeviceProtectedStorageContext()
+            .getSharedPreferences(file, Context.MODE_PRIVATE)
+            : mContext.getSharedPreferences(file, Context.MODE_PRIVATE);
+    }
+
+    /**
+     * Add the package into the archived packages list.
+     */
+    private void addInArchivedPackagesInfo(@UserIdInt int userId, @NonNull String packageName) {
+        String user = Integer.toString(userId);
+        SharedPreferences sp = getArchivedPackagesSp(getArchivedPackagesFile());
+        Set<String> packageNames = new ArraySet<>(sp.getStringSet(user, new ArraySet<>()));
+        if (DEBUG) {
+            Slog.d(TAG, "addInArchivedPackagesInfo before packageNames: " + packageNames
+                    + ", packageName: " + packageName);
+        }
+        if (packageNames.add(packageName)) {
+            // commit and log the result.
+            if (!sp.edit().putStringSet(user, packageNames).commit()) {
+                Slog.e(TAG, "failed to add the package");
+            }
+        }
+    }
+
+    /**
+     * Remove the package from the archived packages list.
+     */
+    private void removeFromArchivedPackagesInfo(@UserIdInt int userId,
+            @NonNull String packageName) {
+        File file = getArchivedPackagesFile();
+        if (file.exists()) {
+            String user = Integer.toString(userId);
+            SharedPreferences sp = getArchivedPackagesSp(getArchivedPackagesFile());
+            Set<String> packageNames = new ArraySet<>(sp.getStringSet(user, new ArraySet<>()));
+            if (DEBUG) {
+                Slog.d(TAG, "removeFromArchivedPackagesInfo before packageNames: " + packageNames
+                        + ", packageName: " + packageName);
+            }
+            if (packageNames.remove(packageName)) {
+                SharedPreferences.Editor editor = sp.edit();
+                if (packageNames.isEmpty()) {
+                    if (!editor.remove(user).commit()) {
+                        Slog.e(TAG, "Failed to remove user");
+                    }
+                    if (sp.getAll().isEmpty()) {
+                        file.delete();
+                    }
+                } else {
+                    // commit and log the result.
+                    if (!editor.putStringSet(user, packageNames).commit()) {
+                        Slog.e(TAG, "failed to remove the package");
+                    }
+                }
+            }
+        }
+    }
+
+    /**
+     * Remove the user from the archived packages list.
+     */
+    private void removeArchivedPackagesForUser(@UserIdInt int userId) {
+        String user = Integer.toString(userId);
+        File file = getArchivedPackagesFile();
+        SharedPreferences sp = getArchivedPackagesSp(file);
+
+        if (sp == null || !sp.contains(user)) {
+            Slog.w(TAG, "The profile is not existed in the archived package info");
+            return;
+        }
+
+        if (!sp.edit().remove(user).commit()) {
+            Slog.e(TAG, "Failed to remove user");
+        }
+
+        if (sp.getAll().isEmpty() && file.exists()) {
+            file.delete();
         }
     }
 
diff --git a/services/tests/servicestests/src/com/android/server/locales/LocaleManagerBackupRestoreTest.java b/services/tests/servicestests/src/com/android/server/locales/LocaleManagerBackupRestoreTest.java
index 7dd1847..50cfa75 100644
--- a/services/tests/servicestests/src/com/android/server/locales/LocaleManagerBackupRestoreTest.java
+++ b/services/tests/servicestests/src/com/android/server/locales/LocaleManagerBackupRestoreTest.java
@@ -20,7 +20,6 @@
 import static junit.framework.Assert.assertNull;
 import static junit.framework.Assert.assertTrue;
 
-import static org.junit.Assert.assertNotNull;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
@@ -52,6 +51,7 @@
 import android.util.SparseArray;
 import android.util.Xml;
 
+import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.internal.content.PackageMonitor;
@@ -70,6 +70,7 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
+import java.io.File;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.charset.StandardCharsets;
@@ -95,21 +96,21 @@
     private static final int DEFAULT_USER_ID = 0;
     private static final int WORK_PROFILE_USER_ID = 10;
     private static final int DEFAULT_UID = Binder.getCallingUid() + 100;
+    private static final int WORK_PROFILE_UID = Binder.getCallingUid() + 1000100;
     private static final long DEFAULT_CREATION_TIME_MILLIS = 1000;
     private static final Duration RETENTION_PERIOD = Duration.ofDays(3);
     private static final LocaleList DEFAULT_LOCALES =
             LocaleList.forLanguageTags(DEFAULT_LOCALE_TAGS);
     private static final Map<String, LocalesInfo> DEFAULT_PACKAGE_LOCALES_INFO_MAP = Map.of(
             DEFAULT_PACKAGE_NAME, new LocalesInfo(DEFAULT_LOCALE_TAGS, false));
-    private static final SparseArray<LocaleManagerBackupHelper.StagedData> STAGE_DATA =
-            new SparseArray<>();
+    private final SparseArray<File> mStagedDataFiles = new SparseArray<>();
+    private File mArchivedPackageFile;
 
     private LocaleManagerBackupHelper mBackupHelper;
     private long mCurrentTimeMillis;
+    private Context mContext = spy(ApplicationProvider.getApplicationContext());
 
     @Mock
-    private Context mMockContext;
-    @Mock
     private PackageManager mMockPackageManager;
     @Mock
     private LocaleManagerService mMockLocaleManagerService;
@@ -138,23 +139,28 @@
 
     @Before
     public void setUp() throws Exception {
-        mMockContext = mock(Context.class);
         mMockPackageManager = mock(PackageManager.class);
         mMockLocaleManagerService = mock(LocaleManagerService.class);
         mMockDelegateAppLocalePackages = mock(SharedPreferences.class);
         mMockSpEditor = mock(SharedPreferences.Editor.class);
         SystemAppUpdateTracker systemAppUpdateTracker = mock(SystemAppUpdateTracker.class);
 
-        doReturn(mMockPackageManager).when(mMockContext).getPackageManager();
+        doReturn(mMockPackageManager).when(mContext).getPackageManager();
         doReturn(mMockSpEditor).when(mMockDelegateAppLocalePackages).edit();
 
         HandlerThread broadcastHandlerThread = new HandlerThread(TAG,
                 Process.THREAD_PRIORITY_BACKGROUND);
         broadcastHandlerThread.start();
 
-        mBackupHelper = spy(new ShadowLocaleManagerBackupHelper(mMockContext,
-                mMockLocaleManagerService, mMockPackageManager, mClock, STAGE_DATA,
-                broadcastHandlerThread, mMockDelegateAppLocalePackages));
+        File file0 = new File(mContext.getCacheDir(), "file_user_0.txt");
+        File file10 = new File(mContext.getCacheDir(), "file_user_10.txt");
+        mStagedDataFiles.put(DEFAULT_USER_ID, file0);
+        mStagedDataFiles.put(WORK_PROFILE_USER_ID, file10);
+        mArchivedPackageFile = new File(mContext.getCacheDir(), "file_archived.txt");
+
+        mBackupHelper = spy(new ShadowLocaleManagerBackupHelper(mContext,
+            mMockLocaleManagerService, mMockPackageManager, mClock, broadcastHandlerThread,
+                mStagedDataFiles, mArchivedPackageFile, mMockDelegateAppLocalePackages));
         doNothing().when(mBackupHelper).notifyBackupManager();
 
         mUserMonitor = mBackupHelper.getUserMonitor();
@@ -165,7 +171,16 @@
 
     @After
     public void tearDown() throws Exception {
-        STAGE_DATA.clear();
+        for (int i = 0; i < mStagedDataFiles.size(); i++) {
+            int userId = mStagedDataFiles.keyAt(i);
+            File file = mStagedDataFiles.get(userId);
+            SharedPreferences sp = mContext.getSharedPreferences(file, Context.MODE_PRIVATE);
+            sp.edit().clear().commit();
+            if (file.exists()) {
+                file.delete();
+            }
+        }
+        mStagedDataFiles.clear();
     }
 
     @Test
@@ -543,17 +558,21 @@
         mPackageMonitor.onPackageAddedWithExtras(pkgNameA, DEFAULT_UID, bundle);
         mPackageMonitor.onPackageAddedWithExtras(pkgNameB, DEFAULT_UID, bundle);
 
+        checkArchivedFileExists();
+
         mBackupHelper.stageAndApplyRestoredPayload(out.toByteArray(), DEFAULT_USER_ID);
 
         verifyNothingRestored();
 
         setUpPackageInstalled(pkgNameA);
 
-        mPackageMonitor.onPackageUpdateFinished(pkgNameA, DEFAULT_UID);
+        mBackupHelper.onPackageUpdateFinished(pkgNameA, DEFAULT_UID);
 
         verify(mMockLocaleManagerService, times(1)).setApplicationLocales(pkgNameA, DEFAULT_USER_ID,
                 LocaleList.forLanguageTags(langTagsA), false, FrameworkStatsLog
                 .APPLICATION_LOCALES_CHANGED__CALLER__CALLER_BACKUP_RESTORE);
+        checkArchivedFileExists();
+
 
         mBackupHelper.persistLocalesModificationInfo(DEFAULT_USER_ID, pkgNameA, false, false);
 
@@ -565,11 +584,12 @@
 
         setUpPackageInstalled(pkgNameB);
 
-        mPackageMonitor.onPackageUpdateFinished(pkgNameB, DEFAULT_UID);
+        mBackupHelper.onPackageUpdateFinished(pkgNameB, DEFAULT_UID);
 
         verify(mMockLocaleManagerService, times(1)).setApplicationLocales(pkgNameB, DEFAULT_USER_ID,
                 LocaleList.forLanguageTags(langTagsB), true, FrameworkStatsLog
                 .APPLICATION_LOCALES_CHANGED__CALLER__CALLER_BACKUP_RESTORE);
+        checkArchivedFileDoesNotExist();
 
         mBackupHelper.persistLocalesModificationInfo(DEFAULT_USER_ID, pkgNameB, true, false);
 
@@ -723,7 +743,7 @@
         Intent intent = new Intent();
         intent.setAction(Intent.ACTION_USER_REMOVED);
         intent.putExtra(Intent.EXTRA_USER_HANDLE, DEFAULT_USER_ID);
-        mUserMonitor.onReceive(mMockContext, intent);
+        mUserMonitor.onReceive(mContext, intent);
 
         // Stage data should be removed only for DEFAULT_USER_ID.
         checkStageDataDoesNotExist(DEFAULT_USER_ID);
@@ -732,6 +752,72 @@
     }
 
     @Test
+    public void testRestore_multipleProfile_restoresFromStage_ArchiveEnabled() throws Exception {
+        final ByteArrayOutputStream outDefault = new ByteArrayOutputStream();
+        writeTestPayload(outDefault, DEFAULT_PACKAGE_LOCALES_INFO_MAP);
+        final ByteArrayOutputStream outWorkProfile = new ByteArrayOutputStream();
+        String anotherPackage = "com.android.anotherapp";
+        String anotherLangTags = "mr,zh";
+        LocalesInfo localesInfo = new LocalesInfo(anotherLangTags, true);
+        HashMap<String, LocalesInfo> pkgLocalesMapWorkProfile = new HashMap<>();
+        pkgLocalesMapWorkProfile.put(anotherPackage, localesInfo);
+        writeTestPayload(outWorkProfile, pkgLocalesMapWorkProfile);
+        // DEFAULT_PACKAGE_NAME is NOT installed on the device.
+        setUpPackageNotInstalled(DEFAULT_PACKAGE_NAME);
+        setUpPackageNotInstalled(anotherPackage);
+        setUpLocalesForPackage(DEFAULT_PACKAGE_NAME, LocaleList.getEmptyLocaleList());
+        setUpLocalesForPackage(anotherPackage, LocaleList.getEmptyLocaleList());
+        setUpPackageNamesForSp(new ArraySet<>());
+
+        Bundle bundle = new Bundle();
+        bundle.putBoolean(Intent.EXTRA_ARCHIVAL, true);
+        mPackageMonitor.onPackageAddedWithExtras(DEFAULT_PACKAGE_NAME, DEFAULT_UID, bundle);
+        mPackageMonitor.onPackageAddedWithExtras(anotherPackage, WORK_PROFILE_UID, bundle);
+
+        checkArchivedFileExists();
+
+        mBackupHelper.stageAndApplyRestoredPayload(outDefault.toByteArray(), DEFAULT_USER_ID);
+        mBackupHelper.stageAndApplyRestoredPayload(outWorkProfile.toByteArray(),
+                WORK_PROFILE_USER_ID);
+
+        verifyNothingRestored();
+        verifyStageDataForUser(DEFAULT_PACKAGE_LOCALES_INFO_MAP,
+                DEFAULT_CREATION_TIME_MILLIS, DEFAULT_USER_ID);
+        verifyStageDataForUser(pkgLocalesMapWorkProfile,
+                DEFAULT_CREATION_TIME_MILLIS, WORK_PROFILE_USER_ID);
+
+        setUpPackageInstalled(DEFAULT_PACKAGE_NAME);
+        mBackupHelper.onPackageUpdateFinished(DEFAULT_PACKAGE_NAME, DEFAULT_UID);
+
+        verify(mMockLocaleManagerService, times(1)).setApplicationLocales(DEFAULT_PACKAGE_NAME,
+                DEFAULT_USER_ID,
+                LocaleList.forLanguageTags(DEFAULT_LOCALE_TAGS), false, FrameworkStatsLog
+                        .APPLICATION_LOCALES_CHANGED__CALLER__CALLER_BACKUP_RESTORE);
+        checkArchivedFileExists();
+        checkStageDataDoesNotExist(DEFAULT_USER_ID);
+
+        mBackupHelper.persistLocalesModificationInfo(DEFAULT_USER_ID, DEFAULT_PACKAGE_NAME, false,
+                false);
+
+        verify(mMockSpEditor, times(0)).putStringSet(anyString(), any());
+
+        setUpPackageInstalled(anotherPackage);
+        mBackupHelper.onPackageUpdateFinished(anotherPackage, WORK_PROFILE_UID);
+
+        verify(mMockLocaleManagerService, times(1)).setApplicationLocales(anotherPackage,
+                WORK_PROFILE_USER_ID,
+                LocaleList.forLanguageTags(anotherLangTags), true, FrameworkStatsLog
+                        .APPLICATION_LOCALES_CHANGED__CALLER__CALLER_BACKUP_RESTORE);
+        checkArchivedFileDoesNotExist();
+
+        mBackupHelper.persistLocalesModificationInfo(DEFAULT_USER_ID, anotherPackage, true, false);
+
+        verify(mMockSpEditor, times(1)).putStringSet(Integer.toString(DEFAULT_USER_ID),
+            new ArraySet<>(Arrays.asList(anotherPackage)));
+        checkStageDataDoesNotExist(WORK_PROFILE_USER_ID);
+    }
+
+    @Test
     public void testPackageRemoved_noInfoInSp() throws Exception {
         String pkgNameA = "com.android.myAppA";
         String pkgNameB = "com.android.myAppB";
@@ -858,10 +944,22 @@
 
     private void verifyStageDataForUser(Map<String, LocalesInfo> expectedPkgLocalesMap,
             long expectedCreationTimeMillis, int userId) {
-        LocaleManagerBackupHelper.StagedData stagedDataForUser = STAGE_DATA.get(userId);
-        assertNotNull(stagedDataForUser);
-        assertEquals(expectedCreationTimeMillis, stagedDataForUser.mCreationTimeMillis);
-        verifyStageData(expectedPkgLocalesMap, stagedDataForUser.mPackageStates);
+        SharedPreferences sp = mContext.getSharedPreferences(mStagedDataFiles.get(userId),
+                Context.MODE_PRIVATE);
+        assertTrue(sp.getAll().size() > 0);
+        assertEquals(expectedCreationTimeMillis, sp.getLong("staged_data_time", -1));
+        verifyStageData(expectedPkgLocalesMap, sp);
+    }
+
+    private static void verifyStageData(Map<String, LocalesInfo> expectedPkgLocalesMap,
+            SharedPreferences sp) {
+        for (String pkg : expectedPkgLocalesMap.keySet()) {
+            assertTrue(!sp.getString(pkg, "").isEmpty());
+            String[] info = sp.getString(pkg, "").split(" s:");
+            assertEquals(expectedPkgLocalesMap.get(pkg).mLocales, info[0]);
+            assertEquals(expectedPkgLocalesMap.get(pkg).mSetFromDelegate,
+                    Boolean.parseBoolean(info[1]));
+        }
     }
 
     private static void verifyStageData(Map<String, LocalesInfo> expectedPkgLocalesMap,
@@ -875,11 +973,19 @@
         }
     }
 
-    private static void checkStageDataExists(int userId) {
-        assertNotNull(STAGE_DATA.get(userId));
+    private void checkStageDataExists(int userId) {
+        assertTrue(mStagedDataFiles.get(userId) != null && mStagedDataFiles.get(userId).exists());
     }
 
-    private static void checkStageDataDoesNotExist(int userId) {
-        assertNull(STAGE_DATA.get(userId));
+    private void checkStageDataDoesNotExist(int userId) {
+        assertTrue(mStagedDataFiles.get(userId) == null || !mStagedDataFiles.get(userId).exists());
     }
-}
+
+    private void checkArchivedFileExists() {
+        assertTrue(mArchivedPackageFile.exists());
+    }
+
+    private void checkArchivedFileDoesNotExist() {
+        assertTrue(!mArchivedPackageFile.exists());
+    }
+}
\ No newline at end of file
diff --git a/services/tests/servicestests/src/com/android/server/locales/ShadowLocaleManagerBackupHelper.java b/services/tests/servicestests/src/com/android/server/locales/ShadowLocaleManagerBackupHelper.java
index 9f7cbe3..b46902d 100644
--- a/services/tests/servicestests/src/com/android/server/locales/ShadowLocaleManagerBackupHelper.java
+++ b/services/tests/servicestests/src/com/android/server/locales/ShadowLocaleManagerBackupHelper.java
@@ -22,6 +22,7 @@
 import android.os.HandlerThread;
 import android.util.SparseArray;
 
+import java.io.File;
 import java.time.Clock;
 
 /**
@@ -33,9 +34,9 @@
     ShadowLocaleManagerBackupHelper(Context context,
             LocaleManagerService localeManagerService,
             PackageManager packageManager, Clock clock,
-            SparseArray<LocaleManagerBackupHelper.StagedData> stagedData,
-            HandlerThread broadcastHandlerThread, SharedPreferences delegateAppLocalePackages) {
-        super(context, localeManagerService, packageManager, clock, stagedData,
-                broadcastHandlerThread, delegateAppLocalePackages);
+            HandlerThread broadcastHandlerThread, SparseArray<File> stagedDataFiles,
+            File archivedPackagesFile, SharedPreferences delegateAppLocalePackages) {
+        super(context, localeManagerService, packageManager, clock, broadcastHandlerThread,
+                stagedDataFiles, archivedPackagesFile, delegateAppLocalePackages);
     }
 }