Merge "Simplify app-locales B&R stage data retention."

3 files changed, 155 insertions, 381 deletions

diff --git a/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java b/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java
index b2bd47b55cfb..134fb967cac5 100644
--- a/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java
+++ b/services/core/java/com/android/server/locales/LocaleManagerBackupHelper.java
@@ -21,7 +21,6 @@ import static android.os.UserHandle.USER_NULL;
 import static com.android.server.locales.LocaleManagerService.DEBUG;
 
 import android.annotation.NonNull;
-import android.annotation.Nullable;
 import android.annotation.UserIdInt;
 import android.app.backup.BackupManager;
 import android.content.BroadcastReceiver;
@@ -34,7 +33,6 @@ import android.content.pm.PackageManager;
 import android.content.pm.PackageManagerInternal;
 import android.os.BestClock;
 import android.os.Binder;
-import android.os.Environment;
 import android.os.HandlerThread;
 import android.os.LocaleList;
 import android.os.Process;
@@ -42,7 +40,6 @@ import android.os.RemoteException;
 import android.os.SystemClock;
 import android.os.UserHandle;
 import android.text.TextUtils;
-import android.util.AtomicFile;
 import android.util.Slog;
 import android.util.SparseArray;
 import android.util.TypedXmlPullParser;
@@ -53,17 +50,12 @@ import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.content.PackageMonitor;
 import com.android.internal.util.XmlUtils;
 
-import libcore.io.IoUtils;
-
 import org.xmlpull.v1.XmlPullParser;
 import org.xmlpull.v1.XmlPullParserException;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
-import java.io.File;
-import java.io.FileOutputStream;
 import java.io.IOException;
-import java.io.InputStream;
 import java.io.OutputStream;
 import java.io.UnsupportedEncodingException;
 import java.nio.charset.StandardCharsets;
@@ -71,8 +63,6 @@ import java.time.Clock;
 import java.time.Duration;
 import java.time.ZoneOffset;
 import java.util.HashMap;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
 
 /**
  * Helper class for managing backup and restore of app-specific locales.
@@ -87,33 +77,28 @@ class LocaleManagerBackupHelper {
     private static final String ATTR_LOCALES = "locales";
     private static final String ATTR_CREATION_TIME_MILLIS = "creationTimeMillis";
 
-    private static final String STAGE_FILE_NAME = "staged_locales";
     private static final String SYSTEM_BACKUP_PACKAGE_KEY = "android";
-
-    private static final Pattern STAGE_FILE_NAME_PATTERN = Pattern.compile(
-            TextUtils.formatSimple("(^%s_)(\\d+)(\\.xml$)", STAGE_FILE_NAME));
-    private static final int USER_ID_GROUP_INDEX_IN_PATTERN = 2;
-    private static final Duration STAGE_FILE_RETENTION_PERIOD = Duration.ofDays(3);
+    // 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);
 
     private final LocaleManagerService mLocaleManagerService;
     private final PackageManagerInternal mPackageManagerInternal;
-    private final File mStagedLocalesDir;
     private final Clock mClock;
     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 = new SparseArray<>();
+    private final SparseArray<StagedData> mStagedData;
 
     private final PackageMonitor mPackageMonitor;
     private final BroadcastReceiver mUserMonitor;
 
     LocaleManagerBackupHelper(LocaleManagerService localeManagerService,
             PackageManagerInternal pmInternal) {
-        this(localeManagerService.mContext, localeManagerService, pmInternal,
-                new File(Environment.getDataSystemCeDirectory(),
-                        "app_locales"), getDefaultClock());
+        this(localeManagerService.mContext, localeManagerService, pmInternal, getDefaultClock(),
+                new SparseArray<>());
     }
 
     private static @NonNull Clock getDefaultClock() {
@@ -123,14 +108,12 @@ class LocaleManagerBackupHelper {
 
     @VisibleForTesting LocaleManagerBackupHelper(Context context,
             LocaleManagerService localeManagerService,
-            PackageManagerInternal pmInternal, File stagedLocalesDir, Clock clock) {
+            PackageManagerInternal pmInternal, Clock clock, SparseArray<StagedData> stagedData) {
         mContext = context;
         mLocaleManagerService = localeManagerService;
         mPackageManagerInternal = pmInternal;
         mClock = clock;
-        mStagedLocalesDir = stagedLocalesDir;
-
-        loadAllStageFiles();
+        mStagedData = stagedData;
 
         HandlerThread broadcastHandlerThread = new HandlerThread(TAG,
                 Process.THREAD_PRIORITY_BACKGROUND);
@@ -158,67 +141,6 @@ class LocaleManagerBackupHelper {
     }
 
     /**
-     * Loads the staged data into memory by reading all the files in the staged directory.
-     *
-     * <p><b>Note:</b> We don't ned to hold the lock here because this is only called in the
-     * constructor (before any broadcast receivers are registered).
-     */
-    private void loadAllStageFiles() {
-        File[] files = mStagedLocalesDir.listFiles();
-        if (files == null) {
-            return;
-        }
-        for (File file : files) {
-            String fileName = file.getName();
-            Matcher matcher = STAGE_FILE_NAME_PATTERN.matcher(fileName);
-            if (!matcher.matches()) {
-                file.delete();
-                Slog.w(TAG, TextUtils.formatSimple("Deleted %s. Reason: %s.", fileName,
-                        "Unrecognized file"));
-                continue;
-            }
-            try {
-                final int userId = Integer.parseInt(matcher.group(USER_ID_GROUP_INDEX_IN_PATTERN));
-                StagedData stagedData = readStageFile(file);
-                if (stagedData != null) {
-                    mStagedData.put(userId, stagedData);
-                } else {
-                    file.delete();
-                    Slog.w(TAG, TextUtils.formatSimple("Deleted %s. Reason: %s.", fileName,
-                            "Could not read file"));
-                }
-            } catch (NumberFormatException e) {
-                file.delete();
-                Slog.w(TAG, TextUtils.formatSimple("Deleted %s. Reason: %s.", fileName,
-                        "Could not parse user id from file name"));
-            }
-        }
-    }
-
-    /**
-     * Loads the stage file from the disk and parses it into a list of app backups.
-     */
-    private @Nullable StagedData readStageFile(@NonNull File file) {
-        InputStream stagedDataInputStream = null;
-        AtomicFile stageFile = new AtomicFile(file);
-        try {
-            stagedDataInputStream = stageFile.openRead();
-            final TypedXmlPullParser parser = Xml.newFastPullParser();
-            parser.setInput(stagedDataInputStream, StandardCharsets.UTF_8.name());
-
-            XmlUtils.beginDocument(parser, LOCALES_XML_TAG);
-            long creationTimeMillis = parser.getAttributeLong(/* namespace= */ null,
-                    ATTR_CREATION_TIME_MILLIS);
-            return new StagedData(creationTimeMillis, readFromXml(parser));
-        } catch (IOException | XmlPullParserException e) {
-            Slog.e(TAG, "Could not parse stage file ", e);
-        } finally {
-            IoUtils.closeQuietly(stagedDataInputStream);
-        }
-        return null;
-    }
-
-    /**
      * @see LocaleManagerInternal#getBackupPayload(int userId)
      */
     public byte[] getBackupPayload(int userId) {
@@ -261,9 +183,7 @@ class LocaleManagerBackupHelper {
 
         final ByteArrayOutputStream out = new ByteArrayOutputStream();
         try {
-            // Passing arbitrary value for creationTimeMillis since it is ignored when forStage
-            // is false.
-            writeToXml(out, pkgStates, /* forStage= */ false, /* creationTimeMillis= */ -1);
+            writeToXml(out, pkgStates);
         } catch (IOException e) {
             Slog.e(TAG, "Could not write to xml for backup ", e);
             return null;
@@ -284,7 +204,7 @@ class LocaleManagerBackupHelper {
             int userId = mStagedData.keyAt(i);
             StagedData stagedData = mStagedData.get(userId);
             if (stagedData.mCreationTimeMillis
-                    < mClock.millis() - STAGE_FILE_RETENTION_PERIOD.toMillis()) {
+                    < mClock.millis() - STAGE_DATA_RETENTION_PERIOD.toMillis()) {
                 deleteStagedDataLocked(userId);
             }
         }
@@ -305,7 +225,7 @@ class LocaleManagerBackupHelper {
 
         final ByteArrayInputStream inputStream = new ByteArrayInputStream(payload);
 
-        HashMap<String, String> pkgStates = new HashMap<>();
+        HashMap<String, String> pkgStates;
         try {
             // Parse the input blob into a list of BackupPackageState.
             final TypedXmlPullParser parser = Xml.newFastPullParser();
@@ -315,6 +235,7 @@ class LocaleManagerBackupHelper {
             pkgStates = readFromXml(parser);
         } catch (IOException | XmlPullParserException e) {
             Slog.e(TAG, "Could not parse payload ", e);
+            return;
         }
 
         // We need a lock here to prevent race conditions when accessing the stage file.
@@ -323,7 +244,7 @@ class LocaleManagerBackupHelper {
         // performed simultaneously.
         synchronized (mStagedDataLock) {
             // Backups for apps which are yet to be installed.
-            mStagedData.put(userId, new StagedData(mClock.millis(), new HashMap<>()));
+            StagedData stagedData = new StagedData(mClock.millis(), new HashMap<>());
 
             for (String pkgName : pkgStates.keySet()) {
                 String languageTags = pkgStates.get(pkgName);
@@ -333,7 +254,7 @@ class LocaleManagerBackupHelper {
                     checkExistingLocalesAndApplyRestore(pkgName, languageTags, userId);
                 } else {
                     // Stage the data if the app isn't installed.
-                    mStagedData.get(userId).mPackageStates.put(pkgName, languageTags);
+                    stagedData.mPackageStates.put(pkgName, languageTags);
                     if (DEBUG) {
                         Slog.d(TAG, "Add locales=" + languageTags
                                 + " package=" + pkgName + " for lazy restore.");
@@ -341,7 +262,9 @@ class LocaleManagerBackupHelper {
                 }
             }
 
-            writeStageFileLocked(userId);
+            if (!stagedData.mPackageStates.isEmpty()) {
+                mStagedData.put(userId, stagedData);
+            }
         }
     }
 
@@ -396,55 +319,10 @@ class LocaleManagerBackupHelper {
         }
     }
 
-    /**
-     * Converts the list of app backups into xml and writes it onto the disk.
-     */
-    private void writeStageFileLocked(int userId) {
-        StagedData stagedData = mStagedData.get(userId);
-        if (stagedData.mPackageStates.isEmpty()) {
-            deleteStagedDataLocked(userId);
-            return;
-        }
-
-        final FileOutputStream stagedDataOutputStream;
-        AtomicFile stageFile = new AtomicFile(
-                new File(mStagedLocalesDir,
-                        TextUtils.formatSimple("%s_%d.xml", STAGE_FILE_NAME, userId)));
-        try {
-            stagedDataOutputStream = stageFile.startWrite();
-        } catch (IOException e) {
-            Slog.e(TAG, "Failed to save stage file");
-            return;
-        }
-
-        try {
-            writeToXml(stagedDataOutputStream, stagedData.mPackageStates,  /* forStage= */ true,
-                    stagedData.mCreationTimeMillis);
-            stageFile.finishWrite(stagedDataOutputStream);
-            if (DEBUG) {
-                Slog.d(TAG, "Stage file written.");
-            }
-        } catch (IOException e) {
-            Slog.e(TAG, "Could not write stage file", e);
-            stageFile.failWrite(stagedDataOutputStream);
-        }
-    }
-
     private void deleteStagedDataLocked(@UserIdInt int userId) {
-        AtomicFile stageFile = getStageFileIfExistsLocked(userId);
-        if (stageFile != null) {
-            stageFile.delete();
-        }
         mStagedData.remove(userId);
     }
 
-    private @Nullable AtomicFile getStageFileIfExistsLocked(@UserIdInt int userId) {
-        final File stageFile = new File(mStagedLocalesDir,
-                TextUtils.formatSimple("%s_%d.xml", STAGE_FILE_NAME, userId));
-        return stageFile.isFile() ? new AtomicFile(stageFile)
-                : null;
-    }
-
     /**
      * Parses the backup data from the serialized xml input stream.
      */
@@ -468,15 +346,8 @@ class LocaleManagerBackupHelper {
 
     /**
      * Converts the list of app backup data into a serialized xml stream.
-     *
-     * @param forStage Flag to indicate whether this method is called for the purpose of
-     * staging the data. Note that if this is false, {@code creationTimeMillis} is ignored because
-     * we only need it for the stage data.
-     * @param creationTimeMillis The timestamp when the stage data was created. This is required
-     * to determine when to delete the stage data.
      */
-    private static void writeToXml(OutputStream stream,
-            @NonNull HashMap<String, String> pkgStates, boolean forStage, long creationTimeMillis)
+    private static void writeToXml(OutputStream stream, @NonNull HashMap<String, String> pkgStates)
             throws IOException {
         if (pkgStates.isEmpty()) {
             // No need to write anything at all if pkgStates is empty.
@@ -488,11 +359,6 @@ class LocaleManagerBackupHelper {
         out.startDocument(/* encoding= */ null, /* standalone= */ true);
         out.startTag(/* namespace= */ null, LOCALES_XML_TAG);
 
-        if (forStage) {
-            out.attribute(/* namespace= */ null, ATTR_CREATION_TIME_MILLIS,
-                    Long.toString(creationTimeMillis));
-        }
-
         for (String pkg : pkgStates.keySet()) {
             out.startTag(/* namespace= */ null, PACKAGE_XML_TAG);
             out.attribute(/* namespace= */ null, ATTR_PACKAGE_NAME, pkg);
@@ -504,7 +370,7 @@ class LocaleManagerBackupHelper {
         out.endDocument();
     }
 
-    private static class StagedData {
+    static class StagedData {
         final long mCreationTimeMillis;
         final HashMap<String, String> mPackageStates;
 
@@ -517,7 +383,7 @@ class LocaleManagerBackupHelper {
     /**
      * Broadcast listener to capture user removed event.
      *
-     * <p>The stage file is deleted when a user is removed.
+     * <p>The stage data is deleted when a user is removed.
      */
     private final class UserMonitor extends BroadcastReceiver {
         @Override
@@ -546,6 +412,8 @@ class LocaleManagerBackupHelper {
         public void onPackageAdded(String packageName, int uid) {
             try {
                 synchronized (mStagedDataLock) {
+                    cleanStagedDataForOldEntriesLocked();
+
                     int userId = UserHandle.getUserId(uid);
                     if (mStagedData.contains(userId)) {
                         // Perform lazy restore only if the staged data exists.
@@ -589,7 +457,7 @@ class LocaleManagerBackupHelper {
         // Check if the package is installed indeed
         if (!isPackageInstalledForUser(packageName, userId)) {
             Slog.e(TAG, packageName + " not installed for user " + userId
-                    + ". Could not restore locales from stage file");
+                    + ". Could not restore locales from stage data");
             return;
         }
 
@@ -603,8 +471,11 @@ class LocaleManagerBackupHelper {
 
                 // Remove the restored entry from the staged data list.
                 stagedData.mPackageStates.remove(pkgName);
-                // Update the file on the disk.
-                writeStageFileLocked(userId);
+
+                // 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.
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 70e78eb41999..c77100045118 100644
--- a/services/tests/servicestests/src/com/android/server/locales/LocaleManagerBackupRestoreTest.java
+++ b/services/tests/servicestests/src/com/android/server/locales/LocaleManagerBackupRestoreTest.java
@@ -17,9 +17,7 @@
 package com.android.server.locales;
 
 import static junit.framework.Assert.assertEquals;
-import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertNull;
-import static junit.framework.Assert.assertTrue;
 
 import static org.junit.Assert.assertNotNull;
 import static org.mockito.ArgumentMatchers.any;
@@ -43,11 +41,10 @@ import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManagerInternal;
 import android.os.Binder;
-import android.os.Environment;
 import android.os.LocaleList;
 import android.os.RemoteException;
 import android.os.SimpleClock;
-import android.util.AtomicFile;
+import android.util.SparseArray;
 import android.util.TypedXmlPullParser;
 import android.util.TypedXmlSerializer;
 import android.util.Xml;
@@ -57,6 +54,7 @@ import androidx.test.ext.junit.runners.AndroidJUnit4;
 import com.android.internal.content.PackageMonitor;
 import com.android.internal.util.XmlUtils;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -65,10 +63,7 @@ import org.xmlpull.v1.XmlPullParserException;
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
-import java.io.File;
-import java.io.FileOutputStream;
 import java.io.IOException;
-import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.charset.StandardCharsets;
 import java.time.Clock;
@@ -95,9 +90,8 @@ public class LocaleManagerBackupRestoreTest {
             LocaleList.forLanguageTags(DEFAULT_LOCALE_TAGS);
     private static final Map<String, String> DEFAULT_PACKAGE_LOCALES_MAP = Map.of(
             DEFAULT_PACKAGE_NAME, DEFAULT_LOCALE_TAGS);
-    private static final File STAGED_LOCALES_DIR = new File(
-            Environment.getExternalStorageDirectory(), "lmsUnitTests");
-
+    private static final SparseArray<LocaleManagerBackupHelper.StagedData> STAGE_DATA =
+            new SparseArray<>();
 
     private LocaleManagerBackupHelper mBackupHelper;
     private long mCurrentTimeMillis;
@@ -138,14 +132,17 @@ public class LocaleManagerBackupRestoreTest {
         doReturn(mMockPackageManager).when(mMockContext).getPackageManager();
 
         mBackupHelper = spy(new ShadowLocaleManagerBackupHelper(mMockContext,
-                mMockLocaleManagerService, mMockPackageManagerInternal,
-                new File(Environment.getExternalStorageDirectory(), "lmsUnitTests"), mClock));
+                mMockLocaleManagerService, mMockPackageManagerInternal, mClock, STAGE_DATA));
         doNothing().when(mBackupHelper).notifyBackupManager();
 
         mUserMonitor = mBackupHelper.getUserMonitor();
         mPackageMonitor = mBackupHelper.getPackageMonitor();
         setCurrentTimeMillis(DEFAULT_CREATION_TIME_MILLIS);
-        cleanStagedFiles();
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        STAGE_DATA.clear();
     }
 
     @Test
@@ -203,25 +200,25 @@ public class LocaleManagerBackupRestoreTest {
     }
 
     @Test
-    public void testRestore_nullPayload_nothingRestoredAndNoStageFile() throws Exception {
+    public void testRestore_nullPayload_nothingRestoredAndNoStageData() throws Exception {
         mBackupHelper.stageAndApplyRestoredPayload(/* payload= */ null, DEFAULT_USER_ID);
 
         verifyNothingRestored();
-        checkStageFileDoesNotExist(DEFAULT_USER_ID);
+        checkStageDataDoesNotExist(DEFAULT_USER_ID);
     }
 
     @Test
-    public void testRestore_zeroLengthPayload_nothingRestoredAndNoStageFile() throws Exception {
+    public void testRestore_zeroLengthPayload_nothingRestoredAndNoStageData() throws Exception {
         final ByteArrayOutputStream out = new ByteArrayOutputStream();
         mBackupHelper.stageAndApplyRestoredPayload(/* payload= */ out.toByteArray(),
                 DEFAULT_USER_ID);
 
         verifyNothingRestored();
-        checkStageFileDoesNotExist(DEFAULT_USER_ID);
+        checkStageDataDoesNotExist(DEFAULT_USER_ID);
     }
 
     @Test
-    public void testRestore_allAppsInstalled_noStageFileCreated() throws Exception {
+    public void testRestore_allAppsInstalled_noStageDataCreated() throws Exception {
         final ByteArrayOutputStream out = new ByteArrayOutputStream();
         writeTestPayload(out, DEFAULT_PACKAGE_LOCALES_MAP);
 
@@ -234,8 +231,7 @@ public class LocaleManagerBackupRestoreTest {
         verify(mMockLocaleManagerService, times(1)).setApplicationLocales(DEFAULT_PACKAGE_NAME,
                 DEFAULT_USER_ID, DEFAULT_LOCALES);
 
-        // Stage file wasn't created.
-        checkStageFileDoesNotExist(DEFAULT_USER_ID);
+        checkStageDataDoesNotExist(DEFAULT_USER_ID);
     }
 
     @Test
@@ -248,8 +244,8 @@ public class LocaleManagerBackupRestoreTest {
         mBackupHelper.stageAndApplyRestoredPayload(out.toByteArray(), DEFAULT_USER_ID);
 
         verifyNothingRestored();
-        verifyStageFileContent(DEFAULT_PACKAGE_LOCALES_MAP,
-                getStageFileIfExists(DEFAULT_USER_ID), DEFAULT_CREATION_TIME_MILLIS);
+        verifyStageDataForUser(DEFAULT_PACKAGE_LOCALES_MAP,
+                DEFAULT_CREATION_TIME_MILLIS, DEFAULT_USER_ID);
     }
 
     @Test
@@ -257,8 +253,10 @@ public class LocaleManagerBackupRestoreTest {
         final ByteArrayOutputStream out = new ByteArrayOutputStream();
         HashMap<String, String> pkgLocalesMap = new HashMap<>();
 
-        String pkgNameA = "com.android.myAppA", pkgNameB = "com.android.myAppB";
-        String langTagsA = "ru", langTagsB = "hi,fr";
+        String pkgNameA = "com.android.myAppA";
+        String pkgNameB = "com.android.myAppB";
+        String langTagsA = "ru";
+        String langTagsB = "hi,fr";
         pkgLocalesMap.put(pkgNameA, langTagsA);
         pkgLocalesMap.put(pkgNameB, langTagsB);
         writeTestPayload(out, pkgLocalesMap);
@@ -273,12 +271,12 @@ public class LocaleManagerBackupRestoreTest {
                 LocaleList.forLanguageTags(langTagsA));
 
         pkgLocalesMap.remove(pkgNameA);
-        verifyStageFileContent(pkgLocalesMap, getStageFileIfExists(DEFAULT_USER_ID),
-                DEFAULT_CREATION_TIME_MILLIS);
+        verifyStageDataForUser(pkgLocalesMap,
+                DEFAULT_CREATION_TIME_MILLIS, DEFAULT_USER_ID);
     }
 
     @Test
-    public void testRestore_appLocalesAlreadySet_nothingRestoredAndNoStageFile() throws Exception {
+    public void testRestore_appLocalesAlreadySet_nothingRestoredAndNoStageData() throws Exception {
         final ByteArrayOutputStream out = new ByteArrayOutputStream();
         writeTestPayload(out, DEFAULT_PACKAGE_LOCALES_MAP);
 
@@ -289,8 +287,7 @@ public class LocaleManagerBackupRestoreTest {
 
         // Since locales are already set, we should not restore anything for it.
         verifyNothingRestored();
-        // Stage file wasn't created
-        checkStageFileDoesNotExist(DEFAULT_USER_ID);
+        checkStageDataDoesNotExist(DEFAULT_USER_ID);
     }
 
     @Test
@@ -299,9 +296,12 @@ public class LocaleManagerBackupRestoreTest {
         final ByteArrayOutputStream out = new ByteArrayOutputStream();
         HashMap<String, String> pkgLocalesMap = new HashMap<>();
 
-        String pkgNameA = "com.android.myAppA", pkgNameB = "com.android.myAppB", pkgNameC =
-                "com.android.myAppC";
-        String langTagsA = "ru", langTagsB = "hi,fr", langTagsC = "zh,es";
+        String pkgNameA = "com.android.myAppA";
+        String pkgNameB = "com.android.myAppB";
+        String pkgNameC = "com.android.myAppC";
+        String langTagsA = "ru";
+        String langTagsB = "hi,fr";
+        String langTagsC = "zh,es";
         pkgLocalesMap.put(pkgNameA, langTagsA);
         pkgLocalesMap.put(pkgNameB, langTagsB);
         pkgLocalesMap.put(pkgNameC, langTagsC);
@@ -328,8 +328,8 @@ public class LocaleManagerBackupRestoreTest {
         // App C is staged.
         pkgLocalesMap.remove(pkgNameA);
         pkgLocalesMap.remove(pkgNameB);
-        verifyStageFileContent(pkgLocalesMap, getStageFileIfExists(DEFAULT_USER_ID),
-                DEFAULT_CREATION_TIME_MILLIS);
+        verifyStageDataForUser(pkgLocalesMap,
+                DEFAULT_CREATION_TIME_MILLIS, DEFAULT_USER_ID);
     }
 
     @Test
@@ -341,15 +341,15 @@ public class LocaleManagerBackupRestoreTest {
 
         mBackupHelper.stageAndApplyRestoredPayload(out.toByteArray(), DEFAULT_USER_ID);
 
-        verifyStageFileContent(DEFAULT_PACKAGE_LOCALES_MAP, getStageFileIfExists(DEFAULT_USER_ID),
-                DEFAULT_CREATION_TIME_MILLIS);
+        verifyStageDataForUser(DEFAULT_PACKAGE_LOCALES_MAP,
+                DEFAULT_CREATION_TIME_MILLIS, DEFAULT_USER_ID);
 
         final long newCreationTime = DEFAULT_CREATION_TIME_MILLIS + 100;
         setCurrentTimeMillis(newCreationTime);
         mBackupHelper.stageAndApplyRestoredPayload(out.toByteArray(), DEFAULT_USER_ID);
 
-        verifyStageFileContent(DEFAULT_PACKAGE_LOCALES_MAP, getStageFileIfExists(DEFAULT_USER_ID),
-                newCreationTime);
+        verifyStageDataForUser(DEFAULT_PACKAGE_LOCALES_MAP,
+                newCreationTime, DEFAULT_USER_ID);
     }
 
     @Test
@@ -357,8 +357,10 @@ public class LocaleManagerBackupRestoreTest {
         final ByteArrayOutputStream out = new ByteArrayOutputStream();
         HashMap<String, String> pkgLocalesMap = new HashMap<>();
 
-        String pkgNameA = "com.android.myAppA", pkgNameB = "com.android.myAppB";
-        String langTagsA = "ru", langTagsB = "hi,fr";
+        String pkgNameA = "com.android.myAppA";
+        String pkgNameB = "com.android.myAppB";
+        String langTagsA = "ru";
+        String langTagsB = "hi,fr";
         pkgLocalesMap.put(pkgNameA, langTagsA);
         pkgLocalesMap.put(pkgNameB, langTagsB);
         writeTestPayload(out, pkgLocalesMap);
@@ -380,8 +382,7 @@ public class LocaleManagerBackupRestoreTest {
                 LocaleList.forLanguageTags(langTagsA));
 
         pkgLocalesMap.remove(pkgNameA);
-        verifyStageFileContent(pkgLocalesMap, getStageFileIfExists(DEFAULT_USER_ID),
-                DEFAULT_CREATION_TIME_MILLIS);
+        verifyStageDataForUser(pkgLocalesMap, DEFAULT_CREATION_TIME_MILLIS, DEFAULT_USER_ID);
 
         setUpPackageInstalled(pkgNameB);
 
@@ -389,7 +390,7 @@ public class LocaleManagerBackupRestoreTest {
 
         verify(mMockLocaleManagerService, times(1)).setApplicationLocales(pkgNameB, DEFAULT_USER_ID,
                 LocaleList.forLanguageTags(langTagsB));
-        checkStageFileDoesNotExist(DEFAULT_USER_ID);
+        checkStageDataDoesNotExist(DEFAULT_USER_ID);
     }
 
     @Test
@@ -404,8 +405,8 @@ public class LocaleManagerBackupRestoreTest {
         mBackupHelper.stageAndApplyRestoredPayload(out.toByteArray(), DEFAULT_USER_ID);
 
         verifyNothingRestored();
-        verifyStageFileContent(DEFAULT_PACKAGE_LOCALES_MAP, getStageFileIfExists(DEFAULT_USER_ID),
-                DEFAULT_CREATION_TIME_MILLIS);
+        verifyStageDataForUser(DEFAULT_PACKAGE_LOCALES_MAP,
+                DEFAULT_CREATION_TIME_MILLIS, DEFAULT_USER_ID);
 
         // App is installed later (post SUW).
         setUpPackageInstalled(DEFAULT_PACKAGE_NAME);
@@ -415,11 +416,11 @@ public class LocaleManagerBackupRestoreTest {
 
         // Since locales are already set, we should not restore anything for it.
         verifyNothingRestored();
-        checkStageFileDoesNotExist(DEFAULT_USER_ID);
+        checkStageDataDoesNotExist(DEFAULT_USER_ID);
     }
 
     @Test
-    public void testStageFileDeletion_backupPassRunAfterRetentionPeriod_stageFileDeleted()
+    public void testStageDataDeletion_backupPassRunAfterRetentionPeriod_stageDataDeleted()
             throws Exception {
         final ByteArrayOutputStream out = new ByteArrayOutputStream();
         writeTestPayload(out, DEFAULT_PACKAGE_LOCALES_MAP);
@@ -429,8 +430,8 @@ public class LocaleManagerBackupRestoreTest {
         mBackupHelper.stageAndApplyRestoredPayload(out.toByteArray(), DEFAULT_USER_ID);
 
         verifyNothingRestored();
-        verifyStageFileContent(DEFAULT_PACKAGE_LOCALES_MAP, getStageFileIfExists(DEFAULT_USER_ID),
-                DEFAULT_CREATION_TIME_MILLIS);
+        verifyStageDataForUser(DEFAULT_PACKAGE_LOCALES_MAP,
+                DEFAULT_CREATION_TIME_MILLIS, DEFAULT_USER_ID);
 
         // Retention period has not elapsed.
         setCurrentTimeMillis(
@@ -439,32 +440,78 @@ public class LocaleManagerBackupRestoreTest {
                 .getInstalledApplications(anyLong(), anyInt(), anyInt());
         assertNull(mBackupHelper.getBackupPayload(DEFAULT_USER_ID));
 
-        // Stage file should NOT be deleted.
-        checkStageFileExists(DEFAULT_USER_ID);
+        checkStageDataExists(DEFAULT_USER_ID);
 
-        // Exactly RETENTION_PERIOD amount of time has passed so stage file should still not be
+        // Exactly RETENTION_PERIOD amount of time has passed so stage data should still not be
         // removed.
         setCurrentTimeMillis(DEFAULT_CREATION_TIME_MILLIS + RETENTION_PERIOD.toMillis());
         doReturn(List.of()).when(mMockPackageManagerInternal)
                 .getInstalledApplications(anyLong(), anyInt(), anyInt());
         assertNull(mBackupHelper.getBackupPayload(DEFAULT_USER_ID));
 
-        // Stage file should NOT be deleted.
-        checkStageFileExists(DEFAULT_USER_ID);
+        checkStageDataExists(DEFAULT_USER_ID);
 
-        // Retention period has now expired, stage file should be deleted.
+        // Retention period has now expired, stage data should be deleted.
         setCurrentTimeMillis(
                 DEFAULT_CREATION_TIME_MILLIS + RETENTION_PERIOD.plusSeconds(1).toMillis());
         doReturn(List.of()).when(mMockPackageManagerInternal)
                 .getInstalledApplications(anyLong(), anyInt(), anyInt());
         assertNull(mBackupHelper.getBackupPayload(DEFAULT_USER_ID));
 
-        // Stage file should be deleted.
-        checkStageFileDoesNotExist(DEFAULT_USER_ID);
+        checkStageDataDoesNotExist(DEFAULT_USER_ID);
     }
 
     @Test
-    public void testUserRemoval_userRemoved_stageFileDeleted() throws Exception {
+    public void testStageDataDeletion_lazyRestoreAfterRetentionPeriod_stageDataDeleted()
+            throws Exception {
+        final ByteArrayOutputStream out = new ByteArrayOutputStream();
+        HashMap<String, String> pkgLocalesMap = new HashMap<>();
+
+        String pkgNameA = "com.android.myAppA";
+        String pkgNameB = "com.android.myAppB";
+        String langTagsA = "ru";
+        String langTagsB = "hi,fr";
+        pkgLocalesMap.put(pkgNameA, langTagsA);
+        pkgLocalesMap.put(pkgNameB, langTagsB);
+        writeTestPayload(out, pkgLocalesMap);
+
+        setUpPackageNotInstalled(pkgNameA);
+        setUpPackageNotInstalled(pkgNameB);
+        setUpLocalesForPackage(pkgNameA, LocaleList.getEmptyLocaleList());
+        setUpLocalesForPackage(pkgNameB, LocaleList.getEmptyLocaleList());
+
+        mBackupHelper.stageAndApplyRestoredPayload(out.toByteArray(), DEFAULT_USER_ID);
+
+        verifyNothingRestored();
+        verifyStageDataForUser(pkgLocalesMap, DEFAULT_CREATION_TIME_MILLIS, DEFAULT_USER_ID);
+
+        // Retention period has not elapsed.
+        setCurrentTimeMillis(
+                DEFAULT_CREATION_TIME_MILLIS + RETENTION_PERIOD.minusHours(1).toMillis());
+
+        setUpPackageInstalled(pkgNameA);
+        mPackageMonitor.onPackageAdded(pkgNameA, DEFAULT_UID);
+
+        verify(mMockLocaleManagerService, times(1)).setApplicationLocales(pkgNameA, DEFAULT_USER_ID,
+                LocaleList.forLanguageTags(langTagsA));
+
+        pkgLocalesMap.remove(pkgNameA);
+        verifyStageDataForUser(pkgLocalesMap, DEFAULT_CREATION_TIME_MILLIS, DEFAULT_USER_ID);
+
+        // Retention period has now expired, stage data should be deleted.
+        setCurrentTimeMillis(
+                DEFAULT_CREATION_TIME_MILLIS + RETENTION_PERIOD.plusSeconds(1).toMillis());
+        setUpPackageInstalled(pkgNameB);
+        mPackageMonitor.onPackageAdded(pkgNameB, DEFAULT_UID);
+
+        verify(mMockLocaleManagerService, times(0)).setApplicationLocales(eq(pkgNameB), anyInt(),
+                any());
+
+        checkStageDataDoesNotExist(DEFAULT_USER_ID);
+    }
+
+    @Test
+    public void testUserRemoval_userRemoved_stageDataDeleted() throws Exception {
         final ByteArrayOutputStream outDefault = new ByteArrayOutputStream();
         writeTestPayload(outDefault, DEFAULT_PACKAGE_LOCALES_MAP);
 
@@ -485,119 +532,20 @@ public class LocaleManagerBackupRestoreTest {
 
         verifyNothingRestored();
 
-        // Verify stage file contents.
-        AtomicFile stageFileDefaultUser = getStageFileIfExists(DEFAULT_USER_ID);
-        verifyStageFileContent(DEFAULT_PACKAGE_LOCALES_MAP, stageFileDefaultUser,
-                DEFAULT_CREATION_TIME_MILLIS);
-
-        AtomicFile stageFileWorkProfile = getStageFileIfExists(WORK_PROFILE_USER_ID);
-        verifyStageFileContent(pkgLocalesMapWorkProfile, stageFileWorkProfile,
-                DEFAULT_CREATION_TIME_MILLIS);
+        verifyStageDataForUser(DEFAULT_PACKAGE_LOCALES_MAP,
+                DEFAULT_CREATION_TIME_MILLIS, DEFAULT_USER_ID);
+        verifyStageDataForUser(pkgLocalesMapWorkProfile,
+                DEFAULT_CREATION_TIME_MILLIS, WORK_PROFILE_USER_ID);
 
         Intent intent = new Intent();
         intent.setAction(Intent.ACTION_USER_REMOVED);
         intent.putExtra(Intent.EXTRA_USER_HANDLE, DEFAULT_USER_ID);
         mUserMonitor.onReceive(mMockContext, intent);
 
-        // Stage file should be removed only for DEFAULT_USER_ID.
-        checkStageFileDoesNotExist(DEFAULT_USER_ID);
-        verifyStageFileContent(pkgLocalesMapWorkProfile, stageFileWorkProfile,
-                DEFAULT_CREATION_TIME_MILLIS);
-    }
-
-    @Test
-    public void testLoadStageFiles_invalidNameFormat_stageFileDeleted() throws Exception {
-        // Stage file name should be : staged_locales_<user_id_int>.xml
-        File stageFile = new File(STAGED_LOCALES_DIR, "xyz.xml");
-        assertTrue(stageFile.createNewFile());
-        assertTrue(stageFile.isFile());
-
-        // Putting valid xml data in file.
-        FileOutputStream out = new FileOutputStream(stageFile);
-        writeTestPayload(out, DEFAULT_PACKAGE_LOCALES_MAP, /* forStage= */
-                true, /* creationTimeMillis= */ 0);
-        out.flush();
-        out.close();
-
-        verifyStageFileContent(DEFAULT_PACKAGE_LOCALES_MAP,
-                new AtomicFile(stageFile), /* creationTimeMillis= */ 0);
-
-        mBackupHelper = new LocaleManagerBackupHelper(mMockContext, mMockLocaleManagerService,
-                mMockPackageManagerInternal, STAGED_LOCALES_DIR, mClock);
-        assertFalse(stageFile.isFile());
-    }
-
-    @Test
-    public void testLoadStageFiles_userIdNotParseable_stageFileDeleted() throws Exception {
-        // Stage file name should be : staged_locales_<user_id_int>.xml
-        File stageFile = new File(STAGED_LOCALES_DIR, "staged_locales_abc.xml");
-        assertTrue(stageFile.createNewFile());
-        assertTrue(stageFile.isFile());
-
-        // Putting valid xml data in file.
-        FileOutputStream out = new FileOutputStream(stageFile);
-        writeTestPayload(out, DEFAULT_PACKAGE_LOCALES_MAP, /* forStage= */
-                true, /* creationTimeMillis= */ 0);
-        out.flush();
-        out.close();
-
-        verifyStageFileContent(DEFAULT_PACKAGE_LOCALES_MAP,
-                new AtomicFile(stageFile), /* creationTimeMillis= */ 0);
-
-        mBackupHelper = new LocaleManagerBackupHelper(mMockContext, mMockLocaleManagerService,
-                mMockPackageManagerInternal, STAGED_LOCALES_DIR, mClock);
-        assertFalse(stageFile.isFile());
-    }
-
-    @Test
-    public void testLoadStageFiles_invalidContent_stageFileDeleted() throws Exception {
-        File stageFile = new File(STAGED_LOCALES_DIR, "staged_locales_0.xml");
-        assertTrue(stageFile.createNewFile());
-        assertTrue(stageFile.isFile());
-
-        FileOutputStream out = new FileOutputStream(stageFile);
-        out.write("some_non_xml_string".getBytes());
-        out.close();
-
-        mBackupHelper = new LocaleManagerBackupHelper(mMockContext, mMockLocaleManagerService,
-                mMockPackageManagerInternal, STAGED_LOCALES_DIR, mClock);
-        assertFalse(stageFile.isFile());
-    }
-
-    @Test
-    public void testLoadStageFiles_validContent_doesLazyRestore() throws Exception {
-        File stageFile = new File(STAGED_LOCALES_DIR, "staged_locales_0.xml");
-        assertTrue(stageFile.createNewFile());
-        assertTrue(stageFile.isFile());
-
-        // Putting valid xml data in file.
-        FileOutputStream out = new FileOutputStream(stageFile);
-        writeTestPayload(out, DEFAULT_PACKAGE_LOCALES_MAP, /* forStage= */
-                true, DEFAULT_CREATION_TIME_MILLIS);
-        out.flush();
-        out.close();
-
-        verifyStageFileContent(DEFAULT_PACKAGE_LOCALES_MAP,
-                new AtomicFile(stageFile), DEFAULT_CREATION_TIME_MILLIS);
-
-        mBackupHelper = new LocaleManagerBackupHelper(mMockContext, mMockLocaleManagerService,
-                mMockPackageManagerInternal, STAGED_LOCALES_DIR, mClock);
-        mPackageMonitor = mBackupHelper.getPackageMonitor();
-
-        // Stage file still exists.
-        assertTrue(stageFile.isFile());
-
-        // App is installed later.
-        setUpPackageInstalled(DEFAULT_PACKAGE_NAME);
-        setUpLocalesForPackage(DEFAULT_PACKAGE_NAME, LocaleList.getEmptyLocaleList());
-
-        mPackageMonitor.onPackageAdded(DEFAULT_PACKAGE_NAME, DEFAULT_UID);
-
-        verify(mMockLocaleManagerService, times(1)).setApplicationLocales(DEFAULT_PACKAGE_NAME,
-                DEFAULT_USER_ID, DEFAULT_LOCALES);
-
-        // Stage file gets deleted here because all staged locales have been applied.
-        assertFalse(stageFile.isFile());
+        // Stage data should be removed only for DEFAULT_USER_ID.
+        checkStageDataDoesNotExist(DEFAULT_USER_ID);
+        verifyStageDataForUser(pkgLocalesMapWorkProfile,
+                DEFAULT_CREATION_TIME_MILLIS, WORK_PROFILE_USER_ID);
     }
 
     private void setUpPackageInstalled(String packageName) throws Exception {
@@ -633,27 +581,15 @@ public class LocaleManagerBackupRestoreTest {
                 any());
     }
 
-
     private static void verifyPayloadForAppLocales(Map<String, String> expectedPkgLocalesMap,
             byte[] payload)
             throws IOException, XmlPullParserException {
-        verifyPayloadForAppLocales(expectedPkgLocalesMap, payload, /* forStage= */ false, -1);
-    }
-
-    private static void verifyPayloadForAppLocales(Map<String, String> expectedPkgLocalesMap,
-            byte[] payload, boolean forStage, long expectedCreationTime)
-            throws IOException, XmlPullParserException {
         final ByteArrayInputStream stream = new ByteArrayInputStream(payload);
         final TypedXmlPullParser parser = Xml.newFastPullParser();
         parser.setInput(stream, StandardCharsets.UTF_8.name());
 
         Map<String, String> backupDataMap = new HashMap<>();
         XmlUtils.beginDocument(parser, TEST_LOCALES_XML_TAG);
-        if (forStage) {
-            long actualCreationTime = parser.getAttributeLong(/* namespace= */ null,
-                    "creationTimeMillis");
-            assertEquals(expectedCreationTime, actualCreationTime);
-        }
         int depth = parser.getDepth();
         while (XmlUtils.nextElementWithin(parser, depth)) {
             if (parser.getName().equals("package")) {
@@ -668,13 +604,6 @@ public class LocaleManagerBackupRestoreTest {
 
     private static void writeTestPayload(OutputStream stream, Map<String, String> pkgLocalesMap)
             throws IOException {
-        writeTestPayload(stream, pkgLocalesMap, /* forStage= */ false, /* creationTimeMillis= */
-                -1);
-    }
-
-    private static void writeTestPayload(OutputStream stream, Map<String, String> pkgLocalesMap,
-            boolean forStage, long creationTimeMillis)
-            throws IOException {
         if (pkgLocalesMap.isEmpty()) {
             return;
         }
@@ -684,11 +613,6 @@ public class LocaleManagerBackupRestoreTest {
         out.startDocument(/* encoding= */ null, /* standalone= */ true);
         out.startTag(/* namespace= */ null, TEST_LOCALES_XML_TAG);
 
-        if (forStage) {
-            out.attribute(/* namespace= */ null, "creationTimeMillis",
-                    Long.toString(creationTimeMillis));
-        }
-
         for (String pkg : pkgLocalesMap.keySet()) {
             out.startTag(/* namespace= */ null, "package");
             out.attribute(/* namespace= */ null, "name", pkg);
@@ -700,41 +624,19 @@ public class LocaleManagerBackupRestoreTest {
         out.endDocument();
     }
 
-    private static void verifyStageFileContent(Map<String, String> expectedPkgLocalesMap,
-            AtomicFile stageFile,
-            long creationTimeMillis)
-            throws Exception {
-        assertNotNull(stageFile);
-        try (InputStream stagedDataInputStream = stageFile.openRead()) {
-            verifyPayloadForAppLocales(expectedPkgLocalesMap, stagedDataInputStream.readAllBytes(),
-                    /* forStage= */ true, creationTimeMillis);
-        } catch (IOException | XmlPullParserException e) {
-            throw e;
-        }
-    }
-
-    private static void checkStageFileDoesNotExist(int userId) {
-        assertNull(getStageFileIfExists(userId));
+    private void verifyStageDataForUser(Map<String, String> expectedPkgLocalesMap,
+            long expectedCreationTimeMillis, int userId) {
+        LocaleManagerBackupHelper.StagedData stagedDataForUser = STAGE_DATA.get(userId);
+        assertNotNull(stagedDataForUser);
+        assertEquals(expectedCreationTimeMillis, stagedDataForUser.mCreationTimeMillis);
+        assertEquals(expectedPkgLocalesMap, stagedDataForUser.mPackageStates);
     }
 
-    private static void checkStageFileExists(int userId) {
-        assertNotNull(getStageFileIfExists(userId));
+    private static void checkStageDataExists(int userId) {
+        assertNotNull(STAGE_DATA.get(userId));
     }
 
-    private static AtomicFile getStageFileIfExists(int userId) {
-        File file = new File(STAGED_LOCALES_DIR, String.format("staged_locales_%d.xml", userId));
-        if (file.isFile()) {
-            return new AtomicFile(file);
-        }
-        return null;
-    }
-
-    private static void cleanStagedFiles() {
-        File[] files = STAGED_LOCALES_DIR.listFiles();
-        if (files != null) {
-            for (File f : files) {
-                f.delete();
-            }
-        }
+    private static void checkStageDataDoesNotExist(int userId) {
+        assertNull(STAGE_DATA.get(userId));
     }
 }
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 93972c32a50f..b0fc6363b701 100644
--- a/services/tests/servicestests/src/com/android/server/locales/ShadowLocaleManagerBackupHelper.java
+++ b/services/tests/servicestests/src/com/android/server/locales/ShadowLocaleManagerBackupHelper.java
@@ -18,8 +18,8 @@ package com.android.server.locales;
 
 import android.content.Context;
 import android.content.pm.PackageManagerInternal;
+import android.util.SparseArray;
 
-import java.io.File;
 import java.time.Clock;
 
 /**
@@ -30,7 +30,8 @@ import java.time.Clock;
 public class ShadowLocaleManagerBackupHelper extends LocaleManagerBackupHelper {
     ShadowLocaleManagerBackupHelper(Context context,
             LocaleManagerService localeManagerService,
-            PackageManagerInternal pmInternal, File stagedLocalesDir, Clock clock) {
-        super(context, localeManagerService, pmInternal, stagedLocalesDir, clock);
+            PackageManagerInternal pmInternal, Clock clock,
+            SparseArray<LocaleManagerBackupHelper.StagedData> stagedData) {
+        super(context, localeManagerService, pmInternal, clock, stagedData);
     }
 }