Merge "AudioService: fix volume code path with a null audio device" into main
diff --git a/AconfigFlags.bp b/AconfigFlags.bp
index c80f2ea..b3e8ea8 100644
--- a/AconfigFlags.bp
+++ b/AconfigFlags.bp
@@ -55,6 +55,7 @@
     ":android.app.flags-aconfig-java{.generated_srcjars}",
     ":android.credentials.flags-aconfig-java{.generated_srcjars}",
     ":android.view.contentprotection.flags-aconfig-java{.generated_srcjars}",
+    ":com.android.server.flags.pinner-aconfig-java{.generated_srcjars}",
     ":android.service.voice.flags-aconfig-java{.generated_srcjars}",
     ":android.media.tv.flags-aconfig-java{.generated_srcjars}",
     ":android.service.autofill.flags-aconfig-java{.generated_srcjars}",
@@ -584,6 +585,19 @@
     defaults: ["framework-minus-apex-aconfig-java-defaults"],
 }
 
+// Pinner Service
+aconfig_declarations {
+    name: "com.android.server.flags.pinner-aconfig",
+    package: "com.android.server.flags",
+    srcs: ["services/core/java/com/android/server/flags/pinner.aconfig"],
+}
+
+java_aconfig_library {
+    name: "com.android.server.flags.pinner-aconfig-java",
+    aconfig_declarations: "com.android.server.flags.pinner-aconfig",
+    defaults: ["framework-minus-apex-aconfig-java-defaults"],
+}
+
 // Voice
 aconfig_declarations {
     name: "android.service.voice.flags-aconfig",
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index ce5752f..fc23f9b 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -258,6 +258,7 @@
     field public static final String PERFORM_IMS_SINGLE_REGISTRATION = "android.permission.PERFORM_IMS_SINGLE_REGISTRATION";
     field public static final String PERFORM_SIM_ACTIVATION = "android.permission.PERFORM_SIM_ACTIVATION";
     field public static final String POWER_SAVER = "android.permission.POWER_SAVER";
+    field @FlaggedApi("android.permission.flags.factory_reset_prep_permission_apis") public static final String PREPARE_FACTORY_RESET = "android.permission.PREPARE_FACTORY_RESET";
     field public static final String PROVIDE_DEFAULT_ENABLED_CREDENTIAL_SERVICE = "android.permission.PROVIDE_DEFAULT_ENABLED_CREDENTIAL_SERVICE";
     field public static final String PROVIDE_RESOLVER_RANKER_SERVICE = "android.permission.PROVIDE_RESOLVER_RANKER_SERVICE";
     field public static final String PROVIDE_TRUST_AGENT = "android.permission.PROVIDE_TRUST_AGENT";
diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index f3bad3a..f4c8429 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -780,6 +780,25 @@
 
 }
 
+package android.app.pinner {
+
+  @FlaggedApi("android.app.pinner_service_client_api") public final class PinnedFileStat implements android.os.Parcelable {
+    ctor @FlaggedApi("android.app.pinner_service_client_api") public PinnedFileStat(@NonNull String, long, @NonNull String);
+    method @FlaggedApi("android.app.pinner_service_client_api") public int describeContents();
+    method @FlaggedApi("android.app.pinner_service_client_api") public long getBytesPinned();
+    method @FlaggedApi("android.app.pinner_service_client_api") @NonNull public String getFilename();
+    method @FlaggedApi("android.app.pinner_service_client_api") @NonNull public String getGroupName();
+    method @FlaggedApi("android.app.pinner_service_client_api") public void writeToParcel(@NonNull android.os.Parcel, int);
+    field @FlaggedApi("android.app.pinner_service_client_api") @NonNull public static final android.os.Parcelable.Creator<android.app.pinner.PinnedFileStat> CREATOR;
+  }
+
+  @FlaggedApi("android.app.pinner_service_client_api") public class PinnerServiceClient {
+    ctor @FlaggedApi("android.app.pinner_service_client_api") public PinnerServiceClient();
+    method @FlaggedApi("android.app.pinner_service_client_api") @NonNull public java.util.List<android.app.pinner.PinnedFileStat> getPinnerStats();
+  }
+
+}
+
 package android.app.prediction {
 
   public final class AppPredictor {
@@ -4201,8 +4220,13 @@
   public static class WindowInfosListenerForTest.WindowInfo {
     field @NonNull public final android.graphics.Rect bounds;
     field public final int displayId;
+    field public final boolean isDuplicateTouchToWallpaper;
+    field public final boolean isFocusable;
+    field public final boolean isPreventSplitting;
+    field public final boolean isTouchable;
     field public final boolean isTrustedOverlay;
     field public final boolean isVisible;
+    field public final boolean isWatchOutsideTouch;
     field @NonNull public final String name;
     field @NonNull public final android.graphics.Matrix transform;
     field @NonNull public final android.os.IBinder windowToken;
diff --git a/core/java/android/app/OWNERS b/core/java/android/app/OWNERS
index d8448dc..772b0b4 100644
--- a/core/java/android/app/OWNERS
+++ b/core/java/android/app/OWNERS
@@ -85,6 +85,9 @@
 per-file IInstantAppResolver.aidl = file:/services/core/java/com/android/server/pm/OWNERS
 per-file InstantAppResolveInfo.aidl = file:/services/core/java/com/android/server/pm/OWNERS
 
+# Pinner
+per-file pinner-client.aconfig = file:/core/java/android/app/pinner/OWNERS
+
 # ResourcesManager
 per-file ResourcesManager.java = file:RESOURCES_OWNERS
 
diff --git a/core/java/android/app/pinner-client.aconfig b/core/java/android/app/pinner-client.aconfig
new file mode 100644
index 0000000..b60ad9e
--- /dev/null
+++ b/core/java/android/app/pinner-client.aconfig
@@ -0,0 +1,8 @@
+package: "android.app"
+
+flag {
+     namespace: "system_performance"
+     name: "pinner_service_client_api"
+     description: "Control exposing PinnerService APIs."
+     bug: "307594624"
+}
\ No newline at end of file
diff --git a/core/java/android/app/pinner/IPinnerService.aidl b/core/java/android/app/pinner/IPinnerService.aidl
new file mode 100644
index 0000000..e5d0a05
--- /dev/null
+++ b/core/java/android/app/pinner/IPinnerService.aidl
@@ -0,0 +1,12 @@
+package android.app.pinner;
+
+import android.app.pinner.PinnedFileStat;
+
+/**
+ * Interface for processes to communicate with system's PinnerService.
+ * @hide
+ */
+interface IPinnerService {
+    @EnforcePermission("DUMP")
+    List<PinnedFileStat> getPinnerStats();
+}
\ No newline at end of file
diff --git a/core/java/android/app/pinner/OWNERS b/core/java/android/app/pinner/OWNERS
new file mode 100644
index 0000000..3e3fa66
--- /dev/null
+++ b/core/java/android/app/pinner/OWNERS
@@ -0,0 +1,10 @@
+carmenjackson@google.com
+dualli@google.com
+edgararriaga@google.com
+kevinjeon@google.com
+philipcuadra@google.com
+shombert@google.com
+timmurray@google.com
+wessam@google.com
+jdduke@google.com
+shayba@google.com
\ No newline at end of file
diff --git a/core/java/android/app/pinner/PinnedFileStat.aidl b/core/java/android/app/pinner/PinnedFileStat.aidl
new file mode 100644
index 0000000..44217cf
--- /dev/null
+++ b/core/java/android/app/pinner/PinnedFileStat.aidl
@@ -0,0 +1,3 @@
+package android.app.pinner;
+
+parcelable PinnedFileStat;
\ No newline at end of file
diff --git a/core/java/android/app/pinner/PinnedFileStat.java b/core/java/android/app/pinner/PinnedFileStat.java
new file mode 100644
index 0000000..2e36330
--- /dev/null
+++ b/core/java/android/app/pinner/PinnedFileStat.java
@@ -0,0 +1,133 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.app.pinner;
+
+import static android.app.Flags.FLAG_PINNER_SERVICE_CLIENT_API;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.annotation.TestApi;
+import android.os.Parcel;
+import android.os.Parcelable;
+
+/**
+ * @hide
+ */
+@TestApi
+@FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+public final class PinnedFileStat implements Parcelable {
+    private String filename;
+    private long bytesPinned;
+    private String groupName;
+
+    /**
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+    public long getBytesPinned() {
+        return bytesPinned;
+    }
+
+    /**
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+    public @NonNull String getFilename() {
+        return filename;
+    }
+
+    /**
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+    public @NonNull String getGroupName() {
+        return groupName;
+    }
+
+    /**
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+    public PinnedFileStat(@NonNull String filename, long bytesPinned, @NonNull String groupName) {
+        this.filename = filename;
+        this.bytesPinned = bytesPinned;
+        this.groupName = groupName;
+    }
+
+    private PinnedFileStat(Parcel source) {
+        readFromParcel(source);
+    }
+
+    /**
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+    @Override
+    public void writeToParcel(@NonNull Parcel dest, int flags) {
+        dest.writeString8(filename);
+        dest.writeLong(bytesPinned);
+        dest.writeString8(groupName);
+    }
+
+    private void readFromParcel(@NonNull Parcel source) {
+        filename = source.readString8();
+        bytesPinned = source.readLong();
+        groupName = source.readString8();
+    }
+
+    /**
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    /**
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+    public static final @NonNull Creator<PinnedFileStat> CREATOR = new Creator<>() {
+        /**
+         * @hide
+         */
+        @TestApi
+        @FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+        @Override
+        public PinnedFileStat createFromParcel(Parcel source) {
+            return new PinnedFileStat(source);
+        }
+
+        /**
+         * @hide
+         */
+        @TestApi
+        @FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+        @Override
+        public PinnedFileStat[] newArray(int size) {
+            return new PinnedFileStat[size];
+        }
+    };
+}
diff --git a/core/java/android/app/pinner/PinnerServiceClient.java b/core/java/android/app/pinner/PinnerServiceClient.java
new file mode 100644
index 0000000..8b7c6cc
--- /dev/null
+++ b/core/java/android/app/pinner/PinnerServiceClient.java
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.app.pinner;
+
+import static android.app.Flags.FLAG_PINNER_SERVICE_CLIENT_API;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.annotation.TestApi;
+import android.app.pinner.IPinnerService;
+import android.app.pinner.PinnedFileStat;
+import android.os.IBinder;
+import android.os.RemoteException;
+import android.os.ServiceManager;
+import android.util.Slog;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Expose PinnerService as an interface to apps.
+ * @hide
+ */
+@TestApi
+@FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+public class PinnerServiceClient {
+    private static String TAG = "PinnerServiceClient";
+    /**
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+    public PinnerServiceClient() {}
+
+    /**
+     * Obtain the pinned file stats used for testing infrastructure.
+     * @return List of pinned files or an empty list if failed to retrieve them.
+     * @throws RuntimeException on failure to retrieve stats.
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_PINNER_SERVICE_CLIENT_API)
+    public @NonNull List<PinnedFileStat> getPinnerStats() {
+        IBinder binder = ServiceManager.getService("pinner");
+        if (binder == null) {
+            Slog.w(TAG,
+                    "Failed to retrieve PinnerService. A common failure reason is due to a lack of selinux permissions.");
+            return new ArrayList<>();
+        }
+        IPinnerService pinnerService = IPinnerService.Stub.asInterface(binder);
+        if (pinnerService == null) {
+            Slog.w(TAG, "Failed to cast PinnerService.");
+            return new ArrayList<>();
+        }
+        List<PinnedFileStat> stats;
+        try {
+            stats = pinnerService.getPinnerStats();
+        } catch (RemoteException e) {
+            throw new RuntimeException("Failed to retrieve stats from PinnerService");
+        }
+        return stats;
+    }
+}
diff --git a/core/java/android/appwidget/AppWidgetManager.java b/core/java/android/appwidget/AppWidgetManager.java
index c6012bb..51a7f1c 100644
--- a/core/java/android/appwidget/AppWidgetManager.java
+++ b/core/java/android/appwidget/AppWidgetManager.java
@@ -52,6 +52,7 @@
 
 import com.android.internal.appwidget.IAppWidgetService;
 import com.android.internal.os.BackgroundThread;
+import com.android.internal.util.FunctionalUtils;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -562,6 +563,40 @@
         });
     }
 
+    private void tryAdapterConversion(
+            FunctionalUtils.RemoteExceptionIgnoringConsumer<RemoteViews> action,
+            RemoteViews original, String failureMsg) {
+        final boolean isConvertingAdapter = RemoteViews.isAdapterConversionEnabled()
+                && (mHasPostedLegacyLists = mHasPostedLegacyLists
+                        || (original != null && original.hasLegacyLists()));
+
+        if (isConvertingAdapter) {
+            final RemoteViews viewsCopy = new RemoteViews(original);
+            Runnable updateWidgetWithTask = () -> {
+                try {
+                    viewsCopy.collectAllIntents().get();
+                    action.acceptOrThrow(viewsCopy);
+                } catch (Exception e) {
+                    Log.e(TAG, failureMsg, e);
+                }
+            };
+
+            if (Looper.getMainLooper() == Looper.myLooper()) {
+                createUpdateExecutorIfNull().execute(updateWidgetWithTask);
+                return;
+            }
+
+            updateWidgetWithTask.run();
+            return;
+        }
+
+        try {
+            action.acceptOrThrow(original);
+        } catch (RemoteException re) {
+            throw re.rethrowFromSystemServer();
+        }
+    }
+
     /**
      * Set the RemoteViews to use for the specified appWidgetIds.
      * <p>
@@ -586,32 +621,8 @@
             return;
         }
 
-        final boolean isConvertingAdapter = RemoteViews.isAdapterConversionEnabled()
-                && (mHasPostedLegacyLists = mHasPostedLegacyLists
-                        || (views != null && views.hasLegacyLists()));
-
-        if (isConvertingAdapter) {
-            views.collectAllIntents();
-
-            if (Looper.getMainLooper() == Looper.myLooper()) {
-                RemoteViews viewsCopy = new RemoteViews(views);
-                createUpdateExecutorIfNull().execute(() -> {
-                    try {
-                        mService.updateAppWidgetIds(mPackageName, appWidgetIds, viewsCopy);
-                    } catch (RemoteException e) {
-                        Log.e(TAG, "Error updating app widget views in background", e);
-                    }
-                });
-
-                return;
-            }
-        }
-
-        try {
-            mService.updateAppWidgetIds(mPackageName, appWidgetIds, views);
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
+        tryAdapterConversion(view -> mService.updateAppWidgetIds(mPackageName, appWidgetIds,
+                view), views, "Error updating app widget views in background");
     }
 
     /**
@@ -716,32 +727,9 @@
             return;
         }
 
-        final boolean isConvertingAdapter = RemoteViews.isAdapterConversionEnabled()
-                && (mHasPostedLegacyLists = mHasPostedLegacyLists
-                        || (views != null && views.hasLegacyLists()));
-
-        if (isConvertingAdapter) {
-            views.collectAllIntents();
-
-            if (Looper.getMainLooper() == Looper.myLooper()) {
-                RemoteViews viewsCopy = new RemoteViews(views);
-                createUpdateExecutorIfNull().execute(() -> {
-                    try {
-                        mService.partiallyUpdateAppWidgetIds(mPackageName, appWidgetIds, viewsCopy);
-                    } catch (RemoteException e) {
-                        Log.e(TAG, "Error partially updating app widget views in background", e);
-                    }
-                });
-
-                return;
-            }
-        }
-
-        try {
-            mService.partiallyUpdateAppWidgetIds(mPackageName, appWidgetIds, views);
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
+        tryAdapterConversion(view -> mService.partiallyUpdateAppWidgetIds(mPackageName,
+                appWidgetIds, view), views,
+                "Error partially updating app widget views in background");
     }
 
     /**
@@ -793,33 +781,8 @@
             return;
         }
 
-        final boolean isConvertingAdapter = RemoteViews.isAdapterConversionEnabled()
-                && (mHasPostedLegacyLists = mHasPostedLegacyLists
-                        || (views != null && views.hasLegacyLists()));
-
-        if (isConvertingAdapter) {
-            views.collectAllIntents();
-
-            if (Looper.getMainLooper() == Looper.myLooper()) {
-                RemoteViews viewsCopy = new RemoteViews(views);
-                createUpdateExecutorIfNull().execute(() -> {
-                    try {
-                        mService.updateAppWidgetProvider(provider, viewsCopy);
-                    } catch (RemoteException e) {
-                        Log.e(TAG, "Error updating app widget view using provider in background",
-                                e);
-                    }
-                });
-
-                return;
-            }
-        }
-
-        try {
-            mService.updateAppWidgetProvider(provider, views);
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
+        tryAdapterConversion(view -> mService.updateAppWidgetProvider(provider, view), views,
+                "Error updating app widget view using provider in background");
     }
 
     /**
diff --git a/core/java/android/companion/virtual/flags.aconfig b/core/java/android/companion/virtual/flags.aconfig
index 10da8b1..f0477d4 100644
--- a/core/java/android/companion/virtual/flags.aconfig
+++ b/core/java/android/companion/virtual/flags.aconfig
@@ -32,6 +32,13 @@
 }
 
 flag {
+  name: "consistent_display_flags"
+  namespace: "virtual_devices"
+  description: "Make virtual display flags consistent with display manager"
+  bug: "300905478"
+}
+
+flag {
   name: "vdm_custom_home"
   namespace: "virtual_devices"
   description: "Enable custom home API"
diff --git a/core/java/android/hardware/camera2/CameraMetadata.java b/core/java/android/hardware/camera2/CameraMetadata.java
index dfc27ca..507e814 100644
--- a/core/java/android/hardware/camera2/CameraMetadata.java
+++ b/core/java/android/hardware/camera2/CameraMetadata.java
@@ -16,6 +16,7 @@
 
 package android.hardware.camera2;
 
+import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.TestApi;
 import android.compat.annotation.UnsupportedAppUsage;
@@ -24,6 +25,8 @@
 import android.hardware.camera2.impl.SyntheticKey;
 import android.util.Log;
 
+import com.android.internal.camera.flags.Flags;
+
 import java.lang.reflect.Field;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
diff --git a/core/java/android/os/BatteryConsumer.java b/core/java/android/os/BatteryConsumer.java
index ca84b35..6a83cee 100644
--- a/core/java/android/os/BatteryConsumer.java
+++ b/core/java/android/os/BatteryConsumer.java
@@ -682,6 +682,7 @@
 
     static class BatteryConsumerDataLayout {
         private static final Key[] KEY_ARRAY = new Key[0];
+        public static final int POWER_MODEL_NOT_INCLUDED = -1;
         public final String[] customPowerComponentNames;
         public final int customPowerComponentCount;
         public final boolean powerModelsIncluded;
@@ -713,7 +714,9 @@
                 // Declare the Key for the power component, ignoring other dimensions.
                 perComponentKeys.add(
                         new Key(componentId, PROCESS_STATE_ANY,
-                                powerModelsIncluded ? columnIndex++ : -1,  // power model
+                                powerModelsIncluded
+                                        ? columnIndex++
+                                        : POWER_MODEL_NOT_INCLUDED,  // power model
                                 columnIndex++,      // power
                                 columnIndex++       // usage duration
                         ));
@@ -736,7 +739,9 @@
 
                             perComponentKeys.add(
                                     new Key(componentId, processState,
-                                            powerModelsIncluded ? columnIndex++ : -1, // power model
+                                            powerModelsIncluded
+                                                    ? columnIndex++
+                                                    : POWER_MODEL_NOT_INCLUDED, // power model
                                             columnIndex++,      // power
                                             columnIndex++       // usage duration
                                     ));
@@ -843,11 +848,27 @@
 
         @SuppressWarnings("unchecked")
         @NonNull
+        public T addConsumedPower(@PowerComponent int componentId, double componentPower,
+                @PowerModel int powerModel) {
+            mPowerComponentsBuilder.addConsumedPower(getKey(componentId, PROCESS_STATE_UNSPECIFIED),
+                    componentPower, powerModel);
+            return (T) this;
+        }
+
+        @SuppressWarnings("unchecked")
+        @NonNull
         public T setConsumedPower(Key key, double componentPower, @PowerModel int powerModel) {
             mPowerComponentsBuilder.setConsumedPower(key, componentPower, powerModel);
             return (T) this;
         }
 
+        @SuppressWarnings("unchecked")
+        @NonNull
+        public T addConsumedPower(Key key, double componentPower, @PowerModel int powerModel) {
+            mPowerComponentsBuilder.addConsumedPower(key, componentPower, powerModel);
+            return (T) this;
+        }
+
         /**
          * Sets the amount of drain attributed to the specified custom drain type.
          *
diff --git a/core/java/android/os/BatteryStats.java b/core/java/android/os/BatteryStats.java
index e85b7bf..16ffaef 100644
--- a/core/java/android/os/BatteryStats.java
+++ b/core/java/android/os/BatteryStats.java
@@ -4029,6 +4029,17 @@
     }
 
     /**
+     * A helper object passed to various dump... methods to integrate with such objects
+     * as BatteryUsageStatsProvider.
+     */
+    public interface BatteryStatsDumpHelper {
+        /**
+         * Generates BatteryUsageStats based on the specified BatteryStats.
+         */
+        BatteryUsageStats getBatteryUsageStats(BatteryStats batteryStats, boolean detailed);
+    }
+
+    /**
      * Dumps the ControllerActivityCounter if it has any data worth dumping.
      * The order of the arguments in the final check in line is:
      *
@@ -4390,7 +4401,7 @@
      * NOTE: all times are expressed in microseconds, unless specified otherwise.
      */
     public final void dumpCheckinLocked(Context context, PrintWriter pw, int which, int reqUid,
-            boolean wifiOnly) {
+            boolean wifiOnly, BatteryStatsDumpHelper dumpHelper) {
 
         if (which != BatteryStats.STATS_SINCE_CHARGED) {
             dumpLine(pw, 0, STAT_NAMES[which], "err",
@@ -4652,7 +4663,7 @@
             }
         }
 
-        final BatteryUsageStats stats = getBatteryUsageStats(context, true /* detailed */);
+        final BatteryUsageStats stats = dumpHelper.getBatteryUsageStats(this, true /* detailed */);
         dumpLine(pw, 0 /* uid */, category, POWER_USE_SUMMARY_DATA,
                 formatCharge(stats.getBatteryCapacity()),
                 formatCharge(stats.getConsumedPower()),
@@ -5127,7 +5138,7 @@
 
     @SuppressWarnings("unused")
     public final void dumpLocked(Context context, PrintWriter pw, String prefix, final int which,
-            int reqUid, boolean wifiOnly) {
+            int reqUid, boolean wifiOnly, BatteryStatsDumpHelper dumpHelper) {
 
         if (which != BatteryStats.STATS_SINCE_CHARGED) {
             pw.println("ERROR: BatteryStats.dump called for which type " + which
@@ -5854,7 +5865,7 @@
         pw.println();
 
 
-        BatteryUsageStats stats = getBatteryUsageStats(context, true /* detailed */);
+        BatteryUsageStats stats = dumpHelper.getBatteryUsageStats(this, true /* detailed */);
         stats.dump(pw, prefix);
 
         List<UidMobileRadioStats> uidMobileRadioStats =
@@ -7642,10 +7653,11 @@
     /**
      * Dumps a human-readable summary of the battery statistics to the given PrintWriter.
      *
-     * @param pw a Printer to receive the dump output.
+     * @param pw         a Printer to receive the dump output.
      */
     @SuppressWarnings("unused")
-    public void dump(Context context, PrintWriter pw, int flags, int reqUid, long histStart) {
+    public void dump(Context context, PrintWriter pw, int flags, int reqUid, long histStart,
+            BatteryStatsDumpHelper dumpHelper) {
         synchronized (this) {
             prepareForDumpLocked();
         }
@@ -7663,12 +7675,12 @@
         }
 
         synchronized (this) {
-            dumpLocked(context, pw, flags, reqUid, filtering);
+            dumpLocked(context, pw, flags, reqUid, filtering, dumpHelper);
         }
     }
 
     private void dumpLocked(Context context, PrintWriter pw, int flags, int reqUid,
-            boolean filtering) {
+            boolean filtering, BatteryStatsDumpHelper dumpHelper) {
         if (!filtering) {
             SparseArray<? extends Uid> uidStats = getUidStats();
             final int NU = uidStats.size();
@@ -7803,15 +7815,15 @@
             pw.println("  System starts: " + getStartCount()
                     + ", currently on battery: " + getIsOnBattery());
             dumpLocked(context, pw, "", STATS_SINCE_CHARGED, reqUid,
-                    (flags&DUMP_DEVICE_WIFI_ONLY) != 0);
+                    (flags & DUMP_DEVICE_WIFI_ONLY) != 0, dumpHelper);
             pw.println();
         }
     }
 
     // This is called from BatteryStatsService.
     @SuppressWarnings("unused")
-    public void dumpCheckin(Context context, PrintWriter pw,
-            List<ApplicationInfo> apps, int flags, long histStart) {
+    public void dumpCheckin(Context context, PrintWriter pw, List<ApplicationInfo> apps, int flags,
+            long histStart, BatteryStatsDumpHelper dumpHelper) {
         synchronized (this) {
             prepareForDumpLocked();
 
@@ -7829,12 +7841,12 @@
         }
 
         synchronized (this) {
-            dumpCheckinLocked(context, pw, apps, flags);
+            dumpCheckinLocked(context, pw, apps, flags, dumpHelper);
         }
     }
 
     private void dumpCheckinLocked(Context context, PrintWriter pw, List<ApplicationInfo> apps,
-            int flags) {
+            int flags, BatteryStatsDumpHelper dumpHelper) {
         if (apps != null) {
             SparseArray<Pair<ArrayList<String>, MutableBoolean>> uids = new SparseArray<>();
             for (int i=0; i<apps.size(); i++) {
@@ -7881,7 +7893,7 @@
                         (Object[])lineArgs);
             }
             dumpCheckinLocked(context, pw, STATS_SINCE_CHARGED, -1,
-                    (flags&DUMP_DEVICE_WIFI_ONLY) != 0);
+                    (flags & DUMP_DEVICE_WIFI_ONLY) != 0, dumpHelper);
         }
     }
 
@@ -7891,7 +7903,7 @@
      * @hide
      */
     public void dumpProtoLocked(Context context, FileDescriptor fd, List<ApplicationInfo> apps,
-            int flags, long histStart) {
+            int flags, long histStart, BatteryStatsDumpHelper dumpHelper) {
         final ProtoOutputStream proto = new ProtoOutputStream(fd);
         prepareForDumpLocked();
 
@@ -7909,7 +7921,8 @@
         proto.write(BatteryStatsProto.END_PLATFORM_VERSION, getEndPlatformVersion());
 
         if ((flags & DUMP_DAILY_ONLY) == 0) {
-            final BatteryUsageStats stats = getBatteryUsageStats(context, false /* detailed */);
+            final BatteryUsageStats stats =
+                    dumpHelper.getBatteryUsageStats(this, false /* detailed */);
             ProportionalAttributionCalculator proportionalAttributionCalculator =
                     new ProportionalAttributionCalculator(context, stats);
             dumpProtoAppsLocked(proto, stats, apps, proportionalAttributionCalculator);
@@ -8856,8 +8869,6 @@
         return !tm.isDataCapable();
     }
 
-    protected abstract BatteryUsageStats getBatteryUsageStats(Context context, boolean detailed);
-
     private boolean shouldHidePowerComponent(int powerComponent) {
         return powerComponent == BatteryConsumer.POWER_COMPONENT_IDLE
                 || powerComponent == BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO
diff --git a/core/java/android/os/BatteryUsageStats.java b/core/java/android/os/BatteryUsageStats.java
index cd52b5c0..ed31002 100644
--- a/core/java/android/os/BatteryUsageStats.java
+++ b/core/java/android/os/BatteryUsageStats.java
@@ -36,6 +36,7 @@
 import java.io.FileDescriptor;
 import java.io.IOException;
 import java.io.PrintWriter;
+import java.io.StringWriter;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
@@ -586,7 +587,8 @@
                             + "(" + BatteryConsumer.processStateToString(key.processState) + ")";
                 }
                 printPowerComponent(pw, prefix, label, devicePowerMah, appsPowerMah,
-                        deviceConsumer.getPowerModel(key),
+                        mIncludesPowerModels ? deviceConsumer.getPowerModel(key)
+                                : BatteryConsumer.POWER_MODEL_UNDEFINED,
                         deviceConsumer.getUsageDurationMillis(key));
             }
         }
@@ -774,6 +776,15 @@
         super.finalize();
     }
 
+    @Override
+    public String toString() {
+        StringWriter sw = new StringWriter();
+        PrintWriter pw = new PrintWriter(sw);
+        dump(pw, "");
+        pw.flush();
+        return sw.toString();
+    }
+
     /**
      * Builder for BatteryUsageStats.
      */
diff --git a/core/java/android/os/PowerComponents.java b/core/java/android/os/PowerComponents.java
index 9e5f539..9c11ad4 100644
--- a/core/java/android/os/PowerComponents.java
+++ b/core/java/android/os/PowerComponents.java
@@ -15,6 +15,7 @@
  */
 package android.os;
 
+import static android.os.BatteryConsumer.BatteryConsumerDataLayout.POWER_MODEL_NOT_INCLUDED;
 import static android.os.BatteryConsumer.POWER_COMPONENT_ANY;
 import static android.os.BatteryConsumer.PROCESS_STATE_ANY;
 import static android.os.BatteryConsumer.PROCESS_STATE_UNSPECIFIED;
@@ -118,7 +119,7 @@
 
     @BatteryConsumer.PowerModel
     int getPowerModel(BatteryConsumer.Key key) {
-        if (key.mPowerModelColumnIndex == -1) {
+        if (key.mPowerModelColumnIndex == POWER_MODEL_NOT_INCLUDED) {
             throw new IllegalStateException(
                     "Power model IDs were not requested in the BatteryUsageStatsQuery");
         }
@@ -468,7 +469,7 @@
             mMinConsumedPowerThreshold = minConsumedPowerThreshold;
             for (BatteryConsumer.Key[] keys : mData.layout.keys) {
                 for (BatteryConsumer.Key key : keys) {
-                    if (key.mPowerModelColumnIndex != -1) {
+                    if (key.mPowerModelColumnIndex != POWER_MODEL_NOT_INCLUDED) {
                         mData.putInt(key.mPowerModelColumnIndex, POWER_MODEL_UNINITIALIZED);
                     }
                 }
@@ -478,11 +479,19 @@
         @NonNull
         public Builder setConsumedPower(BatteryConsumer.Key key, double componentPower,
                 int powerModel) {
-            if (Math.abs(componentPower) < mMinConsumedPowerThreshold) {
-                componentPower = 0;
-            }
             mData.putDouble(key.mPowerColumnIndex, componentPower);
-            if (key.mPowerModelColumnIndex != -1) {
+            if (key.mPowerModelColumnIndex != POWER_MODEL_NOT_INCLUDED) {
+                mData.putInt(key.mPowerModelColumnIndex, powerModel);
+            }
+            return this;
+        }
+
+        @NonNull
+        public Builder addConsumedPower(BatteryConsumer.Key key, double componentPower,
+                int powerModel) {
+            mData.putDouble(key.mPowerColumnIndex,
+                    mData.getDouble(key.mPowerColumnIndex) + componentPower);
+            if (key.mPowerModelColumnIndex != POWER_MODEL_NOT_INCLUDED) {
                 mData.putInt(key.mPowerModelColumnIndex, powerModel);
             }
             return this;
@@ -496,9 +505,6 @@
          */
         @NonNull
         public Builder setConsumedPowerForCustomComponent(int componentId, double componentPower) {
-            if (Math.abs(componentPower) < mMinConsumedPowerThreshold) {
-                componentPower = 0;
-            }
             final int index = componentId - BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID;
             if (index < 0 || index >= mData.layout.customPowerComponentCount) {
                 throw new IllegalArgumentException(
@@ -575,12 +581,12 @@
                             mData.getLong(key.mDurationColumnIndex)
                                     + otherData.getLong(otherKey.mDurationColumnIndex));
 
-                    if (key.mPowerModelColumnIndex == -1) {
+                    if (key.mPowerModelColumnIndex == POWER_MODEL_NOT_INCLUDED) {
                         continue;
                     }
 
                     boolean undefined = false;
-                    if (otherKey.mPowerModelColumnIndex == -1) {
+                    if (otherKey.mPowerModelColumnIndex == POWER_MODEL_NOT_INCLUDED) {
                         undefined = true;
                     } else {
                         final int powerModel = mData.getInt(key.mPowerModelColumnIndex);
@@ -641,19 +647,26 @@
          */
         @NonNull
         public PowerComponents build() {
-            mData.putDouble(mData.layout.totalConsumedPowerColumnIndex, getTotalPower());
-
             for (BatteryConsumer.Key[] keys : mData.layout.keys) {
                 for (BatteryConsumer.Key key : keys) {
-                    if (key.mPowerModelColumnIndex != -1) {
+                    if (key.mPowerModelColumnIndex != POWER_MODEL_NOT_INCLUDED) {
                         if (mData.getInt(key.mPowerModelColumnIndex) == POWER_MODEL_UNINITIALIZED) {
                             mData.putInt(key.mPowerModelColumnIndex,
                                     BatteryConsumer.POWER_MODEL_UNDEFINED);
                         }
                     }
+
+                    if (mMinConsumedPowerThreshold != 0) {
+                        if (mData.getDouble(key.mPowerColumnIndex) < mMinConsumedPowerThreshold) {
+                            mData.putDouble(key.mPowerColumnIndex, 0);
+                        }
+                    }
                 }
             }
 
+            if (mData.getDouble(mData.layout.totalConsumedPowerColumnIndex) == 0) {
+                mData.putDouble(mData.layout.totalConsumedPowerColumnIndex, getTotalPower());
+            }
             return new PowerComponents(this);
         }
     }
diff --git a/core/java/android/permission/flags.aconfig b/core/java/android/permission/flags.aconfig
index dc86e3f5..5cbc18e 100644
--- a/core/java/android/permission/flags.aconfig
+++ b/core/java/android/permission/flags.aconfig
@@ -56,4 +56,11 @@
   namespace: "permissions"
   description: "enables logging of the OP_ENABLE_MOBILE_DATA_BY_USER"
   bug: "222650148"
-}
\ No newline at end of file
+}
+
+flag {
+  name: "factory_reset_prep_permission_apis"
+  namespace: "wallet_integration"
+  description: "enable Permission PREPARE_FACTORY_RESET."
+  bug: "302016478"
+}
diff --git a/core/java/android/telephony/TelephonyRegistryManager.java b/core/java/android/telephony/TelephonyRegistryManager.java
index 0063d13..886727e 100644
--- a/core/java/android/telephony/TelephonyRegistryManager.java
+++ b/core/java/android/telephony/TelephonyRegistryManager.java
@@ -490,16 +490,32 @@
     /**
      * Notify changes to activity state changes on certain subscription.
      *
+     * @param subId for which data activity state changed.
+     * @param dataActivityType indicates the latest data activity type e.g. {@link
+     * TelephonyManager#DATA_ACTIVITY_IN}
+     */
+    public void notifyDataActivityChanged(int subId, @DataActivityType int dataActivityType) {
+        try {
+            sRegistry.notifyDataActivityForSubscriber(subId, dataActivityType);
+        } catch (RemoteException ex) {
+            // system process is dead
+            throw ex.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Notify changes to activity state changes on certain subscription.
+     *
      * @param slotIndex for which data activity changed. Can be derived from subId except
      * when subId is invalid.
      * @param subId for which data activity state changed.
-     * @param dataActivityType indicates the latest data activity type e.g, {@link
+     * @param dataActivityType indicates the latest data activity type e.g. {@link
      * TelephonyManager#DATA_ACTIVITY_IN}
      */
     public void notifyDataActivityChanged(int slotIndex, int subId,
             @DataActivityType int dataActivityType) {
         try {
-            sRegistry.notifyDataActivityForSubscriber(slotIndex, subId, dataActivityType);
+            sRegistry.notifyDataActivityForSubscriberWithSlot(slotIndex, subId, dataActivityType);
         } catch (RemoteException ex) {
             // system process is dead
             throw ex.rethrowFromSystemServer();
diff --git a/core/java/android/view/SurfaceControl.java b/core/java/android/view/SurfaceControl.java
index 8befe8a..cbbe785 100644
--- a/core/java/android/view/SurfaceControl.java
+++ b/core/java/android/view/SurfaceControl.java
@@ -503,9 +503,6 @@
     // be dumped as additional context
     private static volatile boolean sDebugUsageAfterRelease = false;
 
-    static GlobalTransactionWrapper sGlobalTransaction;
-    static long sTransactionNestCount = 0;
-
     private static final NativeAllocationRegistry sRegistry =
             NativeAllocationRegistry.createMalloced(SurfaceControl.class.getClassLoader(),
                     nativeGetNativeSurfaceControlFinalizer());
@@ -859,33 +856,47 @@
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
     @IntDef(prefix = {"FRAME_RATE_SELECTION_STRATEGY_"},
-            value = {FRAME_RATE_SELECTION_STRATEGY_SELF,
-                    FRAME_RATE_SELECTION_STRATEGY_OVERRIDE_CHILDREN})
+            value = {FRAME_RATE_SELECTION_STRATEGY_PROPAGATE,
+                    FRAME_RATE_SELECTION_STRATEGY_OVERRIDE_CHILDREN,
+                    FRAME_RATE_SELECTION_STRATEGY_SELF})
     public @interface FrameRateSelectionStrategy {}
 
     // From window.h. Keep these in sync.
     /**
      * Default value. The layer uses its own frame rate specifications, assuming it has any
-     * specifications, instead of its parent's.
+     * specifications, instead of its parent's. If it does not have its own frame rate
+     * specifications, it will try to use its parent's. It will propagate its specifications to any
+     * descendants that do not have their own.
+     *
      * However, {@link #FRAME_RATE_SELECTION_STRATEGY_OVERRIDE_CHILDREN} on an ancestor layer
-     * supersedes this behavior, meaning that this layer will inherit the frame rate specifications
-     * of that ancestor layer.
+     * supersedes this behavior, meaning that this layer will inherit frame rate specifications
+     * regardless of whether it has its own.
      * @hide
      */
-    public static final int FRAME_RATE_SELECTION_STRATEGY_SELF = 0;
+    public static final int FRAME_RATE_SELECTION_STRATEGY_PROPAGATE = 0;
 
     /**
      * The layer's frame rate specifications will propagate to and override those of its descendant
      * layers.
-     * The layer with this strategy has the {@link #FRAME_RATE_SELECTION_STRATEGY_SELF} behavior
-     * for itself. This does mean that any parent or ancestor layer that also has the strategy
-     * {@link FRAME_RATE_SELECTION_STRATEGY_OVERRIDE_CHILDREN} will override this layer's
+     *
+     * The layer itself has the {@link #FRAME_RATE_SELECTION_STRATEGY_PROPAGATE} behavior.
+     * Thus, ancestor layer that also has the strategy
+     * {@link #FRAME_RATE_SELECTION_STRATEGY_OVERRIDE_CHILDREN} will override this layer's
      * frame rate specifications.
      * @hide
      */
     public static final int FRAME_RATE_SELECTION_STRATEGY_OVERRIDE_CHILDREN = 1;
 
     /**
+     * The layer's frame rate specifications will not propagate to its descendant
+     * layers, even if the descendant layer has no frame rate specifications.
+     * However, {@link #FRAME_RATE_SELECTION_STRATEGY_OVERRIDE_CHILDREN} on an ancestor
+     * layer supersedes this behavior.
+     * @hide
+     */
+    public static final int FRAME_RATE_SELECTION_STRATEGY_SELF = 2;
+
+    /**
      * Builder class for {@link SurfaceControl} objects.
      *
      * By default the surface will be hidden, and have "unset" bounds, meaning it can
@@ -1576,54 +1587,30 @@
         return mNativeObject != 0;
     }
 
-    /*
-     * set surface parameters.
-     * needs to be inside open/closeTransaction block
-     */
-
     /** start a transaction
      * @hide
-     */
-    @UnsupportedAppUsage
-    public static void openTransaction() {
-        synchronized (SurfaceControl.class) {
-            if (sGlobalTransaction == null) {
-                sGlobalTransaction = new GlobalTransactionWrapper();
-            }
-            synchronized(SurfaceControl.class) {
-                sTransactionNestCount++;
-            }
-        }
-    }
-
-    /**
-     * Merge the supplied transaction in to the deprecated "global" transaction.
-     * This clears the supplied transaction in an identical fashion to {@link Transaction#merge}.
-     * <p>
-     * This is a utility for interop with legacy-code and will go away with the Global Transaction.
-     * @hide
+     * @deprecated Use regular Transaction instead.
      */
     @Deprecated
-    public static void mergeToGlobalTransaction(Transaction t) {
-        synchronized(SurfaceControl.class) {
-            sGlobalTransaction.merge(t);
-        }
+    @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.VANILLA_ICE_CREAM,
+            publicAlternatives = "Use {@code SurfaceControl.Transaction} instead",
+            trackingBug = 247078497)
+    public static void openTransaction() {
+        // TODO(b/247078497): It was used for global transaction (all usages are removed).
+        //  Keep the method declaration to avoid breaking reference from legacy access.
     }
 
     /** end a transaction
      * @hide
+     * @deprecated Use regular Transaction instead.
      */
-    @UnsupportedAppUsage
+    @Deprecated
+    @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.VANILLA_ICE_CREAM,
+            publicAlternatives = "Use {@code SurfaceControl.Transaction} instead",
+            trackingBug = 247078497)
     public static void closeTransaction() {
-        synchronized(SurfaceControl.class) {
-            if (sTransactionNestCount == 0) {
-                Log.e(TAG,
-                        "Call to SurfaceControl.closeTransaction without matching openTransaction");
-            } else if (--sTransactionNestCount > 0) {
-                return;
-            }
-            sGlobalTransaction.applyGlobalTransaction(false);
-        }
+        // TODO(b/247078497): It was used for global transaction (all usages are removed).
+        //  Keep the method declaration to avoid breaking reference from legacy access.
     }
 
     /**
@@ -4499,39 +4486,6 @@
     }
 
     /**
-     * As part of eliminating usage of the global Transaction we expose
-     * a SurfaceControl.getGlobalTransaction function. However calling
-     * apply on this global transaction (rather than using closeTransaction)
-     * would be very dangerous. So for the global transaction we use this
-     * subclass of Transaction where the normal apply throws an exception.
-     */
-    private static class GlobalTransactionWrapper extends SurfaceControl.Transaction {
-        void applyGlobalTransaction(boolean sync) {
-            applyResizedSurfaces();
-            notifyReparentedSurfaces();
-            nativeApplyTransaction(mNativeObject, sync, /*oneWay*/ false);
-        }
-
-        @Override
-        public void apply(boolean sync) {
-            throw new RuntimeException("Global transaction must be applied from closeTransaction");
-        }
-    }
-
-    /**
-     * This is a refactoring utility function to enable lower levels of code to be refactored
-     * from using the global transaction (and instead use a passed in Transaction) without
-     * having to refactor the higher levels at the same time.
-     * The returned global transaction can't be applied, it must be applied from closeTransaction
-     * Unless you are working on removing Global Transaction usage in the WindowManager, this
-     * probably isn't a good function to use.
-     * @hide
-     */
-    public static Transaction getGlobalTransaction() {
-        return sGlobalTransaction;
-    }
-
-    /**
      * @hide
      */
     public void resize(int w, int h) {
diff --git a/core/java/android/view/View.java b/core/java/android/view/View.java
index d591f89..d58c07d 100644
--- a/core/java/android/view/View.java
+++ b/core/java/android/view/View.java
@@ -33018,7 +33018,7 @@
     }
 
     private float getSizePercentage() {
-        if (mResources == null || getAlpha() == 0 || getVisibility() != VISIBLE) {
+        if (mResources == null || getVisibility() != VISIBLE) {
             return 0;
         }
 
diff --git a/core/java/android/widget/RemoteViews.java b/core/java/android/widget/RemoteViews.java
index a919c00..1acebf4 100644
--- a/core/java/android/widget/RemoteViews.java
+++ b/core/java/android/widget/RemoteViews.java
@@ -1054,8 +1054,7 @@
     }
 
     private class SetRemoteCollectionItemListAdapterAction extends Action {
-        @NonNull
-        private CompletableFuture<RemoteCollectionItems> mItemsFuture;
+        private @Nullable RemoteCollectionItems mItems;
         final Intent mServiceIntent;
         int mIntentId = -1;
         boolean mIsReplacedIntoAction = false;
@@ -1064,92 +1063,46 @@
                 @NonNull RemoteCollectionItems items) {
             mViewId = id;
             items.setHierarchyRootData(getHierarchyRootData());
-            mItemsFuture = CompletableFuture.completedFuture(items);
+            mItems = items;
             mServiceIntent = null;
         }
 
         SetRemoteCollectionItemListAdapterAction(@IdRes int id, Intent intent) {
             mViewId = id;
-            mItemsFuture = getItemsFutureFromIntentWithTimeout(intent);
-            setHierarchyRootData(getHierarchyRootData());
+            mItems = null;
             mServiceIntent = intent;
         }
 
-        private static CompletableFuture<RemoteCollectionItems> getItemsFutureFromIntentWithTimeout(
-                Intent intent) {
-            if (intent == null) {
-                Log.e(LOG_TAG, "Null intent received when generating adapter future");
-                return CompletableFuture.completedFuture(new RemoteCollectionItems
-                        .Builder().build());
-            }
-
-            final Context context = ActivityThread.currentApplication();
-            final CompletableFuture<RemoteCollectionItems> result = new CompletableFuture<>();
-
-            context.bindService(intent, Context.BindServiceFlags.of(Context.BIND_AUTO_CREATE),
-                    result.defaultExecutor(), new ServiceConnection() {
-                        @Override
-                        public void onServiceConnected(ComponentName componentName,
-                                IBinder iBinder) {
-                            RemoteCollectionItems items;
-                            try {
-                                items = IRemoteViewsFactory.Stub.asInterface(iBinder)
-                                        .getRemoteCollectionItems();
-                            } catch (RemoteException re) {
-                                items = new RemoteCollectionItems.Builder().build();
-                                Log.e(LOG_TAG, "Error getting collection items from the factory",
-                                        re);
-                            } finally {
-                                context.unbindService(this);
-                            }
-
-                            result.complete(items);
-                        }
-
-                        @Override
-                        public void onServiceDisconnected(ComponentName componentName) { }
-                    });
-
-            result.completeOnTimeout(
-                    new RemoteCollectionItems.Builder().build(),
-                    MAX_ADAPTER_CONVERSION_WAITING_TIME_MS, TimeUnit.MILLISECONDS);
-
-            return result;
-        }
-
         SetRemoteCollectionItemListAdapterAction(Parcel parcel) {
             mViewId = parcel.readInt();
             mIntentId = parcel.readInt();
-            mItemsFuture = CompletableFuture.completedFuture(mIntentId != -1
-                    ? null
-                    : new RemoteCollectionItems(parcel, getHierarchyRootData()));
             mServiceIntent = parcel.readTypedObject(Intent.CREATOR);
+            mItems = mServiceIntent != null
+                    ? null
+                    : new RemoteCollectionItems(parcel, getHierarchyRootData());
         }
 
         @Override
         public void setHierarchyRootData(HierarchyRootData rootData) {
-            if (mIntentId == -1) {
-                mItemsFuture = mItemsFuture
-                        .thenApply(rc -> {
-                            rc.setHierarchyRootData(rootData);
-                            return rc;
-                        });
+            if (mItems != null) {
+                mItems.setHierarchyRootData(rootData);
                 return;
             }
 
-            // Set the root data for items in the cache instead
-            mCollectionCache.setHierarchyDataForId(mIntentId, rootData);
+            if (mIntentId != -1) {
+                // Set the root data for items in the cache instead
+                mCollectionCache.setHierarchyDataForId(mIntentId, rootData);
+            }
         }
 
         @Override
         public void writeToParcel(Parcel dest, int flags) {
             dest.writeInt(mViewId);
             dest.writeInt(mIntentId);
-            if (mIntentId == -1) {
-                RemoteCollectionItems items = getCollectionItemsFromFuture(mItemsFuture);
-                items.writeToParcel(dest, flags, /* attached= */ true);
-            }
             dest.writeTypedObject(mServiceIntent, flags);
+            if (mItems != null) {
+                mItems.writeToParcel(dest, flags, /* attached= */ true);
+            }
         }
 
         @Override
@@ -1159,7 +1112,9 @@
             if (target == null) return;
 
             RemoteCollectionItems items = mIntentId == -1
-                    ? getCollectionItemsFromFuture(mItemsFuture)
+                    ? mItems == null
+                            ? new RemoteCollectionItems.Builder().build()
+                            : mItems
                     : mCollectionCache.getItemsForId(mIntentId);
 
             // Ensure that we are applying to an AppWidget root
@@ -1216,51 +1171,32 @@
 
         @Override
         public void visitUris(@NonNull Consumer<Uri> visitor) {
-            RemoteCollectionItems items = getCollectionItemsFromFuture(mItemsFuture);
-            items.visitUris(visitor);
-        }
-    }
+            if (mIntentId != -1 || mItems == null) {
+                return;
+            }
 
-    private static RemoteCollectionItems getCollectionItemsFromFuture(
-            CompletableFuture<RemoteCollectionItems> itemsFuture) {
-        RemoteCollectionItems items;
-        try {
-            items = itemsFuture.get();
-        } catch (Exception e) {
-            Log.e(LOG_TAG, "Error getting collection items from future", e);
-            items = new RemoteCollectionItems.Builder().build();
+            mItems.visitUris(visitor);
         }
-
-        return items;
     }
 
     /**
      * @hide
      */
-    public void collectAllIntents() {
-        mCollectionCache.collectAllIntentsNoComplete(this);
+    public CompletableFuture<Void> collectAllIntents() {
+        return mCollectionCache.collectAllIntentsNoComplete(this);
     }
 
     private class RemoteCollectionCache {
         private SparseArray<String> mIdToUriMapping = new SparseArray<>();
         private HashMap<String, RemoteCollectionItems> mUriToCollectionMapping = new HashMap<>();
 
-        // We don't put this into the parcel
-        private HashMap<String, CompletableFuture<RemoteCollectionItems>> mTempUriToFutureMapping =
-                new HashMap<>();
-
         RemoteCollectionCache() { }
 
         RemoteCollectionCache(RemoteCollectionCache src) {
-            boolean isWaitingCache = src.mTempUriToFutureMapping.size() != 0;
             for (int i = 0; i < src.mIdToUriMapping.size(); i++) {
                 String uri = src.mIdToUriMapping.valueAt(i);
                 mIdToUriMapping.put(src.mIdToUriMapping.keyAt(i), uri);
-                if (isWaitingCache) {
-                    mTempUriToFutureMapping.put(uri, src.mTempUriToFutureMapping.get(uri));
-                } else {
-                    mUriToCollectionMapping.put(uri, src.mUriToCollectionMapping.get(uri));
-                }
+                mUriToCollectionMapping.put(uri, src.mUriToCollectionMapping.get(uri));
             }
         }
 
@@ -1281,14 +1217,8 @@
 
         void setHierarchyDataForId(int intentId, HierarchyRootData data) {
             String uri = mIdToUriMapping.get(intentId);
-            if (mTempUriToFutureMapping.get(uri) != null) {
-                CompletableFuture<RemoteCollectionItems> itemsFuture =
-                        mTempUriToFutureMapping.get(uri);
-                mTempUriToFutureMapping.put(uri, itemsFuture.thenApply(rc -> {
-                    rc.setHierarchyRootData(data);
-                    return rc;
-                }));
-
+            if (mUriToCollectionMapping.get(uri) == null) {
+                Log.e(LOG_TAG, "Error setting hierarchy data for id=" + intentId);
                 return;
             }
 
@@ -1301,14 +1231,17 @@
             return mUriToCollectionMapping.get(uri);
         }
 
-        void collectAllIntentsNoComplete(@NonNull RemoteViews inViews) {
+        CompletableFuture<Void> collectAllIntentsNoComplete(@NonNull RemoteViews inViews) {
+            CompletableFuture<Void> collectionFuture = CompletableFuture.completedFuture(null);
             if (inViews.hasSizedRemoteViews()) {
                 for (RemoteViews remoteViews : inViews.mSizedRemoteViews) {
-                    remoteViews.collectAllIntents();
+                    collectionFuture = CompletableFuture.allOf(collectionFuture,
+                            collectAllIntentsNoComplete(remoteViews));
                 }
             } else if (inViews.hasLandscapeAndPortraitLayouts()) {
-                inViews.mLandscape.collectAllIntents();
-                inViews.mPortrait.collectAllIntents();
+                collectionFuture = CompletableFuture.allOf(
+                        collectAllIntentsNoComplete(inViews.mLandscape),
+                        collectAllIntentsNoComplete(inViews.mPortrait));
             } else if (inViews.mActions != null) {
                 for (Action action : inViews.mActions) {
                     if (action instanceof SetRemoteCollectionItemListAdapterAction rca) {
@@ -1318,40 +1251,95 @@
                         }
 
                         if (rca.mIntentId != -1 && rca.mIsReplacedIntoAction) {
-                            String uri = mIdToUriMapping.get(rca.mIntentId);
-                            mTempUriToFutureMapping.put(uri, rca.mItemsFuture);
-                            rca.mItemsFuture = CompletableFuture.completedFuture(null);
+                            final String uri = mIdToUriMapping.get(rca.mIntentId);
+                            collectionFuture = CompletableFuture.allOf(collectionFuture,
+                                    getItemsFutureFromIntentWithTimeout(rca.mServiceIntent)
+                                            .thenAccept(rc -> {
+                                                rc.setHierarchyRootData(getHierarchyRootData());
+                                                mUriToCollectionMapping.put(uri, rc);
+                                            }));
+                            rca.mItems = null;
                             continue;
                         }
 
                         // Differentiate between the normal collection actions and the ones with
                         // intents.
                         if (rca.mServiceIntent != null) {
-                            String uri = rca.mServiceIntent.toUri(0);
+                            final String uri = rca.mServiceIntent.toUri(0);
                             int index = mIdToUriMapping.indexOfValue(uri);
                             if (index == -1) {
                                 int newIntentId = mIdToUriMapping.size();
                                 rca.mIntentId = newIntentId;
                                 mIdToUriMapping.put(newIntentId, uri);
-                                // mUriToIntentMapping.put(uri, mServiceIntent);
-                                mTempUriToFutureMapping.put(uri, rca.mItemsFuture);
                             } else {
                                 rca.mIntentId = mIdToUriMapping.keyAt(index);
+                                rca.mItems = null;
+                                continue;
                             }
-                            rca.mItemsFuture = CompletableFuture.completedFuture(null);
+                            collectionFuture = CompletableFuture.allOf(collectionFuture,
+                                    getItemsFutureFromIntentWithTimeout(rca.mServiceIntent)
+                                            .thenAccept(rc -> {
+                                                rc.setHierarchyRootData(getHierarchyRootData());
+                                                mUriToCollectionMapping.put(uri, rc);
+                                            }));
+                            rca.mItems = null;
                         } else {
-                            RemoteCollectionItems items = getCollectionItemsFromFuture(
-                                    rca.mItemsFuture);
-                            for (RemoteViews views : items.mViews) {
-                                views.collectAllIntents();
+                            for (RemoteViews views : rca.mItems.mViews) {
+                                collectionFuture = CompletableFuture.allOf(collectionFuture,
+                                        collectAllIntentsNoComplete(views));
                             }
                         }
                     } else if (action instanceof ViewGroupActionAdd vgaa
                             && vgaa.mNestedViews != null) {
-                        vgaa.mNestedViews.collectAllIntents();
+                        collectionFuture = CompletableFuture.allOf(collectionFuture,
+                                collectAllIntentsNoComplete(vgaa.mNestedViews));
                     }
                 }
             }
+
+            return collectionFuture;
+        }
+
+        private static CompletableFuture<RemoteCollectionItems> getItemsFutureFromIntentWithTimeout(
+                Intent intent) {
+            if (intent == null) {
+                Log.e(LOG_TAG, "Null intent received when generating adapter future");
+                return CompletableFuture.completedFuture(new RemoteCollectionItems
+                    .Builder().build());
+            }
+
+            final Context context = ActivityThread.currentApplication();
+            final CompletableFuture<RemoteCollectionItems> result = new CompletableFuture<>();
+
+            context.bindService(intent, Context.BindServiceFlags.of(Context.BIND_AUTO_CREATE),
+                    result.defaultExecutor(), new ServiceConnection() {
+                        @Override
+                        public void onServiceConnected(ComponentName componentName,
+                                IBinder iBinder) {
+                            RemoteCollectionItems items;
+                            try {
+                                items = IRemoteViewsFactory.Stub.asInterface(iBinder)
+                                    .getRemoteCollectionItems();
+                            } catch (RemoteException re) {
+                                items = new RemoteCollectionItems.Builder().build();
+                                Log.e(LOG_TAG, "Error getting collection items from the factory",
+                                        re);
+                            } finally {
+                                context.unbindService(this);
+                            }
+
+                            result.complete(items);
+                        }
+
+                        @Override
+                        public void onServiceDisconnected(ComponentName componentName) { }
+                    });
+
+            result.completeOnTimeout(
+                    new RemoteCollectionItems.Builder().build(),
+                    MAX_ADAPTER_CONVERSION_WAITING_TIME_MS, TimeUnit.MILLISECONDS);
+
+            return result;
         }
 
         public void writeToParcel(Parcel out, int flags) {
@@ -1360,10 +1348,7 @@
                 out.writeInt(mIdToUriMapping.keyAt(i));
                 String intentUri = mIdToUriMapping.valueAt(i);
                 out.writeString8(intentUri);
-                RemoteCollectionItems items = mTempUriToFutureMapping.get(intentUri) != null
-                        ? getCollectionItemsFromFuture(mTempUriToFutureMapping.get(intentUri))
-                        : mUriToCollectionMapping.get(intentUri);
-                items.writeToParcel(out, flags, true);
+                mUriToCollectionMapping.get(intentUri).writeToParcel(out, flags, true);
             }
         }
     }
diff --git a/core/java/android/window/SystemPerformanceHinter.java b/core/java/android/window/SystemPerformanceHinter.java
index 5b0d8d1..cc2329fc 100644
--- a/core/java/android/window/SystemPerformanceHinter.java
+++ b/core/java/android/window/SystemPerformanceHinter.java
@@ -20,7 +20,7 @@
 import static android.view.Surface.FRAME_RATE_CATEGORY_DEFAULT;
 import static android.view.Surface.FRAME_RATE_CATEGORY_HIGH;
 import static android.view.SurfaceControl.FRAME_RATE_SELECTION_STRATEGY_OVERRIDE_CHILDREN;
-import static android.view.SurfaceControl.FRAME_RATE_SELECTION_STRATEGY_SELF;
+import static android.view.SurfaceControl.FRAME_RATE_SELECTION_STRATEGY_PROPAGATE;
 
 import android.annotation.IntDef;
 import android.annotation.NonNull;
@@ -303,7 +303,7 @@
             SurfaceControl displaySurfaceControl = mDisplayRootProvider.getRootForDisplay(
                     session.displayId);
             mTransaction.setFrameRateSelectionStrategy(displaySurfaceControl,
-                    FRAME_RATE_SELECTION_STRATEGY_SELF);
+                    FRAME_RATE_SELECTION_STRATEGY_PROPAGATE);
             // smoothSwitchOnly is false to request a higher framerate, even if it means switching
             // the display mode will cause would jank on non-VRR devices because keeping a lower
             // refresh rate would mean a poorer user experience.
diff --git a/core/java/android/window/WindowInfosListenerForTest.java b/core/java/android/window/WindowInfosListenerForTest.java
index 35ce726..34c6399 100644
--- a/core/java/android/window/WindowInfosListenerForTest.java
+++ b/core/java/android/window/WindowInfosListenerForTest.java
@@ -19,6 +19,7 @@
 import android.Manifest;
 import android.annotation.NonNull;
 import android.annotation.RequiresPermission;
+import android.annotation.SuppressLint;
 import android.annotation.TestApi;
 import android.graphics.Matrix;
 import android.graphics.Rect;
@@ -87,6 +88,38 @@
         @NonNull
         public final Matrix transform;
 
+        /**
+         * True if the window is touchable.
+         */
+        @SuppressLint("UnflaggedApi") // The API is only used for tests.
+        public final boolean isTouchable;
+
+        /**
+         * True if the window is focusable.
+         */
+        @SuppressLint("UnflaggedApi") // The API is only used for tests.
+        public final boolean isFocusable;
+
+        /**
+         * True if the window is preventing splitting
+         */
+        @SuppressLint("UnflaggedApi") // The API is only used for tests.
+        public final boolean isPreventSplitting;
+
+        /**
+         * True if the window duplicates touches received to wallpaper.
+         */
+        @SuppressLint("UnflaggedApi") // The API is only used for tests.
+        public final boolean isDuplicateTouchToWallpaper;
+
+        /**
+         * True if the window is listening for when there is a touch DOWN event
+         * occurring outside its touchable bounds. When such an event occurs,
+         * this window will receive a MotionEvent with ACTION_OUTSIDE.
+         */
+        @SuppressLint("UnflaggedApi") // The API is only used for tests.
+        public final boolean isWatchOutsideTouch;
+
         WindowInfo(@NonNull IBinder windowToken, @NonNull String name, int displayId,
                 @NonNull Rect bounds, int inputConfig, @NonNull Matrix transform) {
             this.windowToken = windowToken;
@@ -96,6 +129,14 @@
             this.isTrustedOverlay = (inputConfig & InputConfig.TRUSTED_OVERLAY) != 0;
             this.isVisible = (inputConfig & InputConfig.NOT_VISIBLE) == 0;
             this.transform = transform;
+            this.isTouchable = (inputConfig & InputConfig.NOT_TOUCHABLE) == 0;
+            this.isFocusable = (inputConfig & InputConfig.NOT_FOCUSABLE) == 0;
+            this.isPreventSplitting = (inputConfig
+                            & InputConfig.PREVENT_SPLITTING) != 0;
+            this.isDuplicateTouchToWallpaper = (inputConfig
+                            & InputConfig.DUPLICATE_TOUCH_TO_WALLPAPER) != 0;
+            this.isWatchOutsideTouch = (inputConfig
+                            & InputConfig.WATCH_OUTSIDE_TOUCH) != 0;
         }
 
         @Override
diff --git a/core/java/android/window/flags/window_surfaces.aconfig b/core/java/android/window/flags/window_surfaces.aconfig
index 7f93213..29932f3 100644
--- a/core/java/android/window/flags/window_surfaces.aconfig
+++ b/core/java/android/window/flags/window_surfaces.aconfig
@@ -48,3 +48,11 @@
     is_fixed_read_only: true
     bug: "262477923"
 }
+
+flag {
+    namespace: "window_surfaces"
+    name: "secure_window_state"
+    description: "Move SC secure flag to WindowState level"
+    is_fixed_read_only: true
+    bug: "308662081"
+}
diff --git a/core/java/com/android/internal/accessibility/util/AccessibilityStatsLogUtils.java b/core/java/com/android/internal/accessibility/util/AccessibilityStatsLogUtils.java
index 7c4252e..6b074a6 100644
--- a/core/java/com/android/internal/accessibility/util/AccessibilityStatsLogUtils.java
+++ b/core/java/com/android/internal/accessibility/util/AccessibilityStatsLogUtils.java
@@ -33,6 +33,7 @@
 import static com.android.internal.util.FrameworkStatsLog.ACCESSIBILITY_SHORTCUT_REPORTED__SHORTCUT_TYPE__A11Y_FLOATING_MENU;
 import static com.android.internal.util.FrameworkStatsLog.ACCESSIBILITY_SHORTCUT_REPORTED__SHORTCUT_TYPE__A11Y_GESTURE;
 import static com.android.internal.util.FrameworkStatsLog.ACCESSIBILITY_SHORTCUT_REPORTED__SHORTCUT_TYPE__TRIPLE_TAP;
+import static com.android.internal.util.FrameworkStatsLog.ACCESSIBILITY_SHORTCUT_REPORTED__SHORTCUT_TYPE__TWO_FINGER_TRIPLE_TAP;
 import static com.android.internal.util.FrameworkStatsLog.ACCESSIBILITY_SHORTCUT_REPORTED__SHORTCUT_TYPE__UNKNOWN_TYPE;
 import static com.android.internal.util.FrameworkStatsLog.ACCESSIBILITY_SHORTCUT_REPORTED__SHORTCUT_TYPE__VOLUME_KEY;
 import static com.android.internal.util.FrameworkStatsLog.MAGNIFICATION_USAGE_REPORTED__ACTIVATED_MODE__MAGNIFICATION_ALL;
@@ -131,6 +132,18 @@
     }
 
     /**
+     * Logs magnification that is assigned to the two finger triple tap shortcut. Calls this when
+     * triggering the magnification two finger triple tap shortcut.
+     */
+    public static void logMagnificationTwoFingerTripleTap(boolean enabled) {
+        FrameworkStatsLog.write(FrameworkStatsLog.ACCESSIBILITY_SHORTCUT_REPORTED,
+                MAGNIFICATION_COMPONENT_NAME.flattenToString(),
+                // jean update
+                ACCESSIBILITY_SHORTCUT_REPORTED__SHORTCUT_TYPE__TWO_FINGER_TRIPLE_TAP,
+                convertToLoggingServiceStatus(enabled));
+    }
+
+    /**
      * Logs accessibility feature name that is assigned to the long pressed accessibility button
      * shortcut. Calls this when clicking the long pressed accessibility button shortcut.
      *
diff --git a/core/java/com/android/internal/os/MonotonicClock.java b/core/java/com/android/internal/os/MonotonicClock.java
index d0d2354..661628a 100644
--- a/core/java/com/android/internal/os/MonotonicClock.java
+++ b/core/java/com/android/internal/os/MonotonicClock.java
@@ -50,6 +50,8 @@
     private final Clock mClock;
     private long mTimeshift;
 
+    public static final long UNDEFINED = -1;
+
     public MonotonicClock(File file) {
         mFile = new AtomicFile(file);
         mClock = Clock.SYSTEM_CLOCK;
diff --git a/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl b/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl
index dadeb2b..aab2242 100644
--- a/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl
+++ b/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl
@@ -60,7 +60,8 @@
     @UnsupportedAppUsage(maxTargetSdk = 28)
     void notifyCallForwardingChanged(boolean cfi);
     void notifyCallForwardingChangedForSubscriber(in int subId, boolean cfi);
-    void notifyDataActivityForSubscriber(int phoneId, int subId, int state);
+    void notifyDataActivityForSubscriber(int subId, int state);
+    void notifyDataActivityForSubscriberWithSlot(int phoneId, int subId, int state);
     void notifyDataConnectionForSubscriber(
             int phoneId, int subId, in PreciseDataConnectionState preciseState);
     // Uses CellIdentity which is Parcelable here; will convert to CellLocation in client.
diff --git a/core/jni/Android.bp b/core/jni/Android.bp
index 440a332..f365dbb 100644
--- a/core/jni/Android.bp
+++ b/core/jni/Android.bp
@@ -462,9 +462,4 @@
             ],
         },
     },
-
-    // Workaround Clang LTO crash.
-    lto: {
-        never: true,
-    },
 }
diff --git a/core/proto/android/os/batteryusagestats.proto b/core/proto/android/os/batteryusagestats.proto
index 2b74220..11b367b 100644
--- a/core/proto/android/os/batteryusagestats.proto
+++ b/core/proto/android/os/batteryusagestats.proto
@@ -92,8 +92,24 @@
     message UidBatteryConsumer {
         optional int32 uid = 1;
         optional BatteryConsumerData battery_consumer_data = 2;
-        optional int64 time_in_foreground_millis = 3;
-        optional int64 time_in_background_millis = 4;
+        // DEPRECATED Use time_in_state instead.
+        optional int64 time_in_foreground_millis = 3 [deprecated = true];
+        // DEPRECATED Use time_in_state instead.
+        optional int64 time_in_background_millis = 4 [deprecated = true];
+
+        message TimeInState {
+            enum ProcessState {
+                UNSPECIFIED = 0;
+                FOREGROUND = 1;
+                BACKGROUND = 2;
+                FOREGROUND_SERVICE = 3;
+            }
+
+            optional ProcessState process_state = 1;
+            optional int64 time_in_state_millis = 2;
+        }
+
+        repeated TimeInState time_in_state = 5;
     }
     repeated UidBatteryConsumer uid_battery_consumers = 5;
 
diff --git a/core/res/AndroidManifest.xml b/core/res/AndroidManifest.xml
index 4d208c6..0021640 100644
--- a/core/res/AndroidManifest.xml
+++ b/core/res/AndroidManifest.xml
@@ -7817,6 +7817,13 @@
     <permission android:name="android.permission.RESET_HOTWORD_TRAINING_DATA_EGRESS_COUNT"
                 android:protectionLevel="signature" />
 
+    <!-- @SystemApi Allows an app to track all preparations for a complete factory reset.
+     <p>Protection level: signature|privileged
+     @FlaggedApi("android.permission.flags.factory_reset_prep_permission_apis")
+     @hide -->
+    <permission android:name="android.permission.PREPARE_FACTORY_RESET"
+        android:protectionLevel="signature|privileged" />
+
     <!-- Attribution for Geofencing service. -->
     <attribution android:tag="GeofencingService" android:label="@string/geofencing_service"/>
     <!-- Attribution for Country Detector. -->
diff --git a/core/res/OWNERS b/core/res/OWNERS
index f24c3f5..332ad2a 100644
--- a/core/res/OWNERS
+++ b/core/res/OWNERS
@@ -47,6 +47,10 @@
 # Wear
 per-file res/*-watch/* = file:/WEAR_OWNERS
 
+# Peformance
+per-file res/values/config.xml = file:/PERFORMANCE_OWNERS
+per-file res/values/symbols.xml = file:/PERFORMANCE_OWNERS
+
 # PowerProfile
 per-file res/xml/power_profile.xml = file:/BATTERY_STATS_OWNERS
 per-file res/xml/power_profile_test.xml = file:/BATTERY_STATS_OWNERS
diff --git a/core/res/res/values/config.xml b/core/res/res/values/config.xml
index 6cd6eb4..98897d8 100644
--- a/core/res/res/values/config.xml
+++ b/core/res/res/values/config.xml
@@ -4336,6 +4336,9 @@
     <!-- True if assistant app should be pinned via Pinner Service -->
     <bool name="config_pinnerAssistantApp">false</bool>
 
+    <!-- Bytes that the PinnerService will pin for WebView -->
+    <integer name="config_pinnerWebviewPinBytes">0</integer>
+
     <!-- Number of days preloaded file cache should be preserved on a device before it can be
          deleted -->
     <integer name="config_keepPreloadsMinDays">7</integer>
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index 38f1f67..017688a 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -3385,6 +3385,7 @@
   <java-symbol type="bool" name="config_pinnerCameraApp" />
   <java-symbol type="bool" name="config_pinnerHomeApp" />
   <java-symbol type="bool" name="config_pinnerAssistantApp" />
+  <java-symbol type="integer" name="config_pinnerWebviewPinBytes" />
 
   <java-symbol type="string" name="config_doubleTouchGestureEnableFile" />
 
diff --git a/core/tests/coretests/src/android/window/SystemPerformanceHinterTests.java b/core/tests/coretests/src/android/window/SystemPerformanceHinterTests.java
index 6229530..3147eac 100644
--- a/core/tests/coretests/src/android/window/SystemPerformanceHinterTests.java
+++ b/core/tests/coretests/src/android/window/SystemPerformanceHinterTests.java
@@ -21,7 +21,7 @@
 import static android.view.Surface.FRAME_RATE_CATEGORY_DEFAULT;
 import static android.view.Surface.FRAME_RATE_CATEGORY_HIGH;
 import static android.view.SurfaceControl.FRAME_RATE_SELECTION_STRATEGY_OVERRIDE_CHILDREN;
-import static android.view.SurfaceControl.FRAME_RATE_SELECTION_STRATEGY_SELF;
+import static android.view.SurfaceControl.FRAME_RATE_SELECTION_STRATEGY_PROPAGATE;
 import static android.window.SystemPerformanceHinter.HINT_ADPF;
 import static android.window.SystemPerformanceHinter.HINT_ALL;
 import static android.window.SystemPerformanceHinter.HINT_SF_EARLY_WAKEUP;
@@ -170,7 +170,7 @@
         // Verify we call SF
         verify(mTransaction).setFrameRateSelectionStrategy(
                 eq(mDefaultDisplayRoot),
-                eq(FRAME_RATE_SELECTION_STRATEGY_SELF));
+                eq(FRAME_RATE_SELECTION_STRATEGY_PROPAGATE));
         verify(mTransaction).setFrameRateCategory(
                 eq(mDefaultDisplayRoot),
                 eq(FRAME_RATE_CATEGORY_DEFAULT),
@@ -262,7 +262,7 @@
         // Verify we call SF and perf manager to clean up
         verify(mTransaction).setFrameRateSelectionStrategy(
                 eq(mDefaultDisplayRoot),
-                eq(FRAME_RATE_SELECTION_STRATEGY_SELF));
+                eq(FRAME_RATE_SELECTION_STRATEGY_PROPAGATE));
         verify(mTransaction).setFrameRateCategory(
                 eq(mDefaultDisplayRoot),
                 eq(FRAME_RATE_CATEGORY_DEFAULT),
@@ -283,7 +283,7 @@
             // Verify we call SF and perf manager to clean up
             verify(mTransaction).setFrameRateSelectionStrategy(
                     eq(mDefaultDisplayRoot),
-                    eq(FRAME_RATE_SELECTION_STRATEGY_SELF));
+                    eq(FRAME_RATE_SELECTION_STRATEGY_PROPAGATE));
             verify(mTransaction).setFrameRateCategory(
                     eq(mDefaultDisplayRoot),
                     eq(FRAME_RATE_CATEGORY_DEFAULT),
@@ -334,7 +334,7 @@
         // Verify we call SF and perf manager to clean up
         verify(mTransaction).setFrameRateSelectionStrategy(
                 eq(mDefaultDisplayRoot),
-                eq(FRAME_RATE_SELECTION_STRATEGY_SELF));
+                eq(FRAME_RATE_SELECTION_STRATEGY_PROPAGATE));
         verify(mTransaction).setFrameRateCategory(
                 eq(mDefaultDisplayRoot),
                 eq(FRAME_RATE_CATEGORY_DEFAULT),
@@ -385,7 +385,7 @@
         session1.close();
         verify(mTransaction).setFrameRateSelectionStrategy(
                 eq(mDefaultDisplayRoot),
-                eq(FRAME_RATE_SELECTION_STRATEGY_SELF));
+                eq(FRAME_RATE_SELECTION_STRATEGY_PROPAGATE));
         verify(mTransaction).setFrameRateCategory(
                 eq(mDefaultDisplayRoot),
                 eq(FRAME_RATE_CATEGORY_DEFAULT),
@@ -410,7 +410,7 @@
                 anyInt());
         verify(mTransaction).setFrameRateSelectionStrategy(
                 eq(mSecondaryDisplayRoot),
-                eq(FRAME_RATE_SELECTION_STRATEGY_SELF));
+                eq(FRAME_RATE_SELECTION_STRATEGY_PROPAGATE));
         verify(mTransaction).setFrameRateCategory(
                 eq(mSecondaryDisplayRoot),
                 eq(FRAME_RATE_CATEGORY_DEFAULT),
diff --git a/data/etc/services.core.protolog.json b/data/etc/services.core.protolog.json
index f19acbe..d36ac39 100644
--- a/data/etc/services.core.protolog.json
+++ b/data/etc/services.core.protolog.json
@@ -991,12 +991,6 @@
       "group": "WM_DEBUG_WINDOW_INSETS",
       "at": "com\/android\/server\/wm\/InsetsSourceProvider.java"
     },
-    "-1176488860": {
-      "message": "SURFACE isSecure=%b: %s",
-      "level": "INFO",
-      "group": "WM_SHOW_TRANSACTIONS",
-      "at": "com\/android\/server\/wm\/WindowSurfaceController.java"
-    },
     "-1164930508": {
       "message": "Moving to RESUMED: %s (starting new instance) callers=%s",
       "level": "VERBOSE",
@@ -3277,6 +3271,12 @@
       "group": "WM_DEBUG_ORIENTATION",
       "at": "com\/android\/server\/wm\/ActivityRecord.java"
     },
+    "810599500": {
+      "message": "SURFACE isSecure=%b: %s",
+      "level": "INFO",
+      "group": "WM_SHOW_TRANSACTIONS",
+      "at": "com\/android\/server\/wm\/WindowState.java"
+    },
     "829434921": {
       "message": "Draw state now committed in %s",
       "level": "VERBOSE",
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskView.java
index ef8393c..35a1fa0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskView.java
@@ -151,7 +151,14 @@
 
     @Override
     public void setResizeBgColor(SurfaceControl.Transaction t, int bgColor) {
-        runOnViewThread(() -> setResizeBackgroundColor(t, bgColor));
+        if (mHandler.getLooper().isCurrentThread()) {
+            // We can only use the transaction if it can updated synchronously, otherwise the tx
+            // will be applied immediately after but also used/updated on the view thread which
+            // will lead to a race and/or crash
+            runOnViewThread(() -> setResizeBackgroundColor(t, bgColor));
+        } else {
+            runOnViewThread(() -> setResizeBackgroundColor(bgColor));
+        }
     }
 
     /**
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestHandler.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestHandler.java
new file mode 100644
index 0000000..b91d6f9
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestHandler.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell;
+
+import android.os.Handler;
+import android.os.Looper;
+import android.os.Message;
+
+
+/**
+ * Basic test handler that immediately executes anything that is posted on it.
+ */
+public class TestHandler extends Handler {
+    public TestHandler(Looper looper) {
+        super(looper);
+    }
+
+    @Override
+    public boolean sendMessageAtTime(Message msg, long uptimeMillis) {
+        dispatchMessage(msg);
+        return true;
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java
index 4afb29e..d7c4610 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java
@@ -40,6 +40,7 @@
 import android.app.ActivityOptions;
 import android.app.PendingIntent;
 import android.content.Context;
+import android.graphics.Color;
 import android.graphics.Insets;
 import android.graphics.Rect;
 import android.graphics.Region;
@@ -58,6 +59,7 @@
 
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
+import com.android.wm.shell.TestHandler;
 import com.android.wm.shell.common.HandlerExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.common.SyncTransactionQueue.TransactionRunnable;
@@ -92,8 +94,7 @@
     Transitions mTransitions;
     @Mock
     Looper mViewLooper;
-    @Mock
-    Handler mViewHandler;
+    TestHandler mViewHandler;
 
     SurfaceSession mSession;
     SurfaceControl mLeash;
@@ -112,7 +113,7 @@
 
         mContext = getContext();
         doReturn(true).when(mViewLooper).isCurrentThread();
-        doReturn(mViewLooper).when(mViewHandler).getLooper();
+        mViewHandler = spy(new TestHandler(mViewLooper));
 
         mTaskInfo = new ActivityManager.RunningTaskInfo();
         mTaskInfo.token = mToken;
@@ -668,4 +669,24 @@
         mTaskViewTaskController.onTaskInfoChanged(mTaskInfo);
         verify(mViewHandler).post(any());
     }
+
+    @Test
+    public void testSetResizeBgOnSameUiThread_expectUsesTransaction() {
+        SurfaceControl.Transaction tx = mock(SurfaceControl.Transaction.class);
+        mTaskView = spy(mTaskView);
+        mTaskView.setResizeBgColor(tx, Color.BLUE);
+        verify(mViewHandler, never()).post(any());
+        verify(mTaskView, never()).setResizeBackgroundColor(eq(Color.BLUE));
+        verify(mTaskView).setResizeBackgroundColor(eq(tx), eq(Color.BLUE));
+    }
+
+    @Test
+    public void testSetResizeBgOnDifferentUiThread_expectDoesNotUseTransaction() {
+        doReturn(false).when(mViewLooper).isCurrentThread();
+        SurfaceControl.Transaction tx = mock(SurfaceControl.Transaction.class);
+        mTaskView = spy(mTaskView);
+        mTaskView.setResizeBgColor(tx, Color.BLUE);
+        verify(mViewHandler).post(any());
+        verify(mTaskView).setResizeBackgroundColor(eq(Color.BLUE));
+    }
 }
diff --git a/packages/PackageInstaller/src/com/android/packageinstaller/DeleteStagedFileOnResult.java b/packages/PackageInstaller/src/com/android/packageinstaller/DeleteStagedFileOnResult.java
index 19d74b3..7b17cbd 100644
--- a/packages/PackageInstaller/src/com/android/packageinstaller/DeleteStagedFileOnResult.java
+++ b/packages/PackageInstaller/src/com/android/packageinstaller/DeleteStagedFileOnResult.java
@@ -16,8 +16,6 @@
 
 package com.android.packageinstaller;
 
-import static android.content.Intent.CATEGORY_LAUNCHER;
-
 import static com.android.packageinstaller.PackageInstallerActivity.EXTRA_STAGED_SESSION_ID;
 
 import android.app.Activity;
@@ -47,9 +45,6 @@
     protected void onActivityResult(int requestCode, int resultCode, Intent data) {
         setResult(resultCode, data);
         finish();
-        if (data != null && data.hasCategory(CATEGORY_LAUNCHER)) {
-            startActivity(data);
-        }
     }
 
     @Override
diff --git a/packages/PackageInstaller/src/com/android/packageinstaller/InstallSuccess.java b/packages/PackageInstaller/src/com/android/packageinstaller/InstallSuccess.java
index 9af88c3b..215ead3 100644
--- a/packages/PackageInstaller/src/com/android/packageinstaller/InstallSuccess.java
+++ b/packages/PackageInstaller/src/com/android/packageinstaller/InstallSuccess.java
@@ -18,6 +18,7 @@
 
 import android.app.Activity;
 import android.app.AlertDialog;
+import android.content.ActivityNotFoundException;
 import android.content.DialogInterface;
 import android.content.Intent;
 import android.content.pm.ApplicationInfo;
@@ -120,7 +121,12 @@
         Button launchButton = mDialog.getButton(DialogInterface.BUTTON_POSITIVE);
         if (enabled) {
             launchButton.setOnClickListener(view -> {
-                setResult(Activity.RESULT_OK, mLaunchIntent);
+                try {
+                    startActivity(mLaunchIntent.addFlags(
+                        Intent.FLAG_ACTIVITY_CLEAR_TOP | Intent.FLAG_ACTIVITY_SINGLE_TOP));
+                } catch (ActivityNotFoundException | SecurityException e) {
+                    Log.e(LOG_TAG, "Could not start activity", e);
+                }
                 finish();
             });
         } else {
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index a9dc145..1762065 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -114,6 +114,13 @@
 }
 
 flag {
+    name: "unfold_animation_background_progress"
+    namespace: "systemui"
+    description: "Moves unfold animation progress calculation to a background thread"
+    bug: "277879146"
+}
+
+flag {
     name: "qs_new_pipeline"
     namespace: "systemui"
     description: "Use the new pipeline for Quick Settings. Should have no behavior changes."
diff --git a/packages/SystemUI/res/drawable/notification_material_bg.xml b/packages/SystemUI/res/drawable/notification_material_bg.xml
index 9c08f5e..355e75d 100644
--- a/packages/SystemUI/res/drawable/notification_material_bg.xml
+++ b/packages/SystemUI/res/drawable/notification_material_bg.xml
@@ -18,7 +18,7 @@
 <layer-list xmlns:android="http://schemas.android.com/apk/res/android"
         xmlns:androidprv="http://schemas.android.com/apk/prv/res/android"
         android:color="?android:attr/colorControlHighlight">
-    <item android:id="@+id/notification_background_color_layer">
+    <item>
         <shape>
             <solid android:color="?androidprv:attr/materialColorSurfaceContainerHigh" />
         </shape>
diff --git a/packages/SystemUI/res/layout/bluetooth_tile_dialog.xml b/packages/SystemUI/res/layout/bluetooth_tile_dialog.xml
index af29cad..50241cd 100644
--- a/packages/SystemUI/res/layout/bluetooth_tile_dialog.xml
+++ b/packages/SystemUI/res/layout/bluetooth_tile_dialog.xml
@@ -111,107 +111,57 @@
                 app:layout_constraintStart_toStartOf="parent"
                 app:layout_constraintEnd_toEndOf="parent"
                 app:layout_constraintTop_toBottomOf="@+id/bluetooth_toggle"
-                app:layout_constraintBottom_toTopOf="@+id/see_all_text" />
+                app:layout_constraintBottom_toTopOf="@+id/see_all_button" />
 
-            <androidx.constraintlayout.widget.Group
-                android:id="@+id/see_all_layout_group"
-                android:layout_width="wrap_content"
-                android:layout_height="wrap_content"
-                android:visibility="gone"
-                app:constraint_referenced_ids="ic_arrow,see_all_text" />
-
-            <View
-                android:id="@+id/see_all_clickable_row"
+            <Button
+                android:id="@+id/see_all_button"
+                style="@style/BluetoothTileDialog.Device"
+                android:paddingEnd="0dp"
+                android:paddingStart="20dp"
+                android:background="@drawable/bluetooth_tile_dialog_bg_off"
                 android:layout_width="0dp"
-                android:layout_height="0dp"
+                android:layout_height="64dp"
                 android:contentDescription="@string/accessibility_bluetooth_device_settings_see_all"
                 app:layout_constraintStart_toStartOf="parent"
                 app:layout_constraintEnd_toEndOf="parent"
                 app:layout_constraintTop_toBottomOf="@+id/device_list"
-                app:layout_constraintBottom_toTopOf="@+id/pair_new_device_text" />
-
-            <ImageView
-                android:id="@+id/ic_arrow"
-                android:layout_marginStart="36dp"
-                android:layout_width="24dp"
-                android:layout_height="24dp"
-                android:importantForAccessibility="no"
-                android:gravity="center_vertical"
-                android:src="@drawable/ic_arrow_forward"
-                app:layout_constraintBottom_toTopOf="@+id/pair_new_device_text"
-                app:layout_constraintStart_toStartOf="parent"
-                app:layout_constraintEnd_toStartOf="@id/see_all_text"
-                app:layout_constraintTop_toBottomOf="@id/device_list" />
-
-            <TextView
-                android:id="@+id/see_all_text"
-                style="@style/BluetoothTileDialog.Device"
-                android:layout_width="0dp"
-                android:layout_height="64dp"
-                android:maxLines="1"
-                android:ellipsize="end"
-                android:gravity="center_vertical"
-                android:importantForAccessibility="no"
-                android:clickable="false"
-                android:layout_marginStart="0dp"
-                android:paddingStart="20dp"
+                app:layout_constraintBottom_toTopOf="@+id/pair_new_device_button"
+                android:drawableStart="@drawable/ic_arrow_forward"
+                android:drawablePadding="20dp"
+                android:drawableTint="?android:attr/textColorPrimary"
                 android:text="@string/see_all_bluetooth_devices"
                 android:textSize="14sp"
                 android:textAppearance="@style/TextAppearance.Dialog.Title"
-                app:layout_constraintBottom_toTopOf="@+id/pair_new_device_text"
-                app:layout_constraintStart_toEndOf="@+id/ic_arrow"
-                app:layout_constraintTop_toBottomOf="@id/device_list"
-                app:layout_constraintEnd_toEndOf="parent" />
+                android:textDirection="locale"
+                android:textAlignment="viewStart"
+                android:maxLines="1"
+                android:ellipsize="end"
+                android:visibility="gone" />
 
-            <androidx.constraintlayout.widget.Group
-                android:id="@+id/pair_new_device_layout_group"
-                android:layout_width="wrap_content"
-                android:layout_height="wrap_content"
-                android:visibility="gone"
-                app:constraint_referenced_ids="ic_add,pair_new_device_text" />
-
-            <View
-                android:id="@+id/pair_new_device_clickable_row"
+            <Button
+                android:id="@+id/pair_new_device_button"
+                style="@style/BluetoothTileDialog.Device"
+                android:paddingEnd="0dp"
+                android:paddingStart="20dp"
+                android:background="@drawable/bluetooth_tile_dialog_bg_off"
                 android:layout_width="0dp"
-                android:layout_height="0dp"
+                android:layout_height="64dp"
                 android:contentDescription="@string/accessibility_bluetooth_device_settings_pair_new_device"
                 app:layout_constraintStart_toStartOf="parent"
                 app:layout_constraintEnd_toEndOf="parent"
-                app:layout_constraintTop_toBottomOf="@+id/see_all_text"
-                app:layout_constraintBottom_toTopOf="@+id/done_button" />
-
-            <ImageView
-                android:id="@+id/ic_add"
-                android:layout_width="24dp"
-                android:layout_height="24dp"
-                android:layout_marginStart="36dp"
-                android:gravity="center_vertical"
-                android:importantForAccessibility="no"
-                android:src="@drawable/ic_add"
-                app:layout_constraintBottom_toTopOf="@id/done_button"
-                app:layout_constraintStart_toStartOf="parent"
-                app:layout_constraintEnd_toStartOf="@id/pair_new_device_text"
-                app:layout_constraintTop_toBottomOf="@id/see_all_text"
-                android:tint="?android:attr/textColorPrimary" />
-
-            <TextView
-                android:id="@+id/pair_new_device_text"
-                style="@style/BluetoothTileDialog.Device"
-                android:layout_width="0dp"
-                android:layout_height="64dp"
-                android:maxLines="1"
-                android:ellipsize="end"
-                android:gravity="center_vertical"
-                android:importantForAccessibility="no"
-                android:clickable="false"
-                android:layout_marginStart="0dp"
-                android:paddingStart="20dp"
+                app:layout_constraintTop_toBottomOf="@+id/see_all_button"
+                app:layout_constraintBottom_toTopOf="@+id/done_button"
+                android:drawableStart="@drawable/ic_add"
+                android:drawablePadding="20dp"
+                android:drawableTint="?android:attr/textColorPrimary"
                 android:text="@string/pair_new_bluetooth_devices"
                 android:textSize="14sp"
                 android:textAppearance="@style/TextAppearance.Dialog.Title"
-                app:layout_constraintStart_toEndOf="@+id/ic_add"
-                app:layout_constraintTop_toBottomOf="@id/see_all_text"
-                app:layout_constraintEnd_toEndOf="parent" />
+                android:textDirection="locale"
+                android:textAlignment="viewStart"
+                android:maxLines="1"
+                android:ellipsize="end"
+                android:visibility="gone" />
 
             <Button
                 android:id="@+id/done_button"
@@ -227,7 +177,7 @@
                 android:maxLines="1"
                 android:text="@string/inline_done_button"
                 app:layout_constraintEnd_toEndOf="parent"
-                app:layout_constraintTop_toBottomOf="@id/pair_new_device_text"
+                app:layout_constraintTop_toBottomOf="@id/pair_new_device_button"
                 app:layout_constraintBottom_toBottomOf="parent" />
         </androidx.constraintlayout.widget.ConstraintLayout>
     </androidx.core.widget.NestedScrollView>
diff --git a/packages/SystemUI/shared/src/com/android/systemui/unfold/system/ActivityManagerActivityTypeProvider.kt b/packages/SystemUI/shared/src/com/android/systemui/unfold/system/ActivityManagerActivityTypeProvider.kt
index c9e57b4..b33f6fa 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/unfold/system/ActivityManagerActivityTypeProvider.kt
+++ b/packages/SystemUI/shared/src/com/android/systemui/unfold/system/ActivityManagerActivityTypeProvider.kt
@@ -32,8 +32,7 @@
     override val isHomeActivity: Boolean?
         get() = _isHomeActivity
 
-    private var _isHomeActivity: Boolean? = null
-
+    @Volatile private var _isHomeActivity: Boolean? = null
 
     override fun init() {
         _isHomeActivity = activityManager.isOnHomeActivity()
diff --git a/packages/SystemUI/shared/src/com/android/systemui/unfold/system/DeviceStateManagerFoldProvider.kt b/packages/SystemUI/shared/src/com/android/systemui/unfold/system/DeviceStateManagerFoldProvider.kt
index 3b8d318..baa8889 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/unfold/system/DeviceStateManagerFoldProvider.kt
+++ b/packages/SystemUI/shared/src/com/android/systemui/unfold/system/DeviceStateManagerFoldProvider.kt
@@ -18,18 +18,19 @@
 import android.hardware.devicestate.DeviceStateManager
 import com.android.systemui.unfold.updates.FoldProvider
 import com.android.systemui.unfold.updates.FoldProvider.FoldCallback
+import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.Executor
 import javax.inject.Inject
 import javax.inject.Singleton
 
 @Singleton
-class DeviceStateManagerFoldProvider @Inject constructor(
-    private val deviceStateManager: DeviceStateManager,
-    private val context: Context
-) : FoldProvider {
+class DeviceStateManagerFoldProvider
+@Inject
+constructor(private val deviceStateManager: DeviceStateManager, private val context: Context) :
+    FoldProvider {
 
-    private val callbacks: MutableMap<FoldCallback,
-            DeviceStateManager.DeviceStateCallback> = hashMapOf()
+    private val callbacks =
+        ConcurrentHashMap<FoldCallback, DeviceStateManager.DeviceStateCallback>()
 
     override fun registerCallback(callback: FoldCallback, executor: Executor) {
         val listener = FoldStateListener(context, callback)
@@ -39,13 +40,9 @@
 
     override fun unregisterCallback(callback: FoldCallback) {
         val listener = callbacks.remove(callback)
-        listener?.let {
-            deviceStateManager.unregisterCallback(it)
-        }
+        listener?.let { deviceStateManager.unregisterCallback(it) }
     }
 
-    private inner class FoldStateListener(
-        context: Context,
-        listener: FoldCallback
-    ) : DeviceStateManager.FoldStateListener(context, { listener.onFoldUpdated(it) })
+    private inner class FoldStateListener(context: Context, listener: FoldCallback) :
+        DeviceStateManager.FoldStateListener(context, { listener.onFoldUpdated(it) })
 }
diff --git a/packages/SystemUI/shared/src/com/android/systemui/unfold/system/SystemUnfoldSharedModule.kt b/packages/SystemUI/shared/src/com/android/systemui/unfold/system/SystemUnfoldSharedModule.kt
index 7b67e3f..7af9917 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/unfold/system/SystemUnfoldSharedModule.kt
+++ b/packages/SystemUI/shared/src/com/android/systemui/unfold/system/SystemUnfoldSharedModule.kt
@@ -15,24 +15,29 @@
 package com.android.systemui.unfold.system
 
 import android.os.Handler
+import android.os.HandlerThread
+import android.os.Looper
+import android.os.Process
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.dagger.qualifiers.UiBackground
 import com.android.systemui.unfold.config.ResourceUnfoldTransitionConfig
 import com.android.systemui.unfold.config.UnfoldTransitionConfig
-import com.android.systemui.unfold.dagger.UnfoldSingleThreadBg
+import com.android.systemui.unfold.dagger.UnfoldBg
 import com.android.systemui.unfold.dagger.UnfoldMain
+import com.android.systemui.unfold.dagger.UnfoldSingleThreadBg
 import com.android.systemui.unfold.updates.FoldProvider
 import com.android.systemui.unfold.util.CurrentActivityTypeProvider
 import dagger.Binds
 import dagger.Module
+import dagger.Provides
 import java.util.concurrent.Executor
+import javax.inject.Singleton
 
 /**
- * Dagger module with system-only dependencies for the unfold animation.
- * The code that is used to calculate unfold transition progress
- * depends on some hidden APIs that are not available in normal
- * apps. In order to re-use this code and use alternative implementations
- * of these classes in other apps and hidden APIs here.
+ * Dagger module with system-only dependencies for the unfold animation. The code that is used to
+ * calculate unfold transition progress depends on some hidden APIs that are not available in normal
+ * apps. In order to re-use this code and use alternative implementations of these classes in other
+ * apps and hidden APIs here.
  */
 @Module
 abstract class SystemUnfoldSharedModule {
@@ -61,4 +66,22 @@
     @Binds
     @UnfoldSingleThreadBg
     abstract fun backgroundExecutor(@UiBackground executor: Executor): Executor
+
+    companion object {
+        @Provides
+        @UnfoldBg
+        @Singleton
+        fun unfoldBgProgressHandler(@UnfoldBg looper: Looper): Handler {
+            return Handler(looper)
+        }
+
+        @Provides
+        @UnfoldBg
+        @Singleton
+        fun provideBgLooper(): Looper {
+            return HandlerThread("UnfoldBg", Process.THREAD_PRIORITY_FOREGROUND)
+                .apply { start() }
+                .looper
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/biometrics/AuthContainerView.java b/packages/SystemUI/src/com/android/systemui/biometrics/AuthContainerView.java
index 1ac4163..ab23564 100644
--- a/packages/SystemUI/src/com/android/systemui/biometrics/AuthContainerView.java
+++ b/packages/SystemUI/src/com/android/systemui/biometrics/AuthContainerView.java
@@ -30,6 +30,7 @@
 import android.graphics.PixelFormat;
 import android.hardware.biometrics.BiometricAuthenticator.Modality;
 import android.hardware.biometrics.BiometricConstants;
+import android.hardware.biometrics.BiometricManager.Authenticators;
 import android.hardware.biometrics.PromptInfo;
 import android.hardware.face.FaceSensorPropertiesInternal;
 import android.hardware.fingerprint.FingerprintSensorPropertiesInternal;
@@ -168,7 +169,7 @@
     // HAT received from LockSettingsService when credential is verified.
     @Nullable private byte[] mCredentialAttestation;
 
-    // TODO(b/287311775): remove when legacy prompt is replaced
+    // TODO(b/313469218): remove when legacy prompt is replaced
     @Deprecated
     static class Config {
         Context mContext;
@@ -220,6 +221,9 @@
             mHandler.postDelayed(() -> {
                 addCredentialView(false /* animatePanel */, true /* animateContents */);
             }, mConfig.mSkipAnimation ? 0 : ANIMATE_CREDENTIAL_START_DELAY_MS);
+
+            // TODO(b/313469218): Remove Config
+            mConfig.mPromptInfo.setAuthenticators(Authenticators.DEVICE_CREDENTIAL);
         }
 
         @Override
diff --git a/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptIconViewModel.kt b/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptIconViewModel.kt
index 11a5d8b..3defec5 100644
--- a/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptIconViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptIconViewModel.kt
@@ -485,8 +485,7 @@
     ): Int =
         if (isPendingConfirmation) {
             when (sensorType) {
-                FingerprintSensorType.POWER_BUTTON ->
-                    R.string.security_settings_sfps_enroll_find_sensor_message
+                FingerprintSensorType.POWER_BUTTON -> -1
                 else -> R.string.fingerprint_dialog_authenticated_confirmation
             }
         } else if (isAuthenticating || isAuthenticated) {
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java b/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
index d9e0629..e7b8773 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
@@ -19,6 +19,7 @@
 import com.android.systemui.BootCompleteCacheImpl;
 import com.android.systemui.CoreStartable;
 import com.android.systemui.Dependency;
+import com.android.systemui.Flags;
 import com.android.systemui.InitController;
 import com.android.systemui.SystemUIAppComponentFactoryBase;
 import com.android.systemui.dagger.qualifiers.PerUser;
@@ -35,6 +36,7 @@
 import com.android.systemui.unfold.FoldStateLoggingProvider;
 import com.android.systemui.unfold.SysUIUnfoldComponent;
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider;
+import com.android.systemui.unfold.dagger.UnfoldBg;
 import com.android.systemui.unfold.progress.UnfoldTransitionProgressForwarder;
 import com.android.wm.shell.back.BackAnimation;
 import com.android.wm.shell.bubbles.Bubbles;
@@ -144,7 +146,15 @@
         getConnectingDisplayViewModel().init();
         getFoldStateLoggingProvider().ifPresent(FoldStateLoggingProvider::init);
         getFoldStateLogger().ifPresent(FoldStateLogger::init);
-        getUnfoldTransitionProgressProvider()
+
+        Optional<UnfoldTransitionProgressProvider> unfoldTransitionProgressProvider;
+
+        if (Flags.unfoldAnimationBackgroundProgress()) {
+            unfoldTransitionProgressProvider = getBgUnfoldTransitionProgressProvider();
+        } else {
+            unfoldTransitionProgressProvider = getUnfoldTransitionProgressProvider();
+        }
+        unfoldTransitionProgressProvider
                 .ifPresent(
                         (progressProvider) ->
                                 getUnfoldTransitionProgressForwarder()
@@ -170,7 +180,14 @@
     ContextComponentHelper getContextComponentHelper();
 
     /**
-     * Creates a UnfoldTransitionProgressProvider.
+     * Creates a UnfoldTransitionProgressProvider that calculates progress in the background.
+     */
+    @SysUISingleton
+    @UnfoldBg
+    Optional<UnfoldTransitionProgressProvider> getBgUnfoldTransitionProgressProvider();
+
+    /**
+     * Creates a UnfoldTransitionProgressProvider that calculates progress in the main thread.
      */
     @SysUISingleton
     Optional<UnfoldTransitionProgressProvider> getUnfoldTransitionProgressProvider();
diff --git a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt
index 715fb17..4cddb9c 100644
--- a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt
@@ -103,8 +103,11 @@
                 initialValue = false,
             )
 
-    // Authenticated by a TrustAgent like trusted device, location, etc or by face auth.
-    private val passivelyAuthenticated =
+    /**
+     * Whether the user is currently authenticated by a TrustAgent like trusted device, location,
+     * etc., or by face auth.
+     */
+    private val isPassivelyAuthenticated =
         merge(
                 trustRepository.isCurrentUserTrusted,
                 deviceEntryFaceAuthRepository.isAuthenticated,
@@ -117,25 +120,31 @@
      * mechanism like face or trust manager. This returns `false` whenever the lockscreen has been
      * dismissed.
      *
+     * A value of `null` is meaningless and is used as placeholder while the actual value is still
+     * being loaded in the background.
+     *
      * Note: `true` doesn't mean the lockscreen is visible. It may be occluded or covered by other
      * UI.
      */
-    val canSwipeToEnter =
+    val canSwipeToEnter: StateFlow<Boolean?> =
         combine(
                 // This is true when the user has chosen to show the lockscreen but has not made it
                 // secure.
                 authenticationInteractor.authenticationMethod.map {
                     it == AuthenticationMethodModel.None && repository.isLockscreenEnabled()
                 },
-                passivelyAuthenticated,
+                isPassivelyAuthenticated,
                 isDeviceEntered
-            ) { isSwipeAuthMethod, passivelyAuthenticated, isDeviceEntered ->
-                (isSwipeAuthMethod || passivelyAuthenticated) && !isDeviceEntered
+            ) { isSwipeAuthMethod, isPassivelyAuthenticated, isDeviceEntered ->
+                (isSwipeAuthMethod || isPassivelyAuthenticated) && !isDeviceEntered
             }
             .stateIn(
                 scope = applicationScope,
                 started = SharingStarted.Eagerly,
-                initialValue = false,
+                // Starts as null to prevent downstream collectors from falsely assuming that the
+                // user can or cannot swipe to enter the device while the real value is being loaded
+                // from upstream data sources.
+                initialValue = null,
             )
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
index 26c5ea6..c93b8e1 100644
--- a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
@@ -29,7 +29,6 @@
 import com.android.app.tracing.traceSection
 import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.display.data.DisplayEvent
 import com.android.systemui.util.Compile
@@ -93,7 +92,7 @@
 constructor(
     private val displayManager: DisplayManager,
     @Background backgroundHandler: Handler,
-    @Application applicationScope: CoroutineScope,
+    @Background bgApplicationScope: CoroutineScope,
     @Background backgroundCoroutineDispatcher: CoroutineDispatcher
 ) : DisplayRepository {
     private val allDisplayEvents: Flow<DisplayEvent> =
@@ -141,8 +140,7 @@
     private val enabledDisplays =
         allDisplayEvents
             .map { getDisplays() }
-            .flowOn(backgroundCoroutineDispatcher)
-            .shareIn(applicationScope, started = SharingStarted.WhileSubscribed(), replay = 1)
+            .shareIn(bgApplicationScope, started = SharingStarted.WhileSubscribed(), replay = 1)
 
     override val displays: Flow<Set<Display>> = enabledDisplays
 
@@ -203,9 +201,8 @@
             }
             .distinctUntilChanged()
             .debugLog("connectedDisplayIds")
-            .flowOn(backgroundCoroutineDispatcher)
             .stateIn(
-                applicationScope,
+                bgApplicationScope,
                 started = SharingStarted.WhileSubscribed(),
                 // The initial value is set to empty, but connected displays are gathered as soon as
                 // the flow starts being collected. This is to ensure the call to get displays (an
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/LifecycleScreenStatusProvider.kt b/packages/SystemUI/src/com/android/systemui/keyguard/LifecycleScreenStatusProvider.kt
index f9b89b1..7354cfc 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/LifecycleScreenStatusProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/LifecycleScreenStatusProvider.kt
@@ -18,6 +18,7 @@
 import com.android.systemui.unfold.updates.screen.ScreenStatusProvider
 import com.android.systemui.unfold.updates.screen.ScreenStatusProvider.ScreenListener
 import com.android.app.tracing.traceSection
+import java.util.concurrent.CopyOnWriteArrayList
 import javax.inject.Inject
 import javax.inject.Singleton
 
@@ -29,7 +30,7 @@
         screenLifecycle.addObserver(this)
     }
 
-    private val listeners: MutableList<ScreenListener> = mutableListOf()
+    private val listeners: MutableList<ScreenListener> = CopyOnWriteArrayList()
 
     override fun removeCallback(listener: ScreenListener) {
         listeners.remove(listener)
diff --git a/packages/SystemUI/src/com/android/systemui/qs/external/CustomTileStatePersister.kt b/packages/SystemUI/src/com/android/systemui/qs/external/CustomTileStatePersister.kt
index a321eef..6f5dea3 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/external/CustomTileStatePersister.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/external/CustomTileStatePersister.kt
@@ -18,17 +18,19 @@
 
 import android.content.ComponentName
 import android.content.Context
+import android.content.SharedPreferences
 import android.service.quicksettings.Tile
 import android.util.Log
 import com.android.internal.annotations.VisibleForTesting
+import javax.inject.Inject
 import org.json.JSONException
 import org.json.JSONObject
-import javax.inject.Inject
 
 data class TileServiceKey(val componentName: ComponentName, val user: Int) {
     private val string = "${componentName.flattenToString()}:$user"
     override fun toString() = string
 }
+
 private const val STATE = "state"
 private const val LABEL = "label"
 private const val SUBTITLE = "subtitle"
@@ -44,12 +46,7 @@
  * It persists the state from a [Tile] necessary to present the view in the same state when
  * retrieved, with the exception of the icon.
  */
-class CustomTileStatePersister @Inject constructor(context: Context) {
-    companion object {
-        private const val FILE_NAME = "custom_tiles_state"
-    }
-
-    private val sharedPreferences = context.getSharedPreferences(FILE_NAME, 0)
+interface CustomTileStatePersister {
 
     /**
      * Read the state from [SharedPreferences].
@@ -58,7 +55,31 @@
      *
      * Any fields that have not been saved will be set to `null`
      */
-    fun readState(key: TileServiceKey): Tile? {
+    fun readState(key: TileServiceKey): Tile?
+    /**
+     * Persists the state into [SharedPreferences].
+     *
+     * The implementation does not store fields that are `null` or icons.
+     */
+    fun persistState(key: TileServiceKey, tile: Tile)
+    /**
+     * Removes the state for a given tile, user pair.
+     *
+     * Used when the tile is removed by the user.
+     */
+    fun removeState(key: TileServiceKey)
+}
+
+// TODO(b/299909989) Merge this class into into CustomTileRepository
+class CustomTileStatePersisterImpl @Inject constructor(context: Context) :
+    CustomTileStatePersister {
+    companion object {
+        private const val FILE_NAME = "custom_tiles_state"
+    }
+
+    private val sharedPreferences: SharedPreferences = context.getSharedPreferences(FILE_NAME, 0)
+
+    override fun readState(key: TileServiceKey): Tile? {
         val state = sharedPreferences.getString(key.toString(), null) ?: return null
         return try {
             readTileFromString(state)
@@ -68,23 +89,13 @@
         }
     }
 
-    /**
-     * Persists the state into [SharedPreferences].
-     *
-     * The implementation does not store fields that are `null` or icons.
-     */
-    fun persistState(key: TileServiceKey, tile: Tile) {
+    override fun persistState(key: TileServiceKey, tile: Tile) {
         val state = writeToString(tile)
 
         sharedPreferences.edit().putString(key.toString(), state).apply()
     }
 
-    /**
-     * Removes the state for a given tile, user pair.
-     *
-     * Used when the tile is removed by the user.
-     */
-    fun removeState(key: TileServiceKey) {
+    override fun removeState(key: TileServiceKey) {
         sharedPreferences.edit().remove(key.toString()).apply()
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/di/QSTilesModule.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/di/QSTilesModule.kt
index 94137c8..4a34276 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/di/QSTilesModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/di/QSTilesModule.kt
@@ -16,6 +16,8 @@
 
 package com.android.systemui.qs.tiles.di
 
+import com.android.systemui.qs.external.CustomTileStatePersister
+import com.android.systemui.qs.external.CustomTileStatePersisterImpl
 import com.android.systemui.qs.tiles.base.actions.QSTileIntentUserInputHandler
 import com.android.systemui.qs.tiles.base.actions.QSTileIntentUserInputHandlerImpl
 import com.android.systemui.qs.tiles.impl.custom.di.CustomTileComponent
@@ -52,4 +54,7 @@
     fun bindQSTileIntentUserInputHandler(
         impl: QSTileIntentUserInputHandlerImpl
     ): QSTileIntentUserInputHandler
+
+    @Binds
+    fun bindCustomTileStatePersister(impl: CustomTileStatePersisterImpl): CustomTileStatePersister
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialog.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialog.kt
index 5bdb592..db3cf0f 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialog.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialog.kt
@@ -81,10 +81,8 @@
     private lateinit var toggleView: Switch
     private lateinit var subtitleTextView: TextView
     private lateinit var doneButton: View
-    private lateinit var seeAllViewGroup: View
-    private lateinit var pairNewDeviceViewGroup: View
-    private lateinit var seeAllRow: View
-    private lateinit var pairNewDeviceRow: View
+    private lateinit var seeAllButton: View
+    private lateinit var pairNewDeviceButton: View
     private lateinit var deviceListView: RecyclerView
 
     override fun onCreate(savedInstanceState: Bundle?) {
@@ -99,10 +97,8 @@
         toggleView = requireViewById(R.id.bluetooth_toggle)
         subtitleTextView = requireViewById(R.id.bluetooth_tile_dialog_subtitle) as TextView
         doneButton = requireViewById(R.id.done_button)
-        seeAllViewGroup = requireViewById(R.id.see_all_layout_group)
-        pairNewDeviceViewGroup = requireViewById(R.id.pair_new_device_layout_group)
-        seeAllRow = requireViewById(R.id.see_all_clickable_row)
-        pairNewDeviceRow = requireViewById(R.id.pair_new_device_clickable_row)
+        seeAllButton = requireViewById(R.id.see_all_button)
+        pairNewDeviceButton = requireViewById(R.id.pair_new_device_button)
         deviceListView = requireViewById<RecyclerView>(R.id.device_list)
 
         setupToggle()
@@ -110,8 +106,8 @@
 
         subtitleTextView.text = context.getString(subtitleResIdInitialValue)
         doneButton.setOnClickListener { dismiss() }
-        seeAllRow.setOnClickListener { bluetoothTileDialogCallback.onSeeAllClicked(it) }
-        pairNewDeviceRow.setOnClickListener {
+        seeAllButton.setOnClickListener { bluetoothTileDialogCallback.onSeeAllClicked(it) }
+        pairNewDeviceButton.setOnClickListener {
             bluetoothTileDialogCallback.onPairNewDeviceClicked(it)
         }
     }
@@ -134,8 +130,8 @@
             }
             if (isActive) {
                 deviceItemAdapter.refreshDeviceItemList(deviceItem) {
-                    seeAllViewGroup.visibility = if (showSeeAll) VISIBLE else GONE
-                    pairNewDeviceViewGroup.visibility = if (showPairNewDevice) VISIBLE else GONE
+                    seeAllButton.visibility = if (showSeeAll) VISIBLE else GONE
+                    pairNewDeviceButton.visibility = if (showPairNewDevice) VISIBLE else GONE
                     lastUiUpdateMs = systemClock.elapsedRealtime()
                     lastItemRow = itemRow
                     logger.logDeviceUiUpdate(lastUiUpdateMs - start)
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogViewModel.kt
index 34c2aba..5d5e747 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogViewModel.kt
@@ -20,7 +20,6 @@
 import android.content.Intent
 import android.os.Bundle
 import android.view.View
-import androidx.annotation.VisibleForTesting
 import com.android.internal.jank.InteractionJankMonitor
 import com.android.internal.logging.UiEventLogger
 import com.android.systemui.animation.DialogCuj
@@ -40,6 +39,8 @@
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.channels.produce
 import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
@@ -63,26 +64,25 @@
 
     private var job: Job? = null
 
-    @VisibleForTesting internal var dialog: BluetoothTileDialog? = null
-
     /**
      * Shows the dialog.
      *
      * @param context The context in which the dialog is displayed.
      * @param view The view from which the dialog is shown.
      */
+    @kotlinx.coroutines.ExperimentalCoroutinesApi
     fun showDialog(context: Context, view: View?) {
-        dismissDialog()
-
-        var updateDeviceItemJob: Job? = null
-        var updateDialogUiJob: Job? = null
+        cancelJob()
 
         job =
             coroutineScope.launch(mainDispatcher) {
-                dialog = createBluetoothTileDialog(context)
+                var updateDeviceItemJob: Job?
+                var updateDialogUiJob: Job? = null
+                val dialog = createBluetoothTileDialog(context)
+
                 view?.let {
                     dialogLaunchAnimator.showFromView(
-                        dialog!!,
+                        dialog,
                         it,
                         animateBackgroundBoundsChange = true,
                         cuj =
@@ -92,9 +92,8 @@
                             )
                     )
                 }
-                    ?: dialog!!.show()
+                    ?: dialog.show()
 
-                updateDeviceItemJob?.cancel()
                 updateDeviceItemJob = launch {
                     deviceItemInteractor.updateDeviceItems(context, DeviceFetchTrigger.FIRST_LOAD)
                 }
@@ -102,7 +101,7 @@
                 bluetoothStateInteractor.bluetoothStateUpdate
                     .filterNotNull()
                     .onEach {
-                        dialog!!.onBluetoothStateUpdated(it, getSubtitleResId(it))
+                        dialog.onBluetoothStateUpdated(it, getSubtitleResId(it))
                         updateDeviceItemJob?.cancel()
                         updateDeviceItemJob = launch {
                             deviceItemInteractor.updateDeviceItems(
@@ -129,7 +128,7 @@
                     .onEach {
                         updateDialogUiJob?.cancel()
                         updateDialogUiJob = launch {
-                            dialog?.onDeviceItemUpdated(
+                            dialog.onDeviceItemUpdated(
                                 it.take(MAX_DEVICE_ITEM_ENTRY),
                                 showSeeAll = it.size > MAX_DEVICE_ITEM_ENTRY,
                                 showPairNewDevice = bluetoothStateInteractor.isBluetoothEnabled
@@ -138,15 +137,15 @@
                     }
                     .launchIn(this)
 
-                dialog!!
-                    .bluetoothStateToggle
+                dialog.bluetoothStateToggle
                     .onEach { bluetoothStateInteractor.isBluetoothEnabled = it }
                     .launchIn(this)
 
-                dialog!!
-                    .deviceItemClick
+                dialog.deviceItemClick
                     .onEach { deviceItemInteractor.updateDeviceItemOnClick(it) }
                     .launchIn(this)
+
+                produce<Unit> { awaitClose { dialog.cancel() } }
             }
     }
 
@@ -161,7 +160,7 @@
                 logger,
                 context
             )
-            .apply { SystemUIDialog.registerDismissListener(this) { dismissDialog() } }
+            .apply { SystemUIDialog.registerDismissListener(this) { cancelJob() } }
     }
 
     override fun onDeviceItemGearClicked(deviceItem: DeviceItem, view: View) {
@@ -188,15 +187,13 @@
         startSettingsActivity(Intent(ACTION_PAIR_NEW_DEVICE), view)
     }
 
-    private fun dismissDialog() {
+    private fun cancelJob() {
         job?.cancel()
         job = null
-        dialog?.dismiss()
-        dialog = null
     }
 
     private fun startSettingsActivity(intent: Intent, view: View) {
-        dialog?.run {
+        if (job?.isActive == true) {
             intent.flags = Intent.FLAG_ACTIVITY_NEW_TASK or Intent.FLAG_ACTIVITY_CLEAR_TOP
             activityStarter.postStartActivityDismissingKeyguard(
                 intent,
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/DeviceItemInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/DeviceItemInteractor.kt
index 76fbf8e..fcd45a6 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/DeviceItemInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/bluetooth/DeviceItemInteractor.kt
@@ -168,26 +168,30 @@
         )
     }
 
-    internal fun updateDeviceItemOnClick(deviceItem: DeviceItem) {
-        logger.logDeviceClick(deviceItem.cachedBluetoothDevice.address, deviceItem.type)
+    internal suspend fun updateDeviceItemOnClick(deviceItem: DeviceItem) {
+        withContext(backgroundDispatcher) {
+            logger.logDeviceClick(deviceItem.cachedBluetoothDevice.address, deviceItem.type)
 
-        deviceItem.cachedBluetoothDevice.apply {
-            when (deviceItem.type) {
-                DeviceItemType.ACTIVE_MEDIA_BLUETOOTH_DEVICE -> {
-                    disconnect()
-                    uiEventLogger.log(BluetoothTileDialogUiEvent.ACTIVE_DEVICE_DISCONNECT)
-                }
-                DeviceItemType.AVAILABLE_MEDIA_BLUETOOTH_DEVICE -> {
-                    setActive()
-                    uiEventLogger.log(BluetoothTileDialogUiEvent.CONNECTED_DEVICE_SET_ACTIVE)
-                }
-                DeviceItemType.CONNECTED_BLUETOOTH_DEVICE -> {
-                    disconnect()
-                    uiEventLogger.log(BluetoothTileDialogUiEvent.CONNECTED_OTHER_DEVICE_DISCONNECT)
-                }
-                DeviceItemType.SAVED_BLUETOOTH_DEVICE -> {
-                    connect()
-                    uiEventLogger.log(BluetoothTileDialogUiEvent.SAVED_DEVICE_CONNECT)
+            deviceItem.cachedBluetoothDevice.apply {
+                when (deviceItem.type) {
+                    DeviceItemType.ACTIVE_MEDIA_BLUETOOTH_DEVICE -> {
+                        disconnect()
+                        uiEventLogger.log(BluetoothTileDialogUiEvent.ACTIVE_DEVICE_DISCONNECT)
+                    }
+                    DeviceItemType.AVAILABLE_MEDIA_BLUETOOTH_DEVICE -> {
+                        setActive()
+                        uiEventLogger.log(BluetoothTileDialogUiEvent.CONNECTED_DEVICE_SET_ACTIVE)
+                    }
+                    DeviceItemType.CONNECTED_BLUETOOTH_DEVICE -> {
+                        disconnect()
+                        uiEventLogger.log(
+                            BluetoothTileDialogUiEvent.CONNECTED_OTHER_DEVICE_DISCONNECT
+                        )
+                    }
+                    DeviceItemType.SAVED_BLUETOOTH_DEVICE -> {
+                        connect()
+                        uiEventLogger.log(BluetoothTileDialogUiEvent.SAVED_DEVICE_CONNECT)
+                    }
                 }
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/commons/TileExt.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/commons/TileExt.kt
new file mode 100644
index 0000000..869f6f32
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/commons/TileExt.kt
@@ -0,0 +1,50 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.tiles.impl.custom.commons
+
+import android.service.quicksettings.Tile
+
+fun Tile.copy(): Tile =
+    Tile().also {
+        it.icon = icon
+        it.label = label
+        it.subtitle = subtitle
+        it.contentDescription = contentDescription
+        it.stateDescription = stateDescription
+        it.activityLaunchForClick = activityLaunchForClick
+        it.state = state
+    }
+
+fun Tile.setFrom(otherTile: Tile) {
+    if (otherTile.icon != null) {
+        icon = otherTile.icon
+    }
+    if (otherTile.customLabel != null) {
+        label = otherTile.customLabel
+    }
+    if (otherTile.subtitle != null) {
+        subtitle = otherTile.subtitle
+    }
+    if (otherTile.contentDescription != null) {
+        contentDescription = otherTile.contentDescription
+    }
+    if (otherTile.stateDescription != null) {
+        stateDescription = otherTile.stateDescription
+    }
+    activityLaunchForClick = otherTile.activityLaunchForClick
+    state = otherTile.state
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/data/repository/CustomTileRepository.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/data/repository/CustomTileRepository.kt
new file mode 100644
index 0000000..ca5302e
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/data/repository/CustomTileRepository.kt
@@ -0,0 +1,196 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.tiles.impl.custom.data.repository
+
+import android.graphics.drawable.Icon
+import android.os.UserHandle
+import android.service.quicksettings.Tile
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.qs.external.CustomTileStatePersister
+import com.android.systemui.qs.external.TileServiceKey
+import com.android.systemui.qs.pipeline.shared.TileSpec
+import com.android.systemui.qs.tiles.impl.custom.commons.copy
+import com.android.systemui.qs.tiles.impl.custom.commons.setFrom
+import com.android.systemui.qs.tiles.impl.custom.data.entity.CustomTileDefaults
+import com.android.systemui.qs.tiles.impl.di.QSTileScope
+import javax.inject.Inject
+import kotlin.coroutines.CoroutineContext
+import kotlinx.coroutines.channels.BufferOverflow
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+import kotlinx.coroutines.withContext
+
+/**
+ * Repository store the [Tile] associated with the custom tile. It lives on [QSTileScope] which
+ * allows it to survive service rebinding. Given that, it provides the last received state when
+ * connected again.
+ */
+interface CustomTileRepository {
+
+    /**
+     * Restores the [Tile] if it's [isPersistable]. Restored [Tile] will be available via [getTile]
+     * (but there is no guarantee that restoration is synchronous) and emitted in [getTiles] for a
+     * corresponding [user].
+     */
+    suspend fun restoreForTheUserIfNeeded(user: UserHandle, isPersistable: Boolean)
+
+    /** Returns [Tile] updates for a [user]. */
+    fun getTiles(user: UserHandle): Flow<Tile>
+
+    /**
+     * Return current [Tile] for a [user] or null if the [user] doesn't match currently cached one.
+     * Suspending until [getTiles] returns something is a way to wait for this to become available.
+     *
+     * @throws IllegalStateException when there is no current tile.
+     */
+    fun getTile(user: UserHandle): Tile?
+
+    /**
+     * Updates tile with the non-null values from [newTile]. Overwrites the current cache when
+     * [user] differs from the cached one. [isPersistable] tile will be persisted to be possibly
+     * loaded when the [restoreForTheUserIfNeeded].
+     */
+    suspend fun updateWithTile(
+        user: UserHandle,
+        newTile: Tile,
+        isPersistable: Boolean,
+    )
+
+    /**
+     * Updates tile with the values from [defaults]. Overwrites the current cache when [user]
+     * differs from the cached one. [isPersistable] tile will be persisted to be possibly loaded
+     * when the [restoreForTheUserIfNeeded].
+     */
+    suspend fun updateWithDefaults(
+        user: UserHandle,
+        defaults: CustomTileDefaults,
+        isPersistable: Boolean,
+    )
+}
+
+@QSTileScope
+class CustomTileRepositoryImpl
+@Inject
+constructor(
+    private val tileSpec: TileSpec.CustomTileSpec,
+    private val customTileStatePersister: CustomTileStatePersister,
+    @Background private val backgroundContext: CoroutineContext,
+) : CustomTileRepository {
+
+    private val tileUpdateMutex = Mutex()
+    private val tileWithUserState =
+        MutableSharedFlow<TileWithUser>(onBufferOverflow = BufferOverflow.DROP_OLDEST, replay = 1)
+
+    override suspend fun restoreForTheUserIfNeeded(user: UserHandle, isPersistable: Boolean) {
+        if (isPersistable && getCurrentTileWithUser()?.user != user) {
+            withContext(backgroundContext) {
+                customTileStatePersister.readState(user.getKey())?.let {
+                    updateWithTile(
+                        user,
+                        it,
+                        true,
+                    )
+                }
+            }
+        }
+    }
+
+    override fun getTiles(user: UserHandle): Flow<Tile> =
+        tileWithUserState.filter { it.user == user }.map { it.tile }
+
+    override fun getTile(user: UserHandle): Tile? {
+        val tileWithUser =
+            getCurrentTileWithUser() ?: throw IllegalStateException("Tile is not set")
+        return if (tileWithUser.user == user) {
+            tileWithUser.tile
+        } else {
+            null
+        }
+    }
+
+    override suspend fun updateWithTile(
+        user: UserHandle,
+        newTile: Tile,
+        isPersistable: Boolean,
+    ) = updateTile(user, isPersistable) { setFrom(newTile) }
+
+    override suspend fun updateWithDefaults(
+        user: UserHandle,
+        defaults: CustomTileDefaults,
+        isPersistable: Boolean,
+    ) {
+        if (defaults is CustomTileDefaults.Result) {
+            updateTile(user, isPersistable) {
+                // Update the icon if it's not set or is the default icon.
+                val updateIcon = (icon == null || icon.isResourceEqual(defaults.icon))
+                if (updateIcon) {
+                    icon = defaults.icon
+                }
+                setDefaultLabel(defaults.label)
+            }
+        }
+    }
+
+    private suspend fun updateTile(
+        user: UserHandle,
+        isPersistable: Boolean,
+        update: Tile.() -> Unit
+    ): Unit =
+        tileUpdateMutex.withLock {
+            val currentTileWithUser = getCurrentTileWithUser()
+            val tileToUpdate =
+                if (currentTileWithUser?.user == user) {
+                    currentTileWithUser.tile.copy()
+                } else {
+                    Tile()
+                }
+            tileToUpdate.update()
+            if (isPersistable) {
+                withContext(backgroundContext) {
+                    customTileStatePersister.persistState(user.getKey(), tileToUpdate)
+                }
+            }
+            tileWithUserState.tryEmit(TileWithUser(user, tileToUpdate))
+        }
+
+    private fun getCurrentTileWithUser(): TileWithUser? = tileWithUserState.replayCache.lastOrNull()
+
+    /** Compare two icons, only works for resources. */
+    private fun Icon.isResourceEqual(icon2: Icon?): Boolean {
+        if (icon2 == null) {
+            return false
+        }
+        if (this === icon2) {
+            return true
+        }
+        if (type != Icon.TYPE_RESOURCE || icon2.type != Icon.TYPE_RESOURCE) {
+            return false
+        }
+        if (resId != icon2.resId) {
+            return false
+        }
+        return resPackage == icon2.resPackage
+    }
+
+    private fun UserHandle.getKey() = TileServiceKey(tileSpec.componentName, this.identifier)
+
+    private data class TileWithUser(val user: UserHandle, val tile: Tile)
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/di/CustomTileModule.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/di/CustomTileModule.kt
index 83767aa..d956fde 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/di/CustomTileModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/di/CustomTileModule.kt
@@ -24,6 +24,8 @@
 import com.android.systemui.qs.tiles.impl.custom.CustomTileUserActionInteractor
 import com.android.systemui.qs.tiles.impl.custom.data.repository.CustomTileDefaultsRepository
 import com.android.systemui.qs.tiles.impl.custom.data.repository.CustomTileDefaultsRepositoryImpl
+import com.android.systemui.qs.tiles.impl.custom.data.repository.CustomTileRepository
+import com.android.systemui.qs.tiles.impl.custom.data.repository.CustomTileRepositoryImpl
 import com.android.systemui.qs.tiles.impl.custom.di.bound.CustomTileBoundComponent
 import com.android.systemui.qs.tiles.impl.custom.domain.entity.CustomTileDataModel
 import dagger.Binds
@@ -50,4 +52,6 @@
     fun bindCustomTileDefaultsRepository(
         impl: CustomTileDefaultsRepositoryImpl
     ): CustomTileDefaultsRepository
+
+    @Binds fun bindCustomTileRepository(impl: CustomTileRepositoryImpl): CustomTileRepository
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/domain/interactor/CustomTileInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/domain/interactor/CustomTileInteractor.kt
new file mode 100644
index 0000000..351bba5
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/custom/domain/interactor/CustomTileInteractor.kt
@@ -0,0 +1,112 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.tiles.impl.custom.domain.interactor
+
+import android.os.UserHandle
+import android.service.quicksettings.Tile
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.qs.external.TileServiceManager
+import com.android.systemui.qs.tiles.impl.custom.data.repository.CustomTileDefaultsRepository
+import com.android.systemui.qs.tiles.impl.custom.data.repository.CustomTileRepository
+import com.android.systemui.qs.tiles.impl.custom.di.bound.CustomTileBoundScope
+import javax.inject.Inject
+import kotlin.coroutines.CoroutineContext
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.BufferOverflow
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.firstOrNull
+import kotlinx.coroutines.flow.flowOn
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
+
+/** Manages updates of the [Tile] assigned for the current custom tile. */
+@CustomTileBoundScope
+class CustomTileInteractor
+@Inject
+constructor(
+    private val user: UserHandle,
+    private val defaultsRepository: CustomTileDefaultsRepository,
+    private val customTileRepository: CustomTileRepository,
+    private val tileServiceManager: TileServiceManager,
+    @CustomTileBoundScope private val boundScope: CoroutineScope,
+    @Background private val backgroundContext: CoroutineContext,
+) {
+
+    private val tileUpdates =
+        MutableSharedFlow<Tile>(replay = 1, onBufferOverflow = BufferOverflow.DROP_OLDEST)
+
+    /** [Tile] updates. [updateTile] to emit a new one. */
+    val tiles: Flow<Tile>
+        get() = customTileRepository.getTiles(user)
+
+    /**
+     * Current [Tile]
+     *
+     * @throws IllegalStateException when the repository stores a tile for another user. This means
+     *   the tile hasn't been updated for the current user. Can happen when this is accessed before
+     *   [init] returns.
+     */
+    val tile: Tile
+        get() =
+            customTileRepository.getTile(user)
+                ?: throw IllegalStateException("Attempt to get a tile for a wrong user")
+
+    /**
+     * Initializes the repository for the current user. Suspends until it's safe to call [tile]
+     * which needs at least one of the following:
+     * - defaults are loaded;
+     * - receive tile update in [updateTile];
+     * - restoration happened for a persisted tile.
+     */
+    suspend fun init() {
+        launchUpdates()
+        customTileRepository.restoreForTheUserIfNeeded(user, tileServiceManager.isActiveTile)
+        // Suspend to make sure it gets the tile from one of the sources: restoration, defaults, or
+        // tile update.
+        customTileRepository.getTiles(user).firstOrNull()
+    }
+
+    private fun launchUpdates() {
+        tileUpdates
+            .onEach {
+                customTileRepository.updateWithTile(
+                    user,
+                    it,
+                    tileServiceManager.isActiveTile,
+                )
+            }
+            .flowOn(backgroundContext)
+            .launchIn(boundScope)
+        defaultsRepository
+            .defaults(user)
+            .onEach {
+                customTileRepository.updateWithDefaults(
+                    user,
+                    it,
+                    tileServiceManager.isActiveTile,
+                )
+            }
+            .flowOn(backgroundContext)
+            .launchIn(boundScope)
+    }
+
+    /** Updates current [Tile]. Emits a new event in [tiles]. */
+    fun updateTile(newTile: Tile) {
+        tileUpdates.tryEmit(newTile)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
index f3f9c91..1c5330e 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
@@ -146,19 +146,21 @@
                     isAnySimLocked -> {
                         switchToScene(
                             targetSceneKey = SceneKey.Bouncer,
-                            loggingReason = "Need to authenticate locked sim card."
+                            loggingReason = "Need to authenticate locked SIM card."
                         )
                     }
-                    isUnlocked && !canSwipeToEnter -> {
+                    isUnlocked && canSwipeToEnter == false -> {
                         switchToScene(
                             targetSceneKey = SceneKey.Gone,
-                            loggingReason = "Sim cards are unlocked."
+                            loggingReason = "All SIM cards unlocked and device already" +
+                                " unlocked and lockscreen doesn't require a swipe to dismiss."
                         )
                     }
                     else -> {
                         switchToScene(
                             targetSceneKey = SceneKey.Lockscreen,
-                            loggingReason = "Sim cards are unlocked."
+                            loggingReason = "All SIM cards unlocked and device still locked" +
+                                " or lockscreen still requires a swipe to dismiss."
                         )
                     }
                 }
@@ -205,11 +207,17 @@
                             //    when the user is passively authenticated, the false value here
                             //    when the unlock state changes indicates this is an active
                             //    authentication attempt.
-                            if (isBypassEnabled || !canSwipeToEnter)
-                                SceneKey.Gone to
-                                    "device has been unlocked on lockscreen with either " +
-                                        "bypass enabled or using an active authentication mechanism"
-                            else null
+                            when {
+                                isBypassEnabled ->
+                                    SceneKey.Gone to
+                                        "device has been unlocked on lockscreen with bypass" +
+                                            " enabled"
+                                canSwipeToEnter == false ->
+                                    SceneKey.Gone to
+                                        "device has been unlocked on lockscreen using an active" +
+                                            " authentication mechanism"
+                                else -> null
+                            }
                         // Not on lockscreen or bouncer, so remain in the current scene.
                         else -> null
                     }
@@ -232,7 +240,7 @@
                 } else {
                     val canSwipeToEnter = deviceEntryInteractor.canSwipeToEnter.value
                     val isUnlocked = deviceEntryInteractor.isUnlocked.value
-                    if (isUnlocked && !canSwipeToEnter) {
+                    if (isUnlocked && canSwipeToEnter == false) {
                         switchToScene(
                             targetSceneKey = SceneKey.Gone,
                             loggingReason =
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModel.kt b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModel.kt
index 2a071de..0065db3 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModel.kt
@@ -66,10 +66,10 @@
 
     private fun upDestinationSceneKey(
         isUnlocked: Boolean,
-        canSwipeToDismiss: Boolean,
+        canSwipeToDismiss: Boolean?,
     ): SceneKey {
         return when {
-            canSwipeToDismiss -> SceneKey.Lockscreen
+            canSwipeToDismiss == true -> SceneKey.Lockscreen
             isUnlocked -> SceneKey.Gone
             else -> SceneKey.Lockscreen
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractor.kt
index e90ddf9..31893b4 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractor.kt
@@ -15,48 +15,37 @@
 
 package com.android.systemui.statusbar.notification.domain.interactor
 
-import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.statusbar.notification.collection.render.NotifStats
 import com.android.systemui.statusbar.notification.data.repository.ActiveNotificationListRepository
 import com.android.systemui.statusbar.notification.shared.ActiveNotificationGroupModel
 import com.android.systemui.statusbar.notification.shared.ActiveNotificationModel
 import javax.inject.Inject
-import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 
 class ActiveNotificationsInteractor
 @Inject
 constructor(
     private val repository: ActiveNotificationListRepository,
-    @Background private val backgroundDispatcher: CoroutineDispatcher,
 ) {
     /** Notifications actively presented to the user in the notification stack, in order. */
     val topLevelRepresentativeNotifications: Flow<List<ActiveNotificationModel>> =
-        repository.activeNotifications
-            .map { store ->
-                store.renderList.map { key ->
-                    val entry =
-                        store[key]
-                            ?: error(
-                                "Could not find notification with key $key in active notif store."
-                            )
-                    when (entry) {
-                        is ActiveNotificationGroupModel -> entry.summary
-                        is ActiveNotificationModel -> entry
-                    }
+        repository.activeNotifications.map { store ->
+            store.renderList.map { key ->
+                val entry =
+                    store[key]
+                        ?: error("Could not find notification with key $key in active notif store.")
+                when (entry) {
+                    is ActiveNotificationGroupModel -> entry.summary
+                    is ActiveNotificationModel -> entry
                 }
             }
-            .flowOn(backgroundDispatcher)
+        }
 
     /** Are any notifications being actively presented in the notification stack? */
     val areAnyNotificationsPresent: Flow<Boolean> =
-        repository.activeNotifications
-            .map { it.renderList.isNotEmpty() }
-            .distinctUntilChanged()
-            .flowOn(backgroundDispatcher)
+        repository.activeNotifications.map { it.renderList.isNotEmpty() }.distinctUntilChanged()
 
     /**
      * The same as [areAnyNotificationsPresent], but without flows, for easy access in synchronous
@@ -70,7 +59,6 @@
         repository.notifStats
             .map { it.hasClearableAlertingNotifs || it.hasClearableSilentNotifs }
             .distinctUntilChanged()
-            .flowOn(backgroundDispatcher)
 
     fun setNotifStats(notifStats: NotifStats) {
         repository.notifStats.value = notifStats
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationsKeyguardInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationsKeyguardInteractor.kt
index 73341db..87b8e55 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationsKeyguardInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationsKeyguardInteractor.kt
@@ -15,24 +15,19 @@
  */
 package com.android.systemui.statusbar.notification.domain.interactor
 
-import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.statusbar.notification.data.repository.NotificationsKeyguardViewStateRepository
 import javax.inject.Inject
-import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.flowOn
 
 /** Domain logic pertaining to notifications on the keyguard. */
 class NotificationsKeyguardInteractor
 @Inject
 constructor(
     repository: NotificationsKeyguardViewStateRepository,
-    @Background backgroundDispatcher: CoroutineDispatcher,
 ) {
     /** Is a pulse expansion occurring? */
-    val isPulseExpanding: Flow<Boolean> = repository.isPulseExpanding.flowOn(backgroundDispatcher)
+    val isPulseExpanding: Flow<Boolean> = repository.isPulseExpanding
 
     /** Are notifications fully hidden from view? */
-    val areNotificationsFullyHidden: Flow<Boolean> =
-        repository.areNotificationsFullyHidden.flowOn(backgroundDispatcher)
+    val areNotificationsFullyHidden: Flow<Boolean> = repository.areNotificationsFullyHidden
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationBackgroundView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationBackgroundView.java
index 64f61d9..8eda96f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationBackgroundView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationBackgroundView.java
@@ -27,7 +27,6 @@
 import android.graphics.drawable.LayerDrawable;
 import android.graphics.drawable.RippleDrawable;
 import android.util.AttributeSet;
-import android.util.Log;
 import android.view.View;
 
 import androidx.annotation.NonNull;
@@ -43,11 +42,9 @@
  * A view that can be used for both the dimmed and normal background of an notification.
  */
 public class NotificationBackgroundView extends View implements Dumpable {
-    private static final String TAG = "NotificationBackgroundView";
 
     private final boolean mDontModifyCorners;
     private Drawable mBackground;
-    private Drawable mBackgroundDrawableToTint;
     private int mClipTopAmount;
     private int mClipBottomAmount;
     private int mTintColor;
@@ -134,7 +131,6 @@
             unscheduleDrawable(mBackground);
         }
         mBackground = background;
-        mBackgroundDrawableToTint = findBackgroundDrawableToTint(mBackground);
         mRippleColor = null;
         mBackground.mutate();
         if (mBackground != null) {
@@ -148,46 +144,25 @@
         invalidate();
     }
 
-    // setCustomBackground should be called from ActivatableNotificationView.initBackground
-    // with R.drawable.notification_material_bg, which is a layer-list with a lower layer
-    // for the background color (annotated with an ID so we can find it) and an upper layer
-    // to blend in the stateful @color/notification_overlay_color.
-    //
-    // If the notification is tinted, we want to set a tint list on *just that lower layer* that
-    // will replace the default materialColorSurfaceContainerHigh *without* wiping out the stateful
-    // tints in the upper layer that make the hovered and pressed states visible.
-    //
-    // This function fishes that lower layer out, or makes a fuss in logcat if it can't find it.
-    private @Nullable Drawable findBackgroundDrawableToTint(@Nullable Drawable background) {
-        if (background == null) {
-            return null;
-        }
-
-        if (!(background instanceof LayerDrawable)) {
-            Log.wtf(TAG, "background is not a LayerDrawable: " + background);
-            return background;
-        }
-
-        final Drawable backgroundColorLayer = ((LayerDrawable) background).findDrawableByLayerId(
-                R.id.notification_background_color_layer);
-
-        if (backgroundColorLayer == null) {
-            Log.wtf(TAG, "background is missing background color layer: " + background);
-            return background;
-        }
-
-        return backgroundColorLayer;
-    }
-
     public void setCustomBackground(int drawableResId) {
         final Drawable d = mContext.getDrawable(drawableResId);
         setCustomBackground(d);
     }
 
     public void setTint(int tintColor) {
-        mBackgroundDrawableToTint.setTint(tintColor);
-        mBackgroundDrawableToTint.setTintMode(PorterDuff.Mode.SRC_ATOP);
+        if (tintColor != 0) {
+            ColorStateList stateList = new ColorStateList(new int[][]{
+                    new int[]{com.android.internal.R.attr.state_pressed},
+                    new int[]{com.android.internal.R.attr.state_hovered},
+                    new int[]{}},
 
+                    new int[]{tintColor, tintColor, tintColor}
+            );
+            mBackground.setTintMode(PorterDuff.Mode.SRC_ATOP);
+            mBackground.setTintList(stateList);
+        } else {
+            mBackground.setTintList(null);
+        }
         mTintColor = tintColor;
         invalidate();
     }
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/SysUIUnfoldModule.kt b/packages/SystemUI/src/com/android/systemui/unfold/SysUIUnfoldModule.kt
index 8ae093a..10fc83c 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/SysUIUnfoldModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/SysUIUnfoldModule.kt
@@ -20,6 +20,7 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.shade.NotificationPanelUnfoldAnimationController
 import com.android.systemui.statusbar.phone.StatusBarMoveFromCenterAnimationController
+import com.android.systemui.unfold.dagger.UnfoldBg
 import com.android.systemui.unfold.util.NaturalRotationUnfoldProgressProvider
 import com.android.systemui.unfold.util.ScopedUnfoldTransitionProgressProvider
 import com.android.systemui.unfold.util.UnfoldKeyguardVisibilityManager
@@ -33,10 +34,7 @@
 import javax.inject.Named
 import javax.inject.Scope
 
-@Scope
-@MustBeDocumented
-@Retention(AnnotationRetention.RUNTIME)
-annotation class SysUIUnfoldScope
+@Scope @MustBeDocumented @Retention(AnnotationRetention.RUNTIME) annotation class SysUIUnfoldScope
 
 /**
  * Creates an injectable [SysUIUnfoldComponent] that provides objects that have been scoped with
@@ -55,20 +53,21 @@
     @Provides
     @SysUISingleton
     fun provideSysUIUnfoldComponent(
-            provider: Optional<UnfoldTransitionProgressProvider>,
-            rotationProvider: Optional<NaturalRotationUnfoldProgressProvider>,
-            @Named(UNFOLD_STATUS_BAR) scopedProvider:
-            Optional<ScopedUnfoldTransitionProgressProvider>,
-            unfoldLatencyTracker: Lazy<UnfoldLatencyTracker>,
-            factory: SysUIUnfoldComponent.Factory
+        provider: Optional<UnfoldTransitionProgressProvider>,
+        rotationProvider: Optional<NaturalRotationUnfoldProgressProvider>,
+        @Named(UNFOLD_STATUS_BAR) scopedProvider: Optional<ScopedUnfoldTransitionProgressProvider>,
+        @UnfoldBg bgProvider: Optional<UnfoldTransitionProgressProvider>,
+        unfoldLatencyTracker: Lazy<UnfoldLatencyTracker>,
+        factory: SysUIUnfoldComponent.Factory
     ): Optional<SysUIUnfoldComponent> {
         val p1 = provider.getOrNull()
         val p2 = rotationProvider.getOrNull()
         val p3 = scopedProvider.getOrNull()
-        return if (p1 == null || p2 == null || p3 == null) {
+        val p4 = bgProvider.getOrNull()
+        return if (p1 == null || p2 == null || p3 == null || p4 == null) {
             Optional.empty()
         } else {
-            Optional.of(factory.create(p1, p2, p3, unfoldLatencyTracker.get()))
+            Optional.of(factory.create(p1, p2, p3, p4, unfoldLatencyTracker.get()))
         }
     }
 }
@@ -76,13 +75,15 @@
 @SysUIUnfoldScope
 @Subcomponent
 interface SysUIUnfoldComponent {
+
     @Subcomponent.Factory
     interface Factory {
         fun create(
-                @BindsInstance p1: UnfoldTransitionProgressProvider,
-                @BindsInstance p2: NaturalRotationUnfoldProgressProvider,
-                @BindsInstance p3: ScopedUnfoldTransitionProgressProvider,
-                @BindsInstance p4: UnfoldLatencyTracker,
+            @BindsInstance p1: UnfoldTransitionProgressProvider,
+            @BindsInstance p2: NaturalRotationUnfoldProgressProvider,
+            @BindsInstance p3: ScopedUnfoldTransitionProgressProvider,
+            @BindsInstance @UnfoldBg p4: UnfoldTransitionProgressProvider,
+            @BindsInstance p5: UnfoldLatencyTracker,
         ): SysUIUnfoldComponent
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
index 36a1e8a..b72c6f1 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
@@ -35,6 +35,9 @@
 import android.view.SurfaceSession
 import android.view.WindowManager
 import android.view.WindowlessWindowManager
+import com.android.app.tracing.traceSection
+import com.android.keyguard.logging.ScrimLogger
+import com.android.systemui.Flags.unfoldAnimationBackgroundProgress
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.flags.Flags
@@ -45,16 +48,16 @@
 import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation.AddOverlayReason.FOLD
 import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation.AddOverlayReason.UNFOLD
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener
+import com.android.systemui.unfold.dagger.UnfoldBg
 import com.android.systemui.unfold.updates.RotationChangeProvider
 import com.android.systemui.unfold.util.ScaleAwareTransitionProgressProvider.Companion.areAnimationsEnabled
 import com.android.systemui.util.concurrency.ThreadFactory
-import com.android.app.tracing.traceSection
-import com.android.keyguard.logging.ScrimLogger
 import com.android.wm.shell.displayareahelper.DisplayAreaHelper
 import java.util.Optional
 import java.util.concurrent.Executor
 import java.util.function.Consumer
 import javax.inject.Inject
+import javax.inject.Provider
 
 @SysUIUnfoldScope
 class UnfoldLightRevealOverlayAnimation
@@ -65,11 +68,14 @@
     private val deviceStateManager: DeviceStateManager,
     private val contentResolver: ContentResolver,
     private val displayManager: DisplayManager,
-    private val unfoldTransitionProgressProvider: UnfoldTransitionProgressProvider,
+    @UnfoldBg
+    private val unfoldTransitionBgProgressProvider: Provider<UnfoldTransitionProgressProvider>,
+    private val unfoldTransitionProgressProvider: Provider<UnfoldTransitionProgressProvider>,
     private val displayAreaHelper: Optional<DisplayAreaHelper>,
     @Main private val executor: Executor,
     private val threadFactory: ThreadFactory,
-    private val rotationChangeProvider: RotationChangeProvider,
+    @UnfoldBg private val rotationChangeProvider: RotationChangeProvider,
+    @UnfoldBg private val unfoldProgressHandler: Handler,
     private val displayTracker: DisplayTracker,
     private val scrimLogger: ScrimLogger,
 ) {
@@ -96,11 +102,15 @@
     fun init() {
         // This method will be called only on devices where this animation is enabled,
         // so normally this thread won't be created
-        bgHandler = threadFactory.buildHandlerOnNewThread(TAG)
+        bgHandler = unfoldProgressHandler
         bgExecutor = threadFactory.buildDelayableExecutorOnHandler(bgHandler)
 
         deviceStateManager.registerCallback(bgExecutor, FoldListener())
-        unfoldTransitionProgressProvider.addCallback(transitionListener)
+        if (unfoldAnimationBackgroundProgress()) {
+            unfoldTransitionBgProgressProvider.get().addCallback(transitionListener)
+        } else {
+            unfoldTransitionProgressProvider.get().addCallback(transitionListener)
+        }
         rotationChangeProvider.addCallback(rotationWatcher)
 
         val containerBuilder =
@@ -169,8 +179,13 @@
 
         overlayAddReason = reason
 
-        val newRoot = SurfaceControlViewHost(context, context.display!!, wwm,
-                "UnfoldLightRevealOverlayAnimation")
+        val newRoot =
+            SurfaceControlViewHost(
+                context,
+                context.display,
+                wwm,
+                "UnfoldLightRevealOverlayAnimation"
+            )
         val params = getLayoutParams()
         val newView =
             LightRevealScrim(
@@ -353,12 +368,13 @@
     }
 
     private fun executeInBackground(f: () -> Unit) {
-        check(Looper.myLooper() != bgHandler.looper) {
-            "Trying to execute using background handler while already running" +
-                " in the background handler"
+        // This is needed to allow progresses to be received both from the main thread (that will
+        // schedule a runnable on the bg thread), and from the bg thread directly (no reposting).
+        if (bgHandler.looper.isCurrentThread) {
+            f()
+        } else {
+            bgHandler.post(f)
         }
-        // The UiBackground executor is not used as it doesn't have a prepared looper.
-        bgHandler.post(f)
     }
 
     private fun ensureInBackground() {
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTraceLogger.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTraceLogger.kt
index 12b8845..94912bf8 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTraceLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTraceLogger.kt
@@ -21,11 +21,14 @@
 import com.android.systemui.CoreStartable
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.unfold.system.DeviceStateRepository
 import com.android.systemui.unfold.updates.FoldStateRepository
 import javax.inject.Inject
+import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.plus
 
 /**
  * Logs several unfold related details in a trace. Mainly used for debugging and investigate
@@ -37,7 +40,8 @@
 constructor(
     private val context: Context,
     private val foldStateRepository: FoldStateRepository,
-    @Application private val applicationScope: CoroutineScope,
+    @Application applicationScope: CoroutineScope,
+    @Background private val coroutineContext: CoroutineContext,
     private val deviceStateRepository: DeviceStateRepository
 ) : CoreStartable {
     private val isFoldable: Boolean
@@ -46,20 +50,22 @@
                 .getIntArray(com.android.internal.R.array.config_foldedDeviceStates)
                 .isNotEmpty()
 
+    private val bgScope = applicationScope.plus(coroutineContext)
+
     override fun start() {
         if (!isFoldable) return
 
-        applicationScope.launch {
+        bgScope.launch {
             val foldUpdateLogger = TraceStateLogger("FoldUpdate")
             foldStateRepository.foldUpdate.collect { foldUpdateLogger.log(it.name) }
         }
 
-        applicationScope.launch {
+        bgScope.launch {
             foldStateRepository.hingeAngle.collect {
                 Trace.traceCounter(Trace.TRACE_TAG_APP, "hingeAngle", it.toInt())
             }
         }
-        applicationScope.launch {
+        bgScope.launch {
             val foldedStateLogger = TraceStateLogger("FoldedState")
             deviceStateRepository.isFolded.collect { isFolded ->
                 foldedStateLogger.log(
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt
index 7b628f8..0531487 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.keyguard.LifecycleScreenStatusProvider
 import com.android.systemui.unfold.config.UnfoldTransitionConfig
+import com.android.systemui.unfold.dagger.UnfoldMain
 import com.android.systemui.unfold.data.repository.UnfoldTransitionRepository
 import com.android.systemui.unfold.data.repository.UnfoldTransitionRepositoryImpl
 import com.android.systemui.unfold.domain.interactor.UnfoldTransitionInteractor
@@ -102,7 +103,7 @@
     @Singleton
     fun provideNaturalRotationProgressProvider(
         context: Context,
-        rotationChangeProvider: RotationChangeProvider,
+        @UnfoldMain rotationChangeProvider: RotationChangeProvider,
         unfoldTransitionProgressProvider: Optional<UnfoldTransitionProgressProvider>
     ): Optional<NaturalRotationUnfoldProgressProvider> =
         unfoldTransitionProgressProvider.map { provider ->
@@ -153,7 +154,8 @@
 
         return resultingProvider?.get()?.orElse(null)?.let { unfoldProgressProvider ->
             UnfoldProgressProvider(unfoldProgressProvider, foldProvider)
-        } ?: ShellUnfoldProgressProvider.NO_PROVIDER
+        }
+            ?: ShellUnfoldProgressProvider.NO_PROVIDER
     }
 
     @Provides
diff --git a/packages/SystemUI/src/com/android/systemui/util/kotlin/CoroutinesModule.kt b/packages/SystemUI/src/com/android/systemui/util/kotlin/CoroutinesModule.kt
index cc9335e..472f0ae 100644
--- a/packages/SystemUI/src/com/android/systemui/util/kotlin/CoroutinesModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/kotlin/CoroutinesModule.kt
@@ -14,6 +14,7 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.plus
 import kotlin.coroutines.CoroutineContext
 import kotlin.coroutines.EmptyCoroutineContext
 
@@ -29,6 +30,14 @@
 
     @Provides
     @SysUISingleton
+    @Background
+    fun bgApplicationScope(
+            @Application applicationScope: CoroutineScope,
+            @Background coroutineContext: CoroutineContext,
+    ): CoroutineScope = applicationScope.plus(coroutineContext)
+
+    @Provides
+    @SysUISingleton
     @Main
     @Deprecated(
         "Use @Main CoroutineContext instead",
diff --git a/packages/SystemUI/tests/src/com/android/systemui/biometrics/AuthContainerViewTest.kt b/packages/SystemUI/tests/src/com/android/systemui/biometrics/AuthContainerViewTest.kt
index f4122d5..ea20d29 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/biometrics/AuthContainerViewTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/biometrics/AuthContainerViewTest.kt
@@ -338,6 +338,13 @@
         waitForIdleSync()
 
         assertThat(container.hasCredentialView()).isTrue()
+        assertThat(container.hasBiometricPrompt()).isFalse()
+
+        // Check credential view persists after new attachment
+        container.onAttachedToWindow()
+
+        assertThat(container.hasCredentialView()).isTrue()
+        assertThat(container.hasBiometricPrompt()).isFalse()
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt
index 0004f52..910097e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorTest.kt
@@ -21,6 +21,7 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.coroutines.collectValues
 import com.android.systemui.deviceentry.data.repository.FakeDeviceEntryRepository
 import com.android.systemui.keyguard.data.repository.FakeDeviceEntryFaceAuthRepository
 import com.android.systemui.keyguard.data.repository.FakeTrustRepository
@@ -56,6 +57,13 @@
         )
 
     @Test
+    fun canSwipeToEnter_startsNull() =
+        testScope.runTest {
+            val values by collectValues(underTest.canSwipeToEnter)
+            assertThat(values[0]).isNull()
+        }
+
+    @Test
     fun isUnlocked_whenAuthMethodIsNoneAndLockscreenDisabled_isTrue() =
         testScope.runTest {
             utils.authenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.None)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/external/CustomTileStatePersisterTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/external/CustomTileStatePersisterTest.kt
index a9f8ea0..81d02b8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/external/CustomTileStatePersisterTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/external/CustomTileStatePersisterTest.kt
@@ -85,7 +85,7 @@
         `when`(sharedPreferences.edit()).thenReturn(editor)
 
         tile = Tile()
-        customTileStatePersister = CustomTileStatePersister(mockContext)
+        customTileStatePersister = CustomTileStatePersisterImpl(mockContext)
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogTest.kt
index 3808c7e..313ccb8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogTest.kt
@@ -228,16 +228,16 @@
                 showPairNewDevice = true
             )
 
-            val seeAllLayout = bluetoothTileDialog.requireViewById<View>(R.id.see_all_layout_group)
-            val pairNewLayout =
-                bluetoothTileDialog.requireViewById<View>(R.id.pair_new_device_layout_group)
+            val seeAllButton = bluetoothTileDialog.requireViewById<View>(R.id.see_all_button)
+            val pairNewButton =
+                bluetoothTileDialog.requireViewById<View>(R.id.pair_new_device_button)
             val recyclerView = bluetoothTileDialog.requireViewById<RecyclerView>(R.id.device_list)
             val adapter = recyclerView?.adapter as BluetoothTileDialog.Adapter
 
-            assertThat(seeAllLayout).isNotNull()
-            assertThat(seeAllLayout.visibility).isEqualTo(GONE)
-            assertThat(pairNewLayout).isNotNull()
-            assertThat(pairNewLayout.visibility).isEqualTo(VISIBLE)
+            assertThat(seeAllButton).isNotNull()
+            assertThat(seeAllButton.visibility).isEqualTo(GONE)
+            assertThat(pairNewButton).isNotNull()
+            assertThat(pairNewButton.visibility).isEqualTo(VISIBLE)
             assertThat(adapter.itemCount).isEqualTo(1)
         }
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogViewModelTest.kt
index fb5dd21..99993f2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/BluetoothTileDialogViewModelTest.kt
@@ -30,7 +30,6 @@
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.nullable
 import com.android.systemui.util.time.FakeSystemClock
-import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -113,9 +112,7 @@
         testScope.runTest {
             bluetoothTileDialogViewModel.showDialog(context, null)
 
-            assertThat(bluetoothTileDialogViewModel.dialog).isNotNull()
             verify(dialogLaunchAnimator, never()).showFromView(any(), any(), any(), any())
-            assertThat(bluetoothTileDialogViewModel.dialog?.isShowing).isTrue()
             verify(uiEventLogger).log(BluetoothTileDialogUiEvent.BLUETOOTH_TILE_DIALOG_SHOWN)
         }
     }
@@ -125,7 +122,6 @@
         testScope.runTest {
             bluetoothTileDialogViewModel.showDialog(mContext, LinearLayout(mContext))
 
-            assertThat(bluetoothTileDialogViewModel.dialog).isNotNull()
             verify(dialogLaunchAnimator).showFromView(any(), any(), nullable(), anyBoolean())
         }
     }
@@ -136,7 +132,6 @@
             backgroundExecutor.execute {
                 bluetoothTileDialogViewModel.showDialog(mContext, LinearLayout(mContext))
 
-                assertThat(bluetoothTileDialogViewModel.dialog).isNotNull()
                 verify(dialogLaunchAnimator).showFromView(any(), any(), nullable(), anyBoolean())
             }
         }
@@ -147,7 +142,6 @@
         testScope.runTest {
             bluetoothTileDialogViewModel.showDialog(context, null)
 
-            assertThat(bluetoothTileDialogViewModel.dialog).isNotNull()
             verify(deviceItemInteractor).deviceItemUpdate
         }
     }
@@ -157,7 +151,6 @@
         testScope.runTest {
             bluetoothTileDialogViewModel.showDialog(context, null)
 
-            assertThat(bluetoothTileDialogViewModel.dialog).isNotNull()
             verify(bluetoothStateInteractor).bluetoothStateUpdate
         }
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/DeviceItemInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/DeviceItemInteractorTest.kt
index 4c173cc..e236f4a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/DeviceItemInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/dialog/bluetooth/DeviceItemInteractorTest.kt
@@ -220,45 +220,57 @@
 
     @Test
     fun testUpdateDeviceItemOnClick_connectedMedia_setActive() {
-        `when`(deviceItem1.type).thenReturn(DeviceItemType.AVAILABLE_MEDIA_BLUETOOTH_DEVICE)
+        testScope.runTest {
+            `when`(deviceItem1.type).thenReturn(DeviceItemType.AVAILABLE_MEDIA_BLUETOOTH_DEVICE)
 
-        interactor.updateDeviceItemOnClick(deviceItem1)
+            interactor.updateDeviceItemOnClick(deviceItem1)
 
-        verify(cachedDevice1).setActive()
-        verify(logger)
-            .logDeviceClick(cachedDevice1.address, DeviceItemType.AVAILABLE_MEDIA_BLUETOOTH_DEVICE)
+            verify(cachedDevice1).setActive()
+            verify(logger)
+                .logDeviceClick(
+                    cachedDevice1.address,
+                    DeviceItemType.AVAILABLE_MEDIA_BLUETOOTH_DEVICE
+                )
+        }
     }
 
     @Test
     fun testUpdateDeviceItemOnClick_activeMedia_disconnect() {
-        `when`(deviceItem1.type).thenReturn(DeviceItemType.ACTIVE_MEDIA_BLUETOOTH_DEVICE)
+        testScope.runTest {
+            `when`(deviceItem1.type).thenReturn(DeviceItemType.ACTIVE_MEDIA_BLUETOOTH_DEVICE)
 
-        interactor.updateDeviceItemOnClick(deviceItem1)
+            interactor.updateDeviceItemOnClick(deviceItem1)
 
-        verify(cachedDevice1).disconnect()
-        verify(logger)
-            .logDeviceClick(cachedDevice1.address, DeviceItemType.ACTIVE_MEDIA_BLUETOOTH_DEVICE)
+            verify(cachedDevice1).disconnect()
+            verify(logger)
+                .logDeviceClick(cachedDevice1.address, DeviceItemType.ACTIVE_MEDIA_BLUETOOTH_DEVICE)
+        }
     }
 
     @Test
     fun testUpdateDeviceItemOnClick_connectedOtherDevice_disconnect() {
-        `when`(deviceItem1.type).thenReturn(DeviceItemType.CONNECTED_BLUETOOTH_DEVICE)
+        testScope.runTest {
+            `when`(deviceItem1.type).thenReturn(DeviceItemType.CONNECTED_BLUETOOTH_DEVICE)
 
-        interactor.updateDeviceItemOnClick(deviceItem1)
+            interactor.updateDeviceItemOnClick(deviceItem1)
 
-        verify(cachedDevice1).disconnect()
-        verify(logger)
-            .logDeviceClick(cachedDevice1.address, DeviceItemType.CONNECTED_BLUETOOTH_DEVICE)
+            verify(cachedDevice1).disconnect()
+            verify(logger)
+                .logDeviceClick(cachedDevice1.address, DeviceItemType.CONNECTED_BLUETOOTH_DEVICE)
+        }
     }
 
     @Test
     fun testUpdateDeviceItemOnClick_saved_connect() {
-        `when`(deviceItem1.type).thenReturn(DeviceItemType.SAVED_BLUETOOTH_DEVICE)
+        testScope.runTest {
+            `when`(deviceItem1.type).thenReturn(DeviceItemType.SAVED_BLUETOOTH_DEVICE)
 
-        interactor.updateDeviceItemOnClick(deviceItem1)
+            interactor.updateDeviceItemOnClick(deviceItem1)
 
-        verify(cachedDevice1).connect()
-        verify(logger).logDeviceClick(cachedDevice1.address, DeviceItemType.SAVED_BLUETOOTH_DEVICE)
+            verify(cachedDevice1).connect()
+            verify(logger)
+                .logDeviceClick(cachedDevice1.address, DeviceItemType.SAVED_BLUETOOTH_DEVICE)
+        }
     }
 
     private fun createFactory(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/impl/custom/data/repository/CustomTileRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/impl/custom/data/repository/CustomTileRepositoryTest.kt
new file mode 100644
index 0000000..cf076c5
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/impl/custom/data/repository/CustomTileRepositoryTest.kt
@@ -0,0 +1,259 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.tiles.impl.custom.data.repository
+
+import android.content.ComponentName
+import android.graphics.drawable.Icon
+import android.os.UserHandle
+import android.service.quicksettings.Tile
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectValues
+import com.android.systemui.qs.external.FakeCustomTileStatePersister
+import com.android.systemui.qs.external.TileServiceKey
+import com.android.systemui.qs.pipeline.shared.TileSpec
+import com.android.systemui.qs.tiles.impl.custom.TileSubject.Companion.assertThat
+import com.android.systemui.qs.tiles.impl.custom.commons.copy
+import com.android.systemui.qs.tiles.impl.custom.data.entity.CustomTileDefaults
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@OptIn(ExperimentalCoroutinesApi::class)
+class CustomTileRepositoryTest : SysuiTestCase() {
+
+    private val testScope = TestScope()
+
+    private val persister = FakeCustomTileStatePersister()
+
+    private val underTest: CustomTileRepository =
+        CustomTileRepositoryImpl(
+            TileSpec.create(TEST_COMPONENT),
+            persister,
+            testScope.testScheduler,
+        )
+
+    @Test
+    fun persistableTileIsRestoredForUser() =
+        testScope.runTest {
+            persister.persistState(TEST_TILE_KEY_1, TEST_TILE_1)
+            persister.persistState(TEST_TILE_KEY_2, TEST_TILE_2)
+
+            underTest.restoreForTheUserIfNeeded(TEST_USER_1, true)
+            runCurrent()
+
+            assertThat(underTest.getTile(TEST_USER_1)).isEqualTo(TEST_TILE_1)
+            assertThat(underTest.getTiles(TEST_USER_1).first()).isEqualTo(TEST_TILE_1)
+        }
+
+    @Test
+    fun notPersistableTileIsNotRestored() =
+        testScope.runTest {
+            persister.persistState(TEST_TILE_KEY_1, TEST_TILE_1)
+            val tiles = collectValues(underTest.getTiles(TEST_USER_1))
+
+            underTest.restoreForTheUserIfNeeded(TEST_USER_1, false)
+            runCurrent()
+
+            assertThat(tiles()).isEmpty()
+        }
+
+    @Test
+    fun emptyPersistedStateIsHandled() =
+        testScope.runTest {
+            val tiles = collectValues(underTest.getTiles(TEST_USER_1))
+
+            underTest.restoreForTheUserIfNeeded(TEST_USER_1, true)
+            runCurrent()
+
+            assertThat(tiles()).isEmpty()
+        }
+
+    @Test
+    fun updatingWithPersistableTilePersists() =
+        testScope.runTest {
+            underTest.updateWithTile(TEST_USER_1, TEST_TILE_1, true)
+            runCurrent()
+
+            assertThat(persister.readState(TEST_TILE_KEY_1)).isEqualTo(TEST_TILE_1)
+        }
+
+    @Test
+    fun updatingWithNotPersistableTileDoesntPersist() =
+        testScope.runTest {
+            underTest.updateWithTile(TEST_USER_1, TEST_TILE_1, false)
+            runCurrent()
+
+            assertThat(persister.readState(TEST_TILE_KEY_1)).isNull()
+        }
+
+    @Test
+    fun updateWithTileEmits() =
+        testScope.runTest {
+            underTest.updateWithTile(TEST_USER_1, TEST_TILE_1, true)
+            runCurrent()
+
+            assertThat(underTest.getTiles(TEST_USER_1).first()).isEqualTo(TEST_TILE_1)
+            assertThat(underTest.getTile(TEST_USER_1)).isEqualTo(TEST_TILE_1)
+        }
+
+    @Test
+    fun updatingPeristableWithDefaultsPersists() =
+        testScope.runTest {
+            underTest.updateWithDefaults(TEST_USER_1, TEST_DEFAULTS_1, true)
+            runCurrent()
+
+            assertThat(persister.readState(TEST_TILE_KEY_1)).isEqualTo(TEST_TILE_1)
+        }
+
+    @Test
+    fun updatingNotPersistableWithDefaultsDoesntPersist() =
+        testScope.runTest {
+            underTest.updateWithDefaults(TEST_USER_1, TEST_DEFAULTS_1, false)
+            runCurrent()
+
+            assertThat(persister.readState(TEST_TILE_KEY_1)).isNull()
+        }
+
+    @Test
+    fun updatingPeristableWithErrorDefaultsDoesntPersist() =
+        testScope.runTest {
+            underTest.updateWithDefaults(TEST_USER_1, CustomTileDefaults.Error, true)
+            runCurrent()
+
+            assertThat(persister.readState(TEST_TILE_KEY_1)).isNull()
+        }
+
+    @Test
+    fun updateWithDefaultsEmits() =
+        testScope.runTest {
+            underTest.updateWithDefaults(TEST_USER_1, TEST_DEFAULTS_1, true)
+            runCurrent()
+
+            assertThat(underTest.getTiles(TEST_USER_1).first()).isEqualTo(TEST_TILE_1)
+            assertThat(underTest.getTile(TEST_USER_1)).isEqualTo(TEST_TILE_1)
+        }
+
+    @Test
+    fun getTileForAnotherUserReturnsNull() =
+        testScope.runTest {
+            underTest.updateWithTile(TEST_USER_1, TEST_TILE_1, true)
+            runCurrent()
+
+            assertThat(underTest.getTile(TEST_USER_2)).isNull()
+        }
+
+    @Test
+    fun getTilesForAnotherUserEmpty() =
+        testScope.runTest {
+            val tiles = collectValues(underTest.getTiles(TEST_USER_2))
+
+            underTest.updateWithTile(TEST_USER_1, TEST_TILE_1, true)
+            runCurrent()
+
+            assertThat(tiles()).isEmpty()
+        }
+
+    @Test
+    fun updatingWithTileForTheSameUserAddsData() =
+        testScope.runTest {
+            underTest.updateWithTile(TEST_USER_1, TEST_TILE_1, true)
+            runCurrent()
+
+            underTest.updateWithTile(TEST_USER_1, Tile().apply { subtitle = "test_subtitle" }, true)
+            runCurrent()
+
+            val expectedTile = TEST_TILE_1.copy().apply { subtitle = "test_subtitle" }
+            assertThat(underTest.getTile(TEST_USER_1)).isEqualTo(expectedTile)
+            assertThat(underTest.getTiles(TEST_USER_1).first()).isEqualTo(expectedTile)
+        }
+
+    @Test
+    fun updatingWithTileForAnotherUserOverridesTile() =
+        testScope.runTest {
+            underTest.updateWithTile(TEST_USER_1, TEST_TILE_1, true)
+            runCurrent()
+
+            val tiles = collectValues(underTest.getTiles(TEST_USER_2))
+            underTest.updateWithTile(TEST_USER_2, TEST_TILE_2, true)
+            runCurrent()
+
+            assertThat(underTest.getTile(TEST_USER_2)).isEqualTo(TEST_TILE_2)
+            assertThat(tiles()).hasSize(1)
+            assertThat(tiles().last()).isEqualTo(TEST_TILE_2)
+        }
+
+    @Test
+    fun updatingWithDefaultsForTheSameUserAddsData() =
+        testScope.runTest {
+            underTest.updateWithTile(TEST_USER_1, Tile().apply { subtitle = "test_subtitle" }, true)
+            runCurrent()
+
+            underTest.updateWithDefaults(TEST_USER_1, TEST_DEFAULTS_1, true)
+            runCurrent()
+
+            val expectedTile = TEST_TILE_1.copy().apply { subtitle = "test_subtitle" }
+            assertThat(underTest.getTile(TEST_USER_1)).isEqualTo(expectedTile)
+            assertThat(underTest.getTiles(TEST_USER_1).first()).isEqualTo(expectedTile)
+        }
+
+    @Test
+    fun updatingWithDefaultsForAnotherUserOverridesTile() =
+        testScope.runTest {
+            underTest.updateWithDefaults(TEST_USER_1, TEST_DEFAULTS_1, true)
+            runCurrent()
+
+            val tiles = collectValues(underTest.getTiles(TEST_USER_2))
+            underTest.updateWithDefaults(TEST_USER_2, TEST_DEFAULTS_2, true)
+            runCurrent()
+
+            assertThat(underTest.getTile(TEST_USER_2)).isEqualTo(TEST_TILE_2)
+            assertThat(tiles()).hasSize(1)
+            assertThat(tiles().last()).isEqualTo(TEST_TILE_2)
+        }
+
+    private companion object {
+
+        val TEST_COMPONENT = ComponentName("test.pkg", "test.cls")
+
+        val TEST_USER_1 = UserHandle.of(1)!!
+        val TEST_TILE_1 =
+            Tile().apply {
+                label = "test_tile_1"
+                icon = Icon.createWithContentUri("file://test_1")
+            }
+        val TEST_TILE_KEY_1 = TileServiceKey(TEST_COMPONENT, TEST_USER_1.identifier)
+        val TEST_DEFAULTS_1 = CustomTileDefaults.Result(TEST_TILE_1.icon, TEST_TILE_1.label)
+
+        val TEST_USER_2 = UserHandle.of(2)!!
+        val TEST_TILE_2 =
+            Tile().apply {
+                label = "test_tile_2"
+                icon = Icon.createWithContentUri("file://test_2")
+            }
+        val TEST_TILE_KEY_2 = TileServiceKey(TEST_COMPONENT, TEST_USER_2.identifier)
+        val TEST_DEFAULTS_2 = CustomTileDefaults.Result(TEST_TILE_2.icon, TEST_TILE_2.label)
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/impl/custom/domain/interactor/CustomTileInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/impl/custom/domain/interactor/CustomTileInteractorTest.kt
new file mode 100644
index 0000000..eebb145
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/impl/custom/domain/interactor/CustomTileInteractorTest.kt
@@ -0,0 +1,187 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.tiles.impl.custom.domain.interactor
+
+import android.content.ComponentName
+import android.graphics.drawable.Icon
+import android.os.UserHandle
+import android.service.quicksettings.Tile
+import android.text.format.DateUtils
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectValues
+import com.android.systemui.qs.external.FakeCustomTileStatePersister
+import com.android.systemui.qs.external.TileServiceKey
+import com.android.systemui.qs.external.TileServiceManager
+import com.android.systemui.qs.pipeline.shared.TileSpec
+import com.android.systemui.qs.tiles.impl.custom.TileSubject.Companion.assertThat
+import com.android.systemui.qs.tiles.impl.custom.data.entity.CustomTileDefaults
+import com.android.systemui.qs.tiles.impl.custom.data.repository.FakeCustomTileDefaultsRepository
+import com.android.systemui.qs.tiles.impl.custom.data.repository.FakeCustomTileRepository
+import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.advanceTimeBy
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@OptIn(ExperimentalCoroutinesApi::class)
+class CustomTileInteractorTest : SysuiTestCase() {
+
+    @Mock private lateinit var tileServiceManager: TileServiceManager
+
+    private val testScope = TestScope()
+
+    private val defaultsRepository = FakeCustomTileDefaultsRepository()
+    private val customTileStatePersister = FakeCustomTileStatePersister()
+    private val customTileRepository =
+        FakeCustomTileRepository(
+            TEST_TILE_SPEC,
+            customTileStatePersister,
+            testScope.testScheduler,
+        )
+
+    private lateinit var underTest: CustomTileInteractor
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+
+        underTest =
+            CustomTileInteractor(
+                TEST_USER,
+                defaultsRepository,
+                customTileRepository,
+                tileServiceManager,
+                testScope.backgroundScope,
+                testScope.testScheduler,
+            )
+    }
+
+    @Test
+    fun activeTileIsAvailableAfterRestored() =
+        testScope.runTest {
+            whenever(tileServiceManager.isActiveTile).thenReturn(true)
+            customTileStatePersister.persistState(
+                TileServiceKey(TEST_COMPONENT, TEST_USER.identifier),
+                TEST_TILE,
+            )
+
+            underTest.init()
+
+            assertThat(underTest.tile).isEqualTo(TEST_TILE)
+            assertThat(underTest.tiles.first()).isEqualTo(TEST_TILE)
+        }
+
+    @Test
+    fun notActiveTileIsAvailableAfterUpdated() =
+        testScope.runTest {
+            whenever(tileServiceManager.isActiveTile).thenReturn(false)
+            customTileStatePersister.persistState(
+                TileServiceKey(TEST_COMPONENT, TEST_USER.identifier),
+                TEST_TILE,
+            )
+            val tiles = collectValues(underTest.tiles)
+            val initJob = launch { underTest.init() }
+
+            underTest.updateTile(TEST_TILE)
+            runCurrent()
+            initJob.join()
+
+            assertThat(tiles()).hasSize(1)
+            assertThat(tiles().last()).isEqualTo(TEST_TILE)
+        }
+
+    @Test
+    fun notActiveTileIsAvailableAfterDefaultsUpdated() =
+        testScope.runTest {
+            whenever(tileServiceManager.isActiveTile).thenReturn(false)
+            customTileStatePersister.persistState(
+                TileServiceKey(TEST_COMPONENT, TEST_USER.identifier),
+                TEST_TILE,
+            )
+            val tiles = collectValues(underTest.tiles)
+            val initJob = launch { underTest.init() }
+
+            defaultsRepository.putDefaults(TEST_USER, TEST_COMPONENT, TEST_DEFAULTS)
+            defaultsRepository.requestNewDefaults(TEST_USER, TEST_COMPONENT)
+            runCurrent()
+            initJob.join()
+
+            assertThat(tiles()).hasSize(1)
+            assertThat(tiles().last()).isEqualTo(TEST_TILE)
+        }
+
+    @Test(expected = IllegalStateException::class)
+    fun getTileBeforeInitThrows() = testScope.runTest { underTest.tile }
+
+    @Test
+    fun initSuspendsForActiveTileNotRestoredAndNotUpdated() =
+        testScope.runTest {
+            whenever(tileServiceManager.isActiveTile).thenReturn(true)
+            val tiles = collectValues(underTest.tiles)
+
+            val initJob = backgroundScope.launch { underTest.init() }
+            advanceTimeBy(1 * DateUtils.DAY_IN_MILLIS)
+
+            // Is still suspended
+            assertThat(initJob.isActive).isTrue()
+            assertThat(tiles()).isEmpty()
+        }
+
+    @Test
+    fun initSuspendedForNotActiveTileWithoutUpdates() =
+        testScope.runTest {
+            whenever(tileServiceManager.isActiveTile).thenReturn(false)
+            customTileStatePersister.persistState(
+                TileServiceKey(TEST_COMPONENT, TEST_USER.identifier),
+                TEST_TILE,
+            )
+            val tiles = collectValues(underTest.tiles)
+
+            val initJob = backgroundScope.launch { underTest.init() }
+            advanceTimeBy(1 * DateUtils.DAY_IN_MILLIS)
+
+            // Is still suspended
+            assertThat(initJob.isActive).isTrue()
+            assertThat(tiles()).isEmpty()
+        }
+
+    private companion object {
+
+        val TEST_COMPONENT = ComponentName("test.pkg", "test.cls")
+        val TEST_TILE_SPEC = TileSpec.create(TEST_COMPONENT)
+        val TEST_USER = UserHandle.of(1)!!
+        val TEST_TILE =
+            Tile().apply {
+                label = "test_tile_1"
+                icon = Icon.createWithContentUri("file://test_1")
+            }
+        val TEST_DEFAULTS = CustomTileDefaults.Result(TEST_TILE.icon, TEST_TILE.label)
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt b/packages/SystemUI/tests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
index c3294ff..c953743 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
@@ -61,7 +61,6 @@
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
-import junit.framework.Assert.assertTrue
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -273,6 +272,9 @@
     }
 
     @Test
+    fun startsInLockscreenScene() = testScope.runTest { assertCurrentScene(SceneKey.Lockscreen) }
+
+    @Test
     fun clickLockButtonAndEnterCorrectPin_unlocksDevice() =
         testScope.runTest {
             emulateUserDrivenTransition(SceneKey.Bouncer)
@@ -336,7 +338,7 @@
         testScope.runTest {
             val upDestinationSceneKey by collectLastValue(shadeSceneViewModel.upDestinationSceneKey)
             setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
-            assertTrue(deviceEntryInteractor.canSwipeToEnter.value)
+            assertThat(deviceEntryInteractor.canSwipeToEnter.value).isTrue()
             assertCurrentScene(SceneKey.Lockscreen)
 
             // Emulate a user swipe to dismiss the lockscreen.
diff --git a/packages/SystemUI/tests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt b/packages/SystemUI/tests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
index c4ec56c..adc1c61 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/scene/domain/startable/SceneContainerStartableTest.kt
@@ -145,6 +145,18 @@
         }
 
     @Test
+    fun startsInLockscreenScene() =
+        testScope.runTest {
+            val currentSceneKey by collectLastValue(sceneInteractor.desiredScene.map { it.key })
+            prepareState()
+
+            underTest.start()
+            runCurrent()
+
+            assertThat(currentSceneKey).isEqualTo(SceneKey.Lockscreen)
+        }
+
+    @Test
     fun switchToLockscreenWhenDeviceLocks() =
         testScope.runTest {
             val currentSceneKey by collectLastValue(sceneInteractor.desiredScene.map { it.key })
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerBaseTest.java
index 62c0ebe..1dbb297 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerBaseTest.java
@@ -22,8 +22,6 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
-import static kotlinx.coroutines.test.TestCoroutineDispatchersKt.StandardTestDispatcher;
-
 import android.content.res.Resources;
 import android.os.Handler;
 import android.os.Looper;
@@ -295,10 +293,8 @@
                 )
         );
 
-        mActiveNotificationsInteractor = new ActiveNotificationsInteractor(
-                        new ActiveNotificationListRepository(),
-                        StandardTestDispatcher(/* scheduler = */ null, /* name = */ null)
-                );
+        mActiveNotificationsInteractor =
+                new ActiveNotificationsInteractor(new ActiveNotificationListRepository());
 
         KeyguardStatusView keyguardStatusView = new KeyguardStatusView(mContext);
         keyguardStatusView.setId(R.id.keyguard_status_view);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt
index 6374d5e..b86f841 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt
@@ -25,19 +25,14 @@
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 
 @SmallTest
 class RenderNotificationsListInteractorTest : SysuiTestCase() {
-    private val backgroundDispatcher = StandardTestDispatcher()
-    private val testScope = TestScope(backgroundDispatcher)
 
     private val notifsRepository = ActiveNotificationListRepository()
-    private val notifsInteractor =
-        ActiveNotificationsInteractor(notifsRepository, backgroundDispatcher)
+    private val notifsInteractor = ActiveNotificationsInteractor(notifsRepository)
     private val underTest =
         RenderNotificationListInteractor(
             notifsRepository,
@@ -45,26 +40,21 @@
         )
 
     @Test
-    fun setRenderedList_preservesOrdering() =
-        testScope.runTest {
-            val notifs by collectLastValue(notifsInteractor.topLevelRepresentativeNotifications)
-            val keys = (1..50).shuffled().map { "$it" }
-            val entries =
-                keys.map {
-                    mock<ListEntry> {
-                        val mockRep =
-                            mock<NotificationEntry> {
-                                whenever(key).thenReturn(it)
-                                whenever(sbn).thenReturn(mock())
-                                whenever(icons).thenReturn(mock())
-                            }
-                        whenever(representativeEntry).thenReturn(mockRep)
+    fun setRenderedList_preservesOrdering() = runTest {
+        val notifs by collectLastValue(notifsInteractor.topLevelRepresentativeNotifications)
+        val keys = (1..50).shuffled().map { "$it" }
+        val entries =
+            keys.map {
+                mock<ListEntry> {
+                    val mockRep = mock<NotificationEntry> {
+                        whenever(key).thenReturn(it)
+                        whenever(sbn).thenReturn(mock())
+                        whenever(icons).thenReturn(mock())
                     }
+                    whenever(representativeEntry).thenReturn(mockRep)
                 }
-            underTest.setRenderedList(entries)
-            assertThat(notifs)
-                .comparingElementsUsing(byKey)
-                .containsExactlyElementsIn(keys)
-                .inOrder()
-        }
+            }
+        underTest.setRenderedList(entries)
+        assertThat(notifs).comparingElementsUsing(byKey).containsExactlyElementsIn(keys).inOrder()
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
index 7558974..ff5c026 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
@@ -35,7 +35,6 @@
 import static org.mockito.Mockito.when;
 
 import static kotlinx.coroutines.flow.FlowKt.emptyFlow;
-import static kotlinx.coroutines.test.TestCoroutineDispatchersKt.StandardTestDispatcher;
 
 import android.metrics.LogMaker;
 import android.testing.AndroidTestingRunner;
@@ -170,8 +169,7 @@
             new ActiveNotificationListRepository();
 
     private final ActiveNotificationsInteractor mActiveNotificationsInteractor =
-            new ActiveNotificationsInteractor(mActiveNotificationsRepository,
-                    StandardTestDispatcher(/* scheduler = */ null, /* name = */ null));
+            new ActiveNotificationsInteractor(mActiveNotificationsRepository);
 
     private final SeenNotificationsInteractor mSeenNotificationsInteractor =
             new SeenNotificationsInteractor(mActiveNotificationsRepository);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProviderTest.kt
index 9fe2f56..14fb054 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProviderTest.kt
@@ -15,9 +15,10 @@
  */
 package com.android.systemui.unfold.progress
 
+import android.os.Handler
+import android.os.HandlerThread
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
-import androidx.test.platform.app.InstrumentationRegistry
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider
 import com.android.systemui.unfold.updates.FOLD_UPDATE_FINISH_CLOSED
@@ -26,6 +27,8 @@
 import com.android.systemui.unfold.updates.FOLD_UPDATE_START_CLOSING
 import com.android.systemui.unfold.updates.FOLD_UPDATE_START_OPENING
 import com.android.systemui.unfold.util.TestFoldStateProvider
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -37,16 +40,28 @@
     private val foldStateProvider: TestFoldStateProvider = TestFoldStateProvider()
     private val listener = TestUnfoldProgressListener()
     private lateinit var progressProvider: UnfoldTransitionProgressProvider
+    private val schedulerFactory =
+        mock<UnfoldFrameCallbackScheduler.Factory>().apply {
+            whenever(create()).then { UnfoldFrameCallbackScheduler() }
+        }
+    private val mockBgHandler = mock<Handler>()
+    private val fakeHandler = Handler(HandlerThread("UnfoldBg").apply { start() }.looper)
 
     @Before
     fun setUp() {
-        progressProvider = PhysicsBasedUnfoldTransitionProgressProvider(context, foldStateProvider)
+        progressProvider =
+            PhysicsBasedUnfoldTransitionProgressProvider(
+                context,
+                schedulerFactory,
+                foldStateProvider = foldStateProvider,
+                progressHandler = fakeHandler
+            )
         progressProvider.addCallback(listener)
     }
 
     @Test
     fun testUnfold_emitsIncreasingTransitionEvents() {
-        runOnMainThreadWithInterval(
+        runOnProgressThreadWithInterval(
             { foldStateProvider.sendFoldUpdate(FOLD_UPDATE_START_OPENING) },
             { foldStateProvider.sendHingeAngleUpdate(10f) },
             { foldStateProvider.sendUnfoldedScreenAvailable() },
@@ -63,7 +78,7 @@
 
     @Test
     fun testUnfold_emitsFinishingEvent() {
-        runOnMainThreadWithInterval(
+        runOnProgressThreadWithInterval(
             { foldStateProvider.sendFoldUpdate(FOLD_UPDATE_START_OPENING) },
             { foldStateProvider.sendHingeAngleUpdate(10f) },
             { foldStateProvider.sendUnfoldedScreenAvailable() },
@@ -77,7 +92,7 @@
 
     @Test
     fun testUnfold_screenAvailableOnlyAfterFullUnfold_emitsIncreasingTransitionEvents() {
-        runOnMainThreadWithInterval(
+        runOnProgressThreadWithInterval(
             { foldStateProvider.sendFoldUpdate(FOLD_UPDATE_START_OPENING) },
             { foldStateProvider.sendHingeAngleUpdate(10f) },
             { foldStateProvider.sendHingeAngleUpdate(90f) },
@@ -94,7 +109,7 @@
 
     @Test
     fun testFold_emitsDecreasingTransitionEvents() {
-        runOnMainThreadWithInterval(
+        runOnProgressThreadWithInterval(
             { foldStateProvider.sendFoldUpdate(FOLD_UPDATE_START_CLOSING) },
             { foldStateProvider.sendHingeAngleUpdate(170f) },
             { foldStateProvider.sendHingeAngleUpdate(90f) },
@@ -110,7 +125,7 @@
 
     @Test
     fun testUnfoldAndStopUnfolding_finishesTheUnfoldTransition() {
-        runOnMainThreadWithInterval(
+        runOnProgressThreadWithInterval(
             { foldStateProvider.sendFoldUpdate(FOLD_UPDATE_START_OPENING) },
             { foldStateProvider.sendUnfoldedScreenAvailable() },
             { foldStateProvider.sendHingeAngleUpdate(10f) },
@@ -126,7 +141,7 @@
 
     @Test
     fun testFoldImmediatelyAfterUnfold_runsFoldAnimation() {
-        runOnMainThreadWithInterval(
+        runOnProgressThreadWithInterval(
             { foldStateProvider.sendFoldUpdate(FOLD_UPDATE_START_OPENING) },
             { foldStateProvider.sendUnfoldedScreenAvailable() },
             { foldStateProvider.sendHingeAngleUpdate(10f) },
@@ -144,9 +159,12 @@
         with(listener.ensureTransitionFinished()) { assertHasFoldAnimationAtTheEnd() }
     }
 
-    private fun runOnMainThreadWithInterval(vararg blocks: () -> Unit, intervalMillis: Long = 60) {
+    private fun runOnProgressThreadWithInterval(
+        vararg blocks: () -> Unit,
+        intervalMillis: Long = 60,
+    ) {
         blocks.forEach {
-            InstrumentationRegistry.getInstrumentation().runOnMainSync { it() }
+            fakeHandler.post(it)
             Thread.sleep(intervalMillis)
         }
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceFoldStateProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceFoldStateProviderTest.kt
index aa49287..552b60c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceFoldStateProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceFoldStateProviderTest.kt
@@ -37,7 +37,6 @@
 import com.android.systemui.unfold.util.UnfoldKeyguardVisibilityProvider
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.capture
-import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
 import junit.framework.Assert.fail
 import java.util.concurrent.Executor
@@ -105,16 +104,15 @@
 
         foldStateProvider =
             DeviceFoldStateProvider(
-                config,
-                testHingeAngleProvider,
-                screenOnStatusProvider,
-                foldProvider,
-                activityTypeProvider,
-                unfoldKeyguardVisibilityProvider,
-                rotationChangeProvider,
-                context,
-                context.mainExecutor,
-                handler
+                    config,
+                    context,
+                    screenOnStatusProvider,
+                    activityTypeProvider,
+                    unfoldKeyguardVisibilityProvider,
+                    foldProvider,
+                    testHingeAngleProvider,
+                    rotationChangeProvider,
+                    handler
             )
 
         foldStateProvider.addCallback(
@@ -151,6 +149,12 @@
             null
         }
 
+        whenever(handler.post(any<Runnable>())).then { invocationOnMock ->
+            val runnable = invocationOnMock.getArgument<Runnable>(0)
+            runnable.run()
+            null
+        }
+
         // By default, we're on launcher.
         setupForegroundActivityType(isHomeActivity = true)
         setIsLargeScreen(true)
@@ -171,7 +175,7 @@
     }
 
     @Test
-    fun testOnUnfold_hingeAngleDecreasesBeforeInnerScreenAvailable_emitsOnlyStartAndInnerScreenAvailableEvents() {
+    fun onUnfold_angleDecrBeforeInnerScrAvailable_emitsOnlyStartAndInnerScrAvailableEvents() {
         setFoldState(folded = true)
         foldUpdates.clear()
 
@@ -187,7 +191,7 @@
     }
 
     @Test
-    fun testOnUnfold_hingeAngleDecreasesAfterInnerScreenAvailable_emitsStartInnerScreenAvailableAndStartClosingEvents() {
+    fun onUnfold_angleDecrAfterInnerScrAvailable_emitsStartInnerScrAvailableAndStartClosingEvnts() {
         setFoldState(folded = true)
         foldUpdates.clear()
 
@@ -690,7 +694,7 @@
             callbacks.forEach { it.onFoldUpdated(isFolded) }
         }
 
-        fun getNumberOfCallbacks(): Int{
+        fun getNumberOfCallbacks(): Int {
             return callbacks.size
         }
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/external/FakeCustomTileStatePersister.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/external/FakeCustomTileStatePersister.kt
new file mode 100644
index 0000000..29702eb
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/external/FakeCustomTileStatePersister.kt
@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.external
+
+import android.service.quicksettings.Tile
+
+class FakeCustomTileStatePersister : CustomTileStatePersister {
+
+    private val tiles: MutableMap<TileServiceKey, Tile> = mutableMapOf()
+
+    override fun readState(key: TileServiceKey): Tile? = tiles[key]
+
+    override fun persistState(key: TileServiceKey, tile: Tile) {
+        tiles[key] = tile
+    }
+
+    override fun removeState(key: TileServiceKey) {
+        tiles.remove(key)
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/custom/TileSubject.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/custom/TileSubject.kt
new file mode 100644
index 0000000..d2351dc
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/custom/TileSubject.kt
@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.tiles.impl.custom
+
+import android.service.quicksettings.Tile
+import com.android.systemui.qs.tiles.impl.custom.TileSubject.Companion.assertThat
+import com.android.systemui.qs.tiles.impl.custom.TileSubject.Companion.tiles
+import com.google.common.truth.FailureMetadata
+import com.google.common.truth.Subject
+import com.google.common.truth.Subject.Factory
+import com.google.common.truth.Truth
+
+/**
+ * [Tile]-specific extension for [Truth]. Use [assertThat] or [tiles] to get an instance of this
+ * subject.
+ */
+class TileSubject private constructor(failureMetadata: FailureMetadata, subject: Tile?) :
+    Subject(failureMetadata, subject) {
+
+    private val actual: Tile? = subject
+
+    /** Asserts if the [Tile] fields are the same. */
+    fun isEqualTo(other: Tile?) {
+        if (actual == null) {
+            check("other").that(other).isNull()
+            return
+        } else {
+            check("other").that(other).isNotNull()
+            other ?: return
+        }
+
+        check("icon").that(actual.icon).isEqualTo(other.icon)
+        check("label").that(actual.label).isEqualTo(other.label)
+        check("subtitle").that(actual.subtitle).isEqualTo(other.subtitle)
+        check("contentDescription")
+            .that(actual.contentDescription)
+            .isEqualTo(other.contentDescription)
+        check("stateDescription").that(actual.stateDescription).isEqualTo(other.stateDescription)
+        check("activityLaunchForClick")
+            .that(actual.activityLaunchForClick)
+            .isEqualTo(other.activityLaunchForClick)
+        check("state").that(actual.state).isEqualTo(other.state)
+    }
+
+    companion object {
+
+        /** Returns a factory to be used with [Truth.assertAbout]. */
+        fun tiles(): Factory<TileSubject, Tile?> {
+            return Factory { failureMetadata: FailureMetadata, subject: Tile? ->
+                TileSubject(failureMetadata, subject)
+            }
+        }
+
+        /** Shortcut for `Truth.assertAbout(tiles()).that(tile)`. */
+        fun assertThat(tile: Tile?): TileSubject = Truth.assertAbout(tiles()).that(tile)
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/custom/data/repository/FakeCustomTileDefaultsRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/custom/data/repository/FakeCustomTileDefaultsRepository.kt
index 13910fd..ccba072 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/custom/data/repository/FakeCustomTileDefaultsRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/custom/data/repository/FakeCustomTileDefaultsRepository.kt
@@ -19,15 +19,20 @@
 import android.content.ComponentName
 import android.os.UserHandle
 import com.android.systemui.qs.tiles.impl.custom.data.entity.CustomTileDefaults
+import kotlinx.coroutines.channels.BufferOverflow
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.mapNotNull
 
 class FakeCustomTileDefaultsRepository : CustomTileDefaultsRepository {
 
     private val defaults: MutableMap<DefaultsKey, CustomTileDefaults> = mutableMapOf()
-    private val defaultsFlow = MutableSharedFlow<DefaultsRequest>()
+    private val defaultsFlow =
+        MutableSharedFlow<DefaultsRequest>(
+            replay = 1,
+            onBufferOverflow = BufferOverflow.DROP_OLDEST
+        )
 
     private val mutableDefaultsRequests: MutableList<DefaultsRequest> = mutableListOf()
     val defaultsRequests: List<DefaultsRequest> = mutableDefaultsRequests
@@ -41,7 +46,7 @@
                     old == new
                 }
             }
-            .map { defaults[DefaultsKey(it.user, it.componentName)]!! }
+            .mapNotNull { defaults[DefaultsKey(it.user, it.componentName)] }
 
     override fun requestNewDefaults(
         user: UserHandle,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/custom/data/repository/FakeCustomTileRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/custom/data/repository/FakeCustomTileRepository.kt
new file mode 100644
index 0000000..ccf0391
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/custom/data/repository/FakeCustomTileRepository.kt
@@ -0,0 +1,58 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.tiles.impl.custom.data.repository
+
+import android.os.UserHandle
+import android.service.quicksettings.Tile
+import com.android.systemui.qs.external.FakeCustomTileStatePersister
+import com.android.systemui.qs.pipeline.shared.TileSpec
+import com.android.systemui.qs.tiles.impl.custom.data.entity.CustomTileDefaults
+import kotlin.coroutines.CoroutineContext
+import kotlinx.coroutines.flow.Flow
+
+class FakeCustomTileRepository(
+    tileSpec: TileSpec.CustomTileSpec,
+    customTileStatePersister: FakeCustomTileStatePersister,
+    testBackgroundContext: CoroutineContext,
+) : CustomTileRepository {
+
+    private val realDelegate: CustomTileRepository =
+        CustomTileRepositoryImpl(
+            tileSpec,
+            customTileStatePersister,
+            testBackgroundContext,
+        )
+
+    override suspend fun restoreForTheUserIfNeeded(user: UserHandle, isPersistable: Boolean) =
+        realDelegate.restoreForTheUserIfNeeded(user, isPersistable)
+
+    override fun getTiles(user: UserHandle): Flow<Tile> = realDelegate.getTiles(user)
+
+    override fun getTile(user: UserHandle): Tile? = realDelegate.getTile(user)
+
+    override suspend fun updateWithTile(
+        user: UserHandle,
+        newTile: Tile,
+        isPersistable: Boolean,
+    ) = realDelegate.updateWithTile(user, newTile, isPersistable)
+
+    override suspend fun updateWithDefaults(
+        user: UserHandle,
+        defaults: CustomTileDefaults,
+        isPersistable: Boolean,
+    ) = realDelegate.updateWithDefaults(user, defaults, isPersistable)
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractorKosmos.kt
index 01f4535..3d7fb6d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/domain/interactor/ActiveNotificationsInteractorKosmos.kt
@@ -17,10 +17,7 @@
 package com.android.systemui.statusbar.notification.domain.interactor
 
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.statusbar.notification.data.repository.activeNotificationListRepository
 
 val Kosmos.activeNotificationsInteractor by
-    Kosmos.Fixture {
-        ActiveNotificationsInteractor(activeNotificationListRepository, testDispatcher)
-    }
+    Kosmos.Fixture { ActiveNotificationsInteractor(activeNotificationListRepository) }
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldRemoteModule.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldRemoteModule.kt
index a639df5..2bc2db3 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldRemoteModule.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldRemoteModule.kt
@@ -16,9 +16,12 @@
 
 package com.android.systemui.unfold
 
+import android.os.Handler
 import com.android.systemui.unfold.config.UnfoldTransitionConfig
+import com.android.systemui.unfold.dagger.UnfoldMain
 import com.android.systemui.unfold.dagger.UseReceivingFilter
 import com.android.systemui.unfold.progress.RemoteUnfoldTransitionReceiver
+import com.android.systemui.unfold.updates.RotationChangeProvider
 import com.android.systemui.unfold.util.ATraceLoggerTransitionProgressListener
 import dagger.Module
 import dagger.Provides
@@ -33,16 +36,25 @@
     @Singleton
     fun provideTransitionProvider(
         config: UnfoldTransitionConfig,
-        traceListener: ATraceLoggerTransitionProgressListener,
+        traceListener: ATraceLoggerTransitionProgressListener.Factory,
         remoteReceiverProvider: Provider<RemoteUnfoldTransitionReceiver>,
     ): Optional<RemoteUnfoldTransitionReceiver> {
         if (!config.isEnabled) {
             return Optional.empty()
         }
         val remoteReceiver = remoteReceiverProvider.get()
-        remoteReceiver.addCallback(traceListener)
+        remoteReceiver.addCallback(traceListener.create("remoteReceiver"))
         return Optional.of(remoteReceiver)
     }
 
     @Provides @UseReceivingFilter fun useReceivingFilter(): Boolean = true
+
+    @Provides
+    @UnfoldMain
+    fun provideMainRotationChangeProvider(
+        rotationChangeProviderFactory: RotationChangeProvider.Factory,
+        @UnfoldMain mainHandler: Handler,
+    ): RotationChangeProvider {
+        return rotationChangeProviderFactory.create(mainHandler)
+    }
 }
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedComponent.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedComponent.kt
index c3a6cf0..31b7ccc 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedComponent.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedComponent.kt
@@ -22,12 +22,12 @@
 import android.hardware.display.DisplayManager
 import android.os.Handler
 import com.android.systemui.unfold.config.UnfoldTransitionConfig
+import com.android.systemui.unfold.dagger.UnfoldBg
 import com.android.systemui.unfold.dagger.UnfoldMain
 import com.android.systemui.unfold.dagger.UnfoldSingleThreadBg
 import com.android.systemui.unfold.progress.RemoteUnfoldTransitionReceiver
 import com.android.systemui.unfold.updates.FoldProvider
 import com.android.systemui.unfold.updates.RotationChangeProvider
-import com.android.systemui.unfold.updates.hinge.HingeAngleProvider
 import com.android.systemui.unfold.updates.screen.ScreenStatusProvider
 import com.android.systemui.unfold.util.CurrentActivityTypeProvider
 import com.android.systemui.unfold.util.UnfoldTransitionATracePrefix
@@ -63,13 +63,12 @@
             @BindsInstance @UnfoldSingleThreadBg singleThreadBgExecutor: Executor,
             @BindsInstance @UnfoldTransitionATracePrefix tracingTagPrefix: String,
             @BindsInstance displayManager: DisplayManager,
-            @BindsInstance contentResolver: ContentResolver = context.contentResolver
+            @BindsInstance @UnfoldBg bgHandler: Handler,
+            @BindsInstance contentResolver: ContentResolver = context.contentResolver,
         ): UnfoldSharedComponent
     }
 
     val unfoldTransitionProvider: Optional<UnfoldTransitionProgressProvider>
-    val hingeAngleProvider: HingeAngleProvider
-    val rotationChangeProvider: RotationChangeProvider
 }
 
 /**
@@ -94,7 +93,8 @@
     }
 
     val remoteTransitionProgress: Optional<RemoteUnfoldTransitionReceiver>
-    val rotationChangeProvider: RotationChangeProvider
+
+    @UnfoldMain fun getRotationChangeProvider(): RotationChangeProvider
 }
 
 /**
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt
index 7473ca6..42d31b3 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt
@@ -16,7 +16,10 @@
 
 package com.android.systemui.unfold
 
+import android.os.Handler
 import com.android.systemui.unfold.config.UnfoldTransitionConfig
+import com.android.systemui.unfold.dagger.UnfoldBg
+import com.android.systemui.unfold.dagger.UnfoldMain
 import com.android.systemui.unfold.progress.FixedTimingTransitionProgressProvider
 import com.android.systemui.unfold.progress.PhysicsBasedUnfoldTransitionProgressProvider
 import com.android.systemui.unfold.progress.UnfoldTransitionProgressForwarder
@@ -24,6 +27,7 @@
 import com.android.systemui.unfold.updates.FoldStateProvider
 import com.android.systemui.unfold.updates.FoldStateRepository
 import com.android.systemui.unfold.updates.FoldStateRepositoryImpl
+import com.android.systemui.unfold.updates.RotationChangeProvider
 import com.android.systemui.unfold.updates.hinge.EmptyHingeAngleProvider
 import com.android.systemui.unfold.updates.hinge.HingeAngleProvider
 import com.android.systemui.unfold.updates.hinge.HingeSensorAngleProvider
@@ -38,16 +42,18 @@
 import javax.inject.Provider
 import javax.inject.Singleton
 
-@Module(includes = [UnfoldSharedInternalModule::class])
+@Module(
+    includes =
+        [
+            UnfoldSharedInternalModule::class,
+            UnfoldRotationProviderInternalModule::class,
+            HingeAngleProviderInternalModule::class,
+            FoldStateProviderModule::class,
+        ]
+)
 class UnfoldSharedModule {
     @Provides
     @Singleton
-    fun provideFoldStateProvider(
-        deviceFoldStateProvider: DeviceFoldStateProvider
-    ): FoldStateProvider = deviceFoldStateProvider
-
-    @Provides
-    @Singleton
     fun unfoldKeyguardVisibilityProvider(
         impl: UnfoldKeyguardVisibilityManagerImpl
     ): UnfoldKeyguardVisibilityProvider = impl
@@ -60,9 +66,7 @@
 
     @Provides
     @Singleton
-    fun foldStateRepository(
-            impl: FoldStateRepositoryImpl
-    ): FoldStateRepository = impl
+    fun foldStateRepository(impl: FoldStateRepositoryImpl): FoldStateRepository = impl
 }
 
 /**
@@ -77,17 +81,69 @@
     fun unfoldTransitionProgressProvider(
         config: UnfoldTransitionConfig,
         scaleAwareProviderFactory: ScaleAwareTransitionProgressProvider.Factory,
+        tracingListener: ATraceLoggerTransitionProgressListener.Factory,
+        physicsBasedUnfoldTransitionProgressProvider:
+            PhysicsBasedUnfoldTransitionProgressProvider.Factory,
+        fixedTimingTransitionProgressProvider: Provider<FixedTimingTransitionProgressProvider>,
+        foldStateProvider: FoldStateProvider,
+        @UnfoldMain mainHandler: Handler,
+    ): Optional<UnfoldTransitionProgressProvider> {
+        return createOptionalUnfoldTransitionProgressProvider(
+            config = config,
+            scaleAwareProviderFactory = scaleAwareProviderFactory,
+            tracingListener = tracingListener.create("MainThread"),
+            physicsBasedUnfoldTransitionProgressProvider =
+                physicsBasedUnfoldTransitionProgressProvider,
+            fixedTimingTransitionProgressProvider = fixedTimingTransitionProgressProvider,
+            foldStateProvider = foldStateProvider,
+            progressHandler = mainHandler,
+        )
+    }
+
+    @Provides
+    @Singleton
+    @UnfoldBg
+    fun unfoldBgTransitionProgressProvider(
+        config: UnfoldTransitionConfig,
+        scaleAwareProviderFactory: ScaleAwareTransitionProgressProvider.Factory,
+        tracingListener: ATraceLoggerTransitionProgressListener.Factory,
+        physicsBasedUnfoldTransitionProgressProvider:
+            PhysicsBasedUnfoldTransitionProgressProvider.Factory,
+        fixedTimingTransitionProgressProvider: Provider<FixedTimingTransitionProgressProvider>,
+        @UnfoldBg bgFoldStateProvider: FoldStateProvider,
+        @UnfoldBg bgHandler: Handler,
+    ): Optional<UnfoldTransitionProgressProvider> {
+        return createOptionalUnfoldTransitionProgressProvider(
+            config = config,
+            scaleAwareProviderFactory = scaleAwareProviderFactory,
+            tracingListener = tracingListener.create("BgThread"),
+            physicsBasedUnfoldTransitionProgressProvider =
+                physicsBasedUnfoldTransitionProgressProvider,
+            fixedTimingTransitionProgressProvider = fixedTimingTransitionProgressProvider,
+            foldStateProvider = bgFoldStateProvider,
+            progressHandler = bgHandler,
+        )
+    }
+
+    private fun createOptionalUnfoldTransitionProgressProvider(
+        config: UnfoldTransitionConfig,
+        scaleAwareProviderFactory: ScaleAwareTransitionProgressProvider.Factory,
         tracingListener: ATraceLoggerTransitionProgressListener,
         physicsBasedUnfoldTransitionProgressProvider:
-            Provider<PhysicsBasedUnfoldTransitionProgressProvider>,
+            PhysicsBasedUnfoldTransitionProgressProvider.Factory,
         fixedTimingTransitionProgressProvider: Provider<FixedTimingTransitionProgressProvider>,
+        foldStateProvider: FoldStateProvider,
+        progressHandler: Handler,
     ): Optional<UnfoldTransitionProgressProvider> {
         if (!config.isEnabled) {
             return Optional.empty()
         }
         val baseProgressProvider =
             if (config.isHingeAngleEnabled) {
-                physicsBasedUnfoldTransitionProgressProvider.get()
+                physicsBasedUnfoldTransitionProgressProvider.create(
+                    foldStateProvider,
+                    progressHandler
+                )
             } else {
                 fixedTimingTransitionProgressProvider.get()
             }
@@ -101,26 +157,105 @@
     }
 
     @Provides
+    @Singleton
+    fun provideProgressForwarder(
+        config: UnfoldTransitionConfig,
+        progressForwarder: Provider<UnfoldTransitionProgressForwarder>
+    ): Optional<UnfoldTransitionProgressForwarder> {
+        if (!config.isEnabled) {
+            return Optional.empty()
+        }
+        return Optional.of(progressForwarder.get())
+    }
+}
+
+/**
+ * Provides [FoldStateProvider]. The [UnfoldBg] annotated binding sends progress in the [UnfoldBg]
+ * handler.
+ */
+@Module
+internal class FoldStateProviderModule {
+    @Provides
+    @Singleton
+    fun provideFoldStateProvider(
+        factory: DeviceFoldStateProvider.Factory,
+        @UnfoldMain hingeAngleProvider: HingeAngleProvider,
+        @UnfoldMain rotationChangeProvider: RotationChangeProvider,
+        @UnfoldMain mainHandler: Handler,
+    ): FoldStateProvider =
+        factory.create(
+            hingeAngleProvider,
+            rotationChangeProvider,
+            progressHandler = mainHandler
+        )
+
+    @Provides
+    @Singleton
+    @UnfoldBg
+    fun provideBgFoldStateProvider(
+        factory: DeviceFoldStateProvider.Factory,
+        @UnfoldBg hingeAngleProvider: HingeAngleProvider,
+        @UnfoldBg rotationChangeProvider: RotationChangeProvider,
+        @UnfoldBg bgHandler: Handler,
+    ): FoldStateProvider =
+        factory.create(
+            hingeAngleProvider,
+            rotationChangeProvider,
+            progressHandler = bgHandler
+        )
+}
+
+/** Provides bindings for both [UnfoldMain] and [UnfoldBg] [HingeAngleProvider]. */
+@Module
+internal class HingeAngleProviderInternalModule {
+    @Provides
+    @UnfoldMain
     fun hingeAngleProvider(
         config: UnfoldTransitionConfig,
-        hingeAngleSensorProvider: Provider<HingeSensorAngleProvider>
+        @UnfoldMain handler: Handler,
+        hingeAngleSensorProvider: HingeSensorAngleProvider.Factory
     ): HingeAngleProvider {
         return if (config.isHingeAngleEnabled) {
-            hingeAngleSensorProvider.get()
+            hingeAngleSensorProvider.create(handler)
         } else {
             EmptyHingeAngleProvider
         }
     }
 
     @Provides
-    @Singleton
-    fun provideProgressForwarder(
-            config: UnfoldTransitionConfig,
-            progressForwarder: Provider<UnfoldTransitionProgressForwarder>
-    ): Optional<UnfoldTransitionProgressForwarder> {
-        if (!config.isEnabled) {
-            return Optional.empty()
+    @UnfoldBg
+    fun hingeAngleProviderBg(
+        config: UnfoldTransitionConfig,
+        @UnfoldBg handler: Handler,
+        hingeAngleSensorProvider: HingeSensorAngleProvider.Factory
+    ): HingeAngleProvider {
+        return if (config.isHingeAngleEnabled) {
+            hingeAngleSensorProvider.create(handler)
+        } else {
+            EmptyHingeAngleProvider
         }
-        return Optional.of(progressForwarder.get())
+    }
+}
+
+@Module
+internal class UnfoldRotationProviderInternalModule {
+    @Provides
+    @Singleton
+    @UnfoldMain
+    fun provideRotationChangeProvider(
+        rotationChangeProviderFactory: RotationChangeProvider.Factory,
+        @UnfoldMain mainHandler: Handler,
+    ): RotationChangeProvider {
+        return rotationChangeProviderFactory.create(mainHandler)
+    }
+
+    @Provides
+    @Singleton
+    @UnfoldBg
+    fun provideBgRotationChangeProvider(
+        rotationChangeProviderFactory: RotationChangeProvider.Factory,
+        @UnfoldBg bgHandler: Handler,
+    ): RotationChangeProvider {
+        return rotationChangeProviderFactory.create(bgHandler)
     }
 }
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldTransitionFactory.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldTransitionFactory.kt
index 1839919..1cbaf31 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldTransitionFactory.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldTransitionFactory.kt
@@ -48,6 +48,7 @@
         singleThreadBgExecutor: Executor,
         tracingTagPrefix: String,
         displayManager: DisplayManager,
+        bgHandler: Handler,
 ): UnfoldSharedComponent =
         DaggerUnfoldSharedComponent.factory()
                 .create(
@@ -62,6 +63,7 @@
                         singleThreadBgExecutor,
                         tracingTagPrefix,
                         displayManager,
+                        bgHandler,
                 )
 
 /**
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/dagger/UnfoldBg.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/dagger/UnfoldBg.kt
new file mode 100644
index 0000000..7cd4419
--- /dev/null
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/dagger/UnfoldBg.kt
@@ -0,0 +1,20 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+package com.android.systemui.unfold.dagger
+
+import javax.inject.Qualifier
+
+/** Annotation for background computations related to unfold lib. */
+@Qualifier @Retention(AnnotationRetention.RUNTIME) annotation class UnfoldBg
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProvider.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProvider.kt
index f8f168b..907bf46 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProvider.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProvider.kt
@@ -20,6 +20,7 @@
 import android.animation.ObjectAnimator
 import android.animation.ValueAnimator
 import android.content.Context
+import android.os.Handler
 import android.os.Trace
 import android.util.FloatProperty
 import android.util.Log
@@ -38,13 +39,25 @@
 import com.android.systemui.unfold.updates.FoldStateProvider.FoldUpdate
 import com.android.systemui.unfold.updates.FoldStateProvider.FoldUpdatesListener
 import com.android.systemui.unfold.updates.name
-import javax.inject.Inject
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 
-/** Maps fold updates to unfold transition progress using DynamicAnimation. */
+/**
+ * Maps fold updates to unfold transition progress using DynamicAnimation.
+ *
+ * Note that all variable accesses must be done in the [Handler] provided in the constructor, that
+ * might be different than [mainHandler]. When a custom handler is provided, the [SpringAnimation]
+ * uses a scheduler different than the default one.
+ */
 class PhysicsBasedUnfoldTransitionProgressProvider
-@Inject
-constructor(context: Context, private val foldStateProvider: FoldStateProvider) :
-    UnfoldTransitionProgressProvider, FoldUpdatesListener, DynamicAnimation.OnAnimationEndListener {
+@AssistedInject
+constructor(
+    context: Context,
+    private val schedulerFactory: UnfoldFrameCallbackScheduler.Factory,
+    @Assisted private val foldStateProvider: FoldStateProvider,
+    @Assisted private val progressHandler: Handler,
+) : UnfoldTransitionProgressProvider, FoldUpdatesListener, DynamicAnimation.OnAnimationEndListener {
 
     private val emphasizedInterpolator =
         loadInterpolator(context, android.R.interpolator.fast_out_extra_slow_in)
@@ -63,6 +76,7 @@
 
     private var transitionProgress: Float = 0.0f
         set(value) {
+            assertInProgressThread()
             if (isTransitionRunning) {
                 listeners.forEach { it.onTransitionProgress(value) }
             }
@@ -72,8 +86,14 @@
     private val listeners: MutableList<TransitionProgressListener> = mutableListOf()
 
     init {
-        foldStateProvider.addCallback(this)
-        foldStateProvider.start()
+        progressHandler.post {
+            // The scheduler needs to be created in the progress handler in order to get the correct
+            // choreographer and frame callbacks. This is because the choreographer can be get only
+            // as a thread local.
+            springAnimation.scheduler = schedulerFactory.create()
+            foldStateProvider.addCallback(this)
+            foldStateProvider.start()
+        }
     }
 
     override fun destroy() {
@@ -81,6 +101,8 @@
     }
 
     override fun onHingeAngleUpdate(angle: Float) {
+        assertInProgressThread()
+
         if (!isTransitionRunning || isAnimatedCancelRunning) return
         val progress = saturate(angle / FINAL_HINGE_ANGLE_POSITION)
         springAnimation.animateToFinalPosition(progress)
@@ -90,6 +112,7 @@
         if (amount < low) low else if (amount > high) high else amount
 
     override fun onFoldUpdate(@FoldUpdate update: Int) {
+        assertInProgressThread()
         when (update) {
             FOLD_UPDATE_FINISH_FULL_OPEN,
             FOLD_UPDATE_FINISH_HALF_OPEN -> {
@@ -148,6 +171,7 @@
     }
 
     private fun cancelTransition(endValue: Float, animate: Boolean) {
+        assertInProgressThread()
         if (isTransitionRunning && animate) {
             if (endValue == 1.0f && !isAnimatedCancelRunning) {
                 listeners.forEach { it.onTransitionFinishing() }
@@ -165,7 +189,6 @@
             isAnimatedCancelRunning = false
             isTransitionRunning = false
             springAnimation.cancel()
-
             cannedAnimator?.removeAllListeners()
             cannedAnimator?.cancel()
             cannedAnimator = null
@@ -182,7 +205,7 @@
         animation: DynamicAnimation<out DynamicAnimation<*>>,
         canceled: Boolean,
         value: Float,
-        velocity: Float
+        velocity: Float,
     ) {
         if (isAnimatedCancelRunning) {
             cancelTransition(value, animate = false)
@@ -202,6 +225,7 @@
     }
 
     private fun startTransition(startValue: Float) {
+        assertInProgressThread()
         if (!isTransitionRunning) onStartTransition()
 
         springAnimation.apply {
@@ -221,14 +245,16 @@
     }
 
     override fun addCallback(listener: TransitionProgressListener) {
-        listeners.add(listener)
+        progressHandler.post { listeners.add(listener) }
     }
 
     override fun removeCallback(listener: TransitionProgressListener) {
-        listeners.remove(listener)
+        progressHandler.post { listeners.remove(listener) }
     }
 
     private fun startCannedCancelAnimation() {
+        assertInProgressThread()
+
         cannedAnimator?.cancel()
         cannedAnimator = null
 
@@ -264,7 +290,7 @@
 
         override fun setValue(
             provider: PhysicsBasedUnfoldTransitionProgressProvider,
-            value: Float
+            value: Float,
         ) {
             provider.transitionProgress = value
         }
@@ -272,6 +298,25 @@
         override fun get(provider: PhysicsBasedUnfoldTransitionProgressProvider): Float =
             provider.transitionProgress
     }
+
+    private fun assertInProgressThread() {
+        check(progressHandler.looper.isCurrentThread) {
+            val progressThread = progressHandler.looper.thread
+            val thisThread = Thread.currentThread()
+            """should be called from the progress thread.
+                progressThread=$progressThread tid=${progressThread.id}
+                Thread.currentThread()=$thisThread tid=${thisThread.id}"""
+                .trimMargin()
+        }
+    }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(
+            foldStateProvider: FoldStateProvider,
+            handler: Handler,
+        ): PhysicsBasedUnfoldTransitionProgressProvider
+    }
 }
 
 private const val TAG = "PhysicsBasedUnfoldTransitionProgressProvider"
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/progress/UnfoldFrameCallbackScheduler.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/progress/UnfoldFrameCallbackScheduler.kt
new file mode 100644
index 0000000..1dffd84
--- /dev/null
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/progress/UnfoldFrameCallbackScheduler.kt
@@ -0,0 +1,59 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.systemui.unfold.progress
+
+import android.os.Looper
+import android.view.Choreographer
+import androidx.dynamicanimation.animation.FrameCallbackScheduler
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+
+/**
+ * Scheduler that posts animation progresses on a thread different than the ui one.
+ *
+ * The following is taken from [AnimationHandler.FrameCallbackScheduler16]. It is extracted here as
+ * there are no guarantees which implementation the [DynamicAnimation] class would use otherwise.
+ * This allows classes using [DynamicAnimation] to be created in any thread, but still use the
+ * scheduler for a specific thread.
+ *
+ * Technically the [AssistedInject] is not needed: it's just to have a nicer factory with a
+ * documentation snippet instead of using a plain dagger provider.
+ */
+class UnfoldFrameCallbackScheduler @AssistedInject constructor() : FrameCallbackScheduler {
+
+    private val choreographer = Choreographer.getInstance()
+    private val looper =
+        Looper.myLooper() ?: error("This should be created in a thread with a looper.")
+
+    override fun postFrameCallback(frameCallback: Runnable) {
+        choreographer.postFrameCallback { frameCallback.run() }
+    }
+
+    override fun isCurrentThread(): Boolean {
+        return looper.isCurrentThread
+    }
+
+    @AssistedFactory
+    interface Factory {
+        /**
+         * Creates a [FrameCallbackScheduler] that uses [Choreographer] to post frame callbacks.
+         *
+         * Note that the choreographer used depends on the thread this [create] is called on, as it
+         * is get from a thread static attribute.
+         */
+        fun create(): UnfoldFrameCallbackScheduler
+    }
+}
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
index 003013e..77f637b 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
@@ -23,37 +23,34 @@
 import androidx.core.util.Consumer
 import com.android.systemui.unfold.compat.INNER_SCREEN_SMALLEST_SCREEN_WIDTH_THRESHOLD_DP
 import com.android.systemui.unfold.config.UnfoldTransitionConfig
-import com.android.systemui.unfold.dagger.UnfoldMain
-import com.android.systemui.unfold.updates.FoldStateProvider.FoldUpdate
-import com.android.systemui.unfold.updates.FoldStateProvider.FoldUpdatesListener
-import com.android.systemui.unfold.updates.RotationChangeProvider.RotationListener
 import com.android.systemui.unfold.updates.hinge.FULLY_CLOSED_DEGREES
 import com.android.systemui.unfold.updates.hinge.FULLY_OPEN_DEGREES
 import com.android.systemui.unfold.updates.hinge.HingeAngleProvider
 import com.android.systemui.unfold.updates.screen.ScreenStatusProvider
 import com.android.systemui.unfold.util.CurrentActivityTypeProvider
 import com.android.systemui.unfold.util.UnfoldKeyguardVisibilityProvider
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+import java.util.concurrent.CopyOnWriteArrayList
 import java.util.concurrent.Executor
-import javax.inject.Inject
 
 class DeviceFoldStateProvider
-@Inject
+@AssistedInject
 constructor(
     config: UnfoldTransitionConfig,
-    private val hingeAngleProvider: HingeAngleProvider,
+    private val context: Context,
     private val screenStatusProvider: ScreenStatusProvider,
-    private val foldProvider: FoldProvider,
     private val activityTypeProvider: CurrentActivityTypeProvider,
     private val unfoldKeyguardVisibilityProvider: UnfoldKeyguardVisibilityProvider,
-    private val rotationChangeProvider: RotationChangeProvider,
-    private val context: Context,
-    @UnfoldMain private val mainExecutor: Executor,
-    @UnfoldMain private val handler: Handler
+    private val foldProvider: FoldProvider,
+    @Assisted private val hingeAngleProvider: HingeAngleProvider,
+    @Assisted private val rotationChangeProvider: RotationChangeProvider,
+    @Assisted private val progressHandler: Handler,
 ) : FoldStateProvider {
+    private val outputListeners = CopyOnWriteArrayList<FoldStateProvider.FoldUpdatesListener>()
 
-    private val outputListeners: MutableList<FoldUpdatesListener> = mutableListOf()
-
-    @FoldUpdate private var lastFoldUpdate: Int? = null
+    @FoldStateProvider.FoldUpdate private var lastFoldUpdate: Int? = null
 
     @FloatRange(from = 0.0, to = 180.0) private var lastHingeAngle: Float = 0f
     @FloatRange(from = 0.0, to = 180.0) private var lastHingeAngleBeforeTransition: Float = 0f
@@ -61,11 +58,9 @@
     private val hingeAngleListener = HingeAngleListener()
     private val screenListener = ScreenStatusListener()
     private val foldStateListener = FoldStateListener()
-    private val mainLooper = handler.looper
     private val timeoutRunnable = Runnable { cancelAnimation() }
-    private val rotationListener = RotationListener {
-        if (isTransitionInProgress) cancelAnimation()
-    }
+    private val rotationListener = FoldRotationListener()
+    private val progressExecutor = Executor { progressHandler.post(it) }
 
     /**
      * Time after which [FOLD_UPDATE_FINISH_HALF_OPEN] is emitted following a
@@ -80,9 +75,9 @@
     private var isStarted = false
 
     override fun start() {
-        assertMainThread()
         if (isStarted) return
-        foldProvider.registerCallback(foldStateListener, mainExecutor)
+        foldProvider.registerCallback(foldStateListener, progressExecutor)
+        // TODO(b/277879146): get callbacks in the background
         screenStatusProvider.addCallback(screenListener)
         hingeAngleProvider.addCallback(hingeAngleListener)
         rotationChangeProvider.addCallback(rotationListener)
@@ -91,7 +86,6 @@
     }
 
     override fun stop() {
-        assertMainThread()
         screenStatusProvider.removeCallback(screenListener)
         foldProvider.unregisterCallback(foldStateListener)
         hingeAngleProvider.removeCallback(hingeAngleListener)
@@ -101,11 +95,11 @@
         isStarted = false
     }
 
-    override fun addCallback(listener: FoldUpdatesListener) {
+    override fun addCallback(listener: FoldStateProvider.FoldUpdatesListener) {
         outputListeners.add(listener)
     }
 
-    override fun removeCallback(listener: FoldUpdatesListener) {
+    override fun removeCallback(listener: FoldStateProvider.FoldUpdatesListener) {
         outputListeners.remove(listener)
     }
 
@@ -121,6 +115,7 @@
                 lastFoldUpdate == FOLD_UPDATE_START_CLOSING
 
     private fun onHingeAngle(angle: Float) {
+        assertInProgressThread()
         if (DEBUG) {
             Log.d(
                 TAG,
@@ -131,14 +126,14 @@
         }
 
         val currentDirection =
-                if (angle < lastHingeAngle) FOLD_UPDATE_START_CLOSING else FOLD_UPDATE_START_OPENING
+            if (angle < lastHingeAngle) FOLD_UPDATE_START_CLOSING else FOLD_UPDATE_START_OPENING
         if (isTransitionInProgress && currentDirection != lastFoldUpdate) {
             lastHingeAngleBeforeTransition = lastHingeAngle
         }
 
         val isClosing = angle < lastHingeAngleBeforeTransition
         val transitionUpdate =
-                if (isClosing) FOLD_UPDATE_START_CLOSING else FOLD_UPDATE_START_OPENING
+            if (isClosing) FOLD_UPDATE_START_CLOSING else FOLD_UPDATE_START_OPENING
         val angleChangeSurpassedThreshold =
             Math.abs(angle - lastHingeAngleBeforeTransition) > HINGE_ANGLE_CHANGE_THRESHOLD_DEGREES
         val isFullyOpened = FULLY_OPEN_DEGREES - angle < FULLY_OPEN_THRESHOLD_DEGREES
@@ -150,12 +145,12 @@
             angleChangeSurpassedThreshold && // Do not react immediately to small changes in angle
                 eventNotAlreadyDispatched && // we haven't sent transition event already
                 !isFullyOpened && // do not send transition event if we are in fully opened hinge
-                                  // angle range as closing threshold could overlap this range
+                // angle range as closing threshold could overlap this range
                 screenAvailableEventSent && // do not send transition event if we are still in the
-                                            // process of turning on the inner display
+                // process of turning on the inner display
                 isClosingThresholdMet(angle) && // hinge angle is below certain threshold.
                 isOnLargeScreen // Avoids sending closing event when on small screen.
-                                // Start event is sent regardless due to hall sensor.
+        // Start event is sent regardless due to hall sensor.
         ) {
             notifyFoldUpdate(transitionUpdate, lastHingeAngle)
         }
@@ -202,6 +197,7 @@
 
     private inner class FoldStateListener : FoldProvider.FoldCallback {
         override fun onFoldUpdated(isFolded: Boolean) {
+            assertInProgressThread()
             this@DeviceFoldStateProvider.isFolded = isFolded
             lastHingeAngle = FULLY_CLOSED_DEGREES
 
@@ -218,7 +214,14 @@
         }
     }
 
-    private fun notifyFoldUpdate(@FoldUpdate update: Int, angle: Float) {
+    private inner class FoldRotationListener : RotationChangeProvider.RotationListener {
+        override fun onRotationChanged(newRotation: Int) {
+            assertInProgressThread()
+            if (isTransitionInProgress) cancelAnimation()
+        }
+    }
+
+    private fun notifyFoldUpdate(@FoldStateProvider.FoldUpdate update: Int, angle: Float) {
         if (DEBUG) {
             Log.d(TAG, update.name())
         }
@@ -236,11 +239,11 @@
         if (isTransitionInProgress) {
             cancelTimeout()
         }
-        handler.postDelayed(timeoutRunnable, halfOpenedTimeoutMillis.toLong())
+        progressHandler.postDelayed(timeoutRunnable, halfOpenedTimeoutMillis.toLong())
     }
 
     private fun cancelTimeout() {
-        handler.removeCallbacks(timeoutRunnable)
+        progressHandler.removeCallbacks(timeoutRunnable)
     }
 
     private fun cancelAnimation(): Unit =
@@ -249,42 +252,61 @@
     private inner class ScreenStatusListener : ScreenStatusProvider.ScreenListener {
 
         override fun onScreenTurnedOn() {
-            // Trigger this event only if we are unfolded and this is the first screen
-            // turned on event since unfold started. This prevents running the animation when
-            // turning on the internal display using the power button.
-            // Initially isUnfoldHandled is true so it will be reset to false *only* when we
-            // receive 'folded' event. If SystemUI started when device is already folded it will
-            // still receive 'folded' event on startup.
-            if (!isFolded && !isUnfoldHandled) {
-                outputListeners.forEach { it.onUnfoldedScreenAvailable() }
-                isUnfoldHandled = true
+            executeInProgressThread {
+                // Trigger this event only if we are unfolded and this is the first screen
+                // turned on event since unfold started. This prevents running the animation when
+                // turning on the internal display using the power button.
+                // Initially isUnfoldHandled is true so it will be reset to false *only* when we
+                // receive 'folded' event. If SystemUI started when device is already folded it will
+                // still receive 'folded' event on startup.
+                if (!isFolded && !isUnfoldHandled) {
+                    outputListeners.forEach { it.onUnfoldedScreenAvailable() }
+                    isUnfoldHandled = true
+                }
             }
         }
 
         override fun markScreenAsTurnedOn() {
-            if (!isFolded) {
-                isUnfoldHandled = true
+            executeInProgressThread {
+                if (!isFolded) {
+                    isUnfoldHandled = true
+                }
             }
         }
 
         override fun onScreenTurningOn() {
-            isScreenOn = true
-            updateHingeAngleProviderState()
+            executeInProgressThread {
+                isScreenOn = true
+                updateHingeAngleProviderState()
+            }
         }
 
         override fun onScreenTurningOff() {
-            isScreenOn = false
-            updateHingeAngleProviderState()
+            executeInProgressThread {
+                isScreenOn = false
+                updateHingeAngleProviderState()
+            }
+        }
+
+        /**
+         * Needed just for compatibility while not all data sources are providing data in the
+         * background.
+         *
+         * TODO(b/277879146): Remove once ScreeStatusProvider provides in the background.
+         */
+        private fun executeInProgressThread(f: () -> Unit) {
+            progressHandler.post { f() }
         }
     }
 
     private fun isOnLargeScreen(): Boolean {
-      return context.resources.configuration.smallestScreenWidthDp >
-          INNER_SCREEN_SMALLEST_SCREEN_WIDTH_THRESHOLD_DP
+        return context.resources.configuration.smallestScreenWidthDp >
+            INNER_SCREEN_SMALLEST_SCREEN_WIDTH_THRESHOLD_DP
     }
 
     /** While the screen is off or the device is folded, hinge angle updates are not needed. */
     private fun updateHingeAngleProviderState() {
+        assertInProgressThread()
         if (isScreenOn && !isFolded) {
             hingeAngleProvider.start()
         } else {
@@ -294,20 +316,34 @@
 
     private inner class HingeAngleListener : Consumer<Float> {
         override fun accept(angle: Float) {
+            assertInProgressThread()
             onHingeAngle(angle)
         }
     }
 
-    private fun assertMainThread() {
-        check(mainLooper.isCurrentThread) {
-            ("should be called from the main thread." +
-                    " sMainLooper.threadName=" + mainLooper.thread.name +
-                    " Thread.currentThread()=" + Thread.currentThread().name)
+    private fun assertInProgressThread() {
+        check(progressHandler.looper.isCurrentThread) {
+            val progressThread = progressHandler.looper.thread
+            val thisThread = Thread.currentThread()
+            """should be called from the progress thread.
+                progressThread=$progressThread tid=${progressThread.id}
+                Thread.currentThread()=$thisThread tid=${thisThread.id}"""
+                .trimMargin()
         }
     }
+
+    @AssistedFactory
+    interface Factory {
+        /** Creates a [DeviceFoldStateProvider] using the provided dependencies. */
+        fun create(
+            hingeAngleProvider: HingeAngleProvider,
+            rotationChangeProvider: RotationChangeProvider,
+            progressHandler: Handler,
+        ): DeviceFoldStateProvider
+    }
 }
 
-fun @receiver:FoldUpdate Int.name() =
+fun @receiver:FoldStateProvider.FoldUpdate Int.name() =
     when (this) {
         FOLD_UPDATE_START_OPENING -> "START_OPENING"
         FOLD_UPDATE_START_CLOSING -> "START_CLOSING"
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/RotationChangeProvider.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/RotationChangeProvider.kt
index ce8f1a1..82ea362 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/RotationChangeProvider.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/RotationChangeProvider.kt
@@ -20,20 +20,21 @@
 import android.hardware.display.DisplayManager
 import android.os.Handler
 import android.os.RemoteException
-import com.android.systemui.unfold.dagger.UnfoldMain
 import com.android.systemui.unfold.util.CallbackController
-import javax.inject.Inject
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 
 /**
- * Allows to subscribe to rotation changes. Updates are provided for the display associated
- * to [context].
+ * Allows to subscribe to rotation changes. Updates are provided for the display associated to
+ * [context].
  */
 class RotationChangeProvider
-@Inject
+@AssistedInject
 constructor(
     private val displayManager: DisplayManager,
     private val context: Context,
-    @UnfoldMain private val mainHandler: Handler,
+    @Assisted private val handler: Handler,
 ) : CallbackController<RotationChangeProvider.RotationListener> {
 
     private val listeners = mutableListOf<RotationListener>()
@@ -42,7 +43,7 @@
     private var lastRotation: Int? = null
 
     override fun addCallback(listener: RotationListener) {
-        mainHandler.post {
+        handler.post {
             if (listeners.isEmpty()) {
                 subscribeToRotation()
             }
@@ -51,7 +52,7 @@
     }
 
     override fun removeCallback(listener: RotationListener) {
-        mainHandler.post {
+        handler.post {
             listeners -= listener
             if (listeners.isEmpty()) {
                 unsubscribeToRotation()
@@ -62,7 +63,7 @@
 
     private fun subscribeToRotation() {
         try {
-            displayManager.registerDisplayListener(displayListener, mainHandler)
+            displayManager.registerDisplayListener(displayListener, handler)
         } catch (e: RemoteException) {
             throw e.rethrowFromSystemServer()
         }
@@ -100,4 +101,10 @@
 
         override fun onDisplayRemoved(displayId: Int) {}
     }
+
+    @AssistedFactory
+    interface Factory {
+        /** Creates a new [RotationChangeProvider] that provides updated using [handler]. */
+        fun create(handler: Handler): RotationChangeProvider
+    }
 }
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/hinge/HingeSensorAngleProvider.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/hinge/HingeSensorAngleProvider.kt
index 89fb12e..14c4cc0 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/hinge/HingeSensorAngleProvider.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/hinge/HingeSensorAngleProvider.kt
@@ -18,21 +18,26 @@
 import android.hardware.SensorEvent
 import android.hardware.SensorEventListener
 import android.hardware.SensorManager
+import android.os.Handler
 import android.os.Trace
 import androidx.core.util.Consumer
 import com.android.systemui.unfold.dagger.UnfoldSingleThreadBg
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+import java.util.concurrent.CopyOnWriteArrayList
 import java.util.concurrent.Executor
-import javax.inject.Inject
 
 internal class HingeSensorAngleProvider
-@Inject
+@AssistedInject
 constructor(
     private val sensorManager: SensorManager,
-    @UnfoldSingleThreadBg private val singleThreadBgExecutor: Executor
+    @UnfoldSingleThreadBg private val singleThreadBgExecutor: Executor,
+    @Assisted private val listenerHandler: Handler,
 ) : HingeAngleProvider {
 
     private val sensorListener = HingeAngleSensorListener()
-    private val listeners: MutableList<Consumer<Float>> = arrayListOf()
+    private val listeners: MutableList<Consumer<Float>> = CopyOnWriteArrayList()
     var started = false
 
     override fun start() {
@@ -43,7 +48,8 @@
             sensorManager.registerListener(
                 sensorListener,
                 sensor,
-                SensorManager.SENSOR_DELAY_FASTEST
+                SensorManager.SENSOR_DELAY_FASTEST,
+                listenerHandler
             )
             Trace.endSection()
 
@@ -75,4 +81,10 @@
             listeners.forEach { it.accept(event.values[0]) }
         }
     }
+
+    @AssistedFactory
+    interface Factory {
+        /** Creates an [HingeSensorAngleProvider] that sends updates using [handler]. */
+        fun create(handler: Handler): HingeSensorAngleProvider
+    }
 }
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/util/ATraceLoggerTransitionProgressListener.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/util/ATraceLoggerTransitionProgressListener.kt
index d8bc018..a31896a 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/util/ATraceLoggerTransitionProgressListener.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/util/ATraceLoggerTransitionProgressListener.kt
@@ -16,7 +16,9 @@
 
 import android.os.Trace
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener
-import javax.inject.Inject
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 import javax.inject.Qualifier
 
 /**
@@ -26,11 +28,11 @@
  * for each fold/unfold: in (1) systemui and (2) launcher process.
  */
 class ATraceLoggerTransitionProgressListener
-@Inject
-internal constructor(@UnfoldTransitionATracePrefix tracePrefix: String) :
+@AssistedInject
+internal constructor(@UnfoldTransitionATracePrefix tracePrefix: String, @Assisted details: String) :
     TransitionProgressListener {
 
-    private val traceName = "$tracePrefix#$UNFOLD_TRANSITION_TRACE_NAME"
+    private val traceName = "$tracePrefix$details#$UNFOLD_TRANSITION_TRACE_NAME"
 
     override fun onTransitionStarted() {
         Trace.beginAsyncSection(traceName, /* cookie= */ 0)
@@ -43,6 +45,12 @@
     override fun onTransitionProgress(progress: Float) {
         Trace.setCounter(traceName, (progress * 100).toLong())
     }
+
+    @AssistedFactory
+    interface Factory {
+        /** Creates an [ATraceLoggerTransitionProgressListener] with [details] in the track name. */
+        fun create(details: String): ATraceLoggerTransitionProgressListener
+    }
 }
 
 private const val UNFOLD_TRANSITION_TRACE_NAME = "FoldUnfoldTransitionInProgress"
diff --git a/services/accessibility/java/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandler.java b/services/accessibility/java/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandler.java
index d31b1ef..e3797c9 100644
--- a/services/accessibility/java/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandler.java
+++ b/services/accessibility/java/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandler.java
@@ -258,6 +258,11 @@
                 public void logMagnificationTripleTap(boolean enabled) {
                     AccessibilityStatsLogUtils.logMagnificationTripleTap(enabled);
                 }
+
+                @Override
+                public void logMagnificationTwoFingerTripleTap(boolean enabled) {
+                    AccessibilityStatsLogUtils.logMagnificationTwoFingerTripleTap(enabled);
+                }
             };
         }
 
@@ -419,6 +424,7 @@
     /** An interface that allows testing magnification log events. */
     interface MagnificationLogger {
         void logMagnificationTripleTap(boolean enabled);
+        void logMagnificationTwoFingerTripleTap(boolean enabled);
     }
 
     interface State {
@@ -987,12 +993,14 @@
                             mDisplayId, event.getX(), event.getY())) {
                         transitionToDelegatingStateAndClear();
 
+                    } else if (isMultiFingerMultiTapTriggered(/* targetTapCount= */ 3, event)) {
+                        // Placing multiple fingers before a single finger, because achieving a
+                        // multi finger multi tap also means achieving a single finger triple tap
+                        onTripleTap(event);
+
                     } else if (isMultiTapTriggered(3 /* taps */)) {
                         onTripleTap(/* up */ event);
 
-                    } else if (isMultiFingerMultiTapTriggered(/* targetTapCount= */ 3, event)) {
-                        onTripleTap(event);
-
                     } else if (
                             // Possible to be false on: 3tap&drag -> scale -> PTR_UP -> UP
                             isFingerDown()
@@ -1026,6 +1034,11 @@
                 mCompletedTapCount++;
                 mIsTwoFingerCountReached = false;
             }
+
+            if (mDetectTwoFingerTripleTap && mCompletedTapCount > 2) {
+                final boolean enabled = !isActivated();
+                mMagnificationLogger.logMagnificationTwoFingerTripleTap(enabled);
+            }
             return mDetectTwoFingerTripleTap && mCompletedTapCount == targetTapCount;
         }
 
@@ -1037,6 +1050,29 @@
             mFirstPointerDownLocation.set(Float.NaN, Float.NaN);
             mSecondPointerDownLocation.set(Float.NaN, Float.NaN);
         }
+
+        void transitionToViewportDraggingStateAndClear(MotionEvent down) {
+
+            if (DEBUG_DETECTING) Slog.i(mLogTag, "onTripleTapAndHold()");
+            final boolean shortcutTriggered = mShortcutTriggered;
+
+            // Only log the 3tap and hold event
+            if (!shortcutTriggered) {
+                final boolean enabled = !isActivated();
+                if (mCompletedTapCount == 2) {
+                    // Two finger triple tap and hold
+                    mMagnificationLogger.logMagnificationTwoFingerTripleTap(enabled);
+                } else {
+                    // Triple tap and hold also belongs to triple tap event
+                    mMagnificationLogger.logMagnificationTripleTap(enabled);
+                }
+            }
+            clear();
+
+            mViewportDraggingState.prepareForZoomInTemporary(shortcutTriggered);
+            zoomInTemporary(down.getX(), down.getY(), shortcutTriggered);
+            transitionTo(mViewportDraggingState);
+        }
     }
 
     /**
@@ -1416,8 +1452,6 @@
 
             // Only log the 3tap and hold event
             if (!shortcutTriggered) {
-                // TODO:(b/309534286): Add metrics for two-finger triple-tap and fix
-                //  the log two-finger bug before enabling the flag
                 // Triple tap and hold also belongs to triple tap event
                 final boolean enabled = !isActivated();
                 mMagnificationLogger.logMagnificationTripleTap(enabled);
diff --git a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
index e6bfeb7..4987fbc 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
@@ -39,7 +39,6 @@
 import android.app.ActivityOptions;
 import android.app.PendingIntent;
 import android.app.admin.DevicePolicyManager;
-import android.app.compat.CompatChanges;
 import android.companion.AssociationInfo;
 import android.companion.virtual.IVirtualDevice;
 import android.companion.virtual.IVirtualDeviceActivityListener;
@@ -55,8 +54,6 @@
 import android.companion.virtual.flags.Flags;
 import android.companion.virtual.sensor.VirtualSensor;
 import android.companion.virtual.sensor.VirtualSensorEvent;
-import android.compat.annotation.ChangeId;
-import android.compat.annotation.EnabledAfter;
 import android.content.AttributionSource;
 import android.content.ComponentName;
 import android.content.Context;
@@ -81,7 +78,6 @@
 import android.hardware.input.VirtualTouchEvent;
 import android.hardware.input.VirtualTouchscreenConfig;
 import android.os.Binder;
-import android.os.Build;
 import android.os.IBinder;
 import android.os.LocaleList;
 import android.os.Looper;
@@ -122,22 +118,6 @@
 
     private static final String TAG = "VirtualDeviceImpl";
 
-    /**
-     * Virtual displays created by a {@code VirtualDeviceManager.VirtualDevice} are more consistent
-     * with virtual displays created via {@link android.hardware.display.DisplayManager} and allow
-     * for the creation of private, auto-mirror, and fixed orientation displays since
-     * {@link android.os.Build.VERSION_CODES#VANILLA_ICE_CREAM}.
-     *
-     * @see DisplayManager#VIRTUAL_DISPLAY_FLAG_PUBLIC
-     * @see DisplayManager#VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY
-     * @see DisplayManager#VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR
-     * @see DisplayManager#VIRTUAL_DISPLAY_FLAG_ROTATES_WITH_CONTENT
-     */
-    @ChangeId
-    @EnabledAfter(targetSdkVersion = Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
-    public static final long MAKE_VIRTUAL_DISPLAY_FLAGS_CONSISTENT_WITH_DISPLAY_MANAGER =
-            294837146L;
-
     private static final int DEFAULT_VIRTUAL_DISPLAY_FLAGS =
             DisplayManager.VIRTUAL_DISPLAY_FLAG_TOUCH_FEEDBACK_DISABLED
                     | DisplayManager.VIRTUAL_DISPLAY_FLAG_DESTROY_CONTENT_ON_REMOVAL
@@ -365,8 +345,7 @@
         }
 
         int flags = DEFAULT_VIRTUAL_DISPLAY_FLAGS;
-        if (!CompatChanges.isChangeEnabled(
-                MAKE_VIRTUAL_DISPLAY_FLAGS_CONSISTENT_WITH_DISPLAY_MANAGER, mOwnerUid)) {
+        if (!Flags.consistentDisplayFlags()) {
             flags |= DEFAULT_VIRTUAL_DISPLAY_FLAGS_PRE_VIC;
         }
         if (mParams.getLockState() == VirtualDeviceParams.LOCK_STATE_ALWAYS_UNLOCKED) {
diff --git a/services/core/java/com/android/server/BinaryTransparencyService.java b/services/core/java/com/android/server/BinaryTransparencyService.java
index 77b6d583..eb3ec24 100644
--- a/services/core/java/com/android/server/BinaryTransparencyService.java
+++ b/services/core/java/com/android/server/BinaryTransparencyService.java
@@ -1466,8 +1466,11 @@
         if (android.security.Flags.binaryTransparencySepolicyHash()) {
             byte[] sepolicyHash = PackageUtils.computeSha256DigestForLargeFileAsBytes(
                     "/sys/fs/selinux/policy", PackageUtils.createLargeFileBuffer());
-            String sepolicyHashEncoded = HexEncoding.encodeToString(sepolicyHash, false);
-            Slog.d(TAG, "sepolicy hash: " + sepolicyHashEncoded);
+            String sepolicyHashEncoded = null;
+            if (sepolicyHash != null) {
+                sepolicyHashEncoded = HexEncoding.encodeToString(sepolicyHash, false);
+                Slog.d(TAG, "sepolicy hash: " + sepolicyHashEncoded);
+            }
             FrameworkStatsLog.write(FrameworkStatsLog.BOOT_INTEGRITY_INFO_REPORTED,
                     sepolicyHashEncoded, mVbmetaDigest);
         }
diff --git a/services/core/java/com/android/server/OWNERS b/services/core/java/com/android/server/OWNERS
index 987507f..7fe0682 100644
--- a/services/core/java/com/android/server/OWNERS
+++ b/services/core/java/com/android/server/OWNERS
@@ -35,7 +35,7 @@
 per-file MmsServiceBroker.java = file:/telephony/OWNERS
 per-file NetIdManager.java = file:/services/core/java/com/android/server/net/OWNERS
 per-file PackageWatchdog.java, RescueParty.java = file:/services/core/java/com/android/server/rollback/OWNERS
-per-file PinnerService.java = file:/apct-tests/perftests/OWNERS
+per-file PinnerService.java = file:/core/java/android/app/pinner/OWNERS
 per-file RescueParty.java = shuc@google.com, ancr@google.com, harshitmahajan@google.com
 per-file SystemClockTime.java = file:/services/core/java/com/android/server/timedetector/OWNERS
 per-file SystemTimeZone.java = file:/services/core/java/com/android/server/timezonedetector/OWNERS
diff --git a/services/core/java/com/android/server/PinnerService.java b/services/core/java/com/android/server/PinnerService.java
index 0e7b4aa..23a30f9 100644
--- a/services/core/java/com/android/server/PinnerService.java
+++ b/services/core/java/com/android/server/PinnerService.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2016 The Android Open Source Project
+ * Copyright (C) 2023 The Android Open Source Project
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -20,6 +20,9 @@
 import static android.app.ActivityManager.UID_OBSERVER_GONE;
 import static android.os.Process.SYSTEM_UID;
 
+import static com.android.server.flags.Flags.pinWebview;
+
+import android.annotation.EnforcePermission;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -28,6 +31,8 @@
 import android.app.IActivityManager;
 import android.app.SearchManager;
 import android.app.UidObserver;
+import android.app.pinner.IPinnerService;
+import android.app.pinner.PinnedFileStat;
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
@@ -48,6 +53,7 @@
 import android.os.ResultReceiver;
 import android.os.ShellCallback;
 import android.os.SystemProperties;
+import android.os.Trace;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.provider.DeviceConfig;
@@ -83,6 +89,8 @@
 import java.lang.annotation.RetentionPolicy;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipFile;
@@ -111,16 +119,15 @@
     private static final int KEY_ASSISTANT = 2;
 
     // Pin using pinlist.meta when pinning apps.
-    private static boolean PROP_PIN_PINLIST = SystemProperties.getBoolean(
-            "pinner.use_pinlist", true);
-    // Pin the whole odex/vdex/etc file when pinning apps.
-    private static boolean PROP_PIN_ODEX = SystemProperties.getBoolean(
-            "pinner.whole_odex", true);
+    private static boolean PROP_PIN_PINLIST =
+            SystemProperties.getBoolean("pinner.use_pinlist", true);
 
     private static final int MAX_CAMERA_PIN_SIZE = 80 * (1 << 20); // 80MB max for camera app.
     private static final int MAX_HOME_PIN_SIZE = 6 * (1 << 20); // 6MB max for home app.
     private static final int MAX_ASSISTANT_PIN_SIZE = 60 * (1 << 20); // 60MB max for assistant app.
 
+    public static final String ANON_REGION_STAT_NAME = "[anon]";
+
     @IntDef({KEY_CAMERA, KEY_HOME, KEY_ASSISTANT})
     @Retention(RetentionPolicy.SOURCE)
     public @interface AppKey {}
@@ -135,8 +142,7 @@
     private SearchManager mSearchManager;
 
     /** The list of the statically pinned files. */
-    @GuardedBy("this")
-    private final ArrayList<PinnedFile> mPinnedFiles = new ArrayList<>();
+    @GuardedBy("this") private final ArrayMap<String, PinnedFile> mPinnedFiles = new ArrayMap<>();
 
     /** The list of the pinned apps. This is a map from {@link AppKey} to a pinned app. */
     @GuardedBy("this")
@@ -175,6 +181,7 @@
     private final boolean mConfiguredToPinCamera;
     private final boolean mConfiguredToPinHome;
     private final boolean mConfiguredToPinAssistant;
+    private final int mConfiguredWebviewPinBytes;
 
     private BinderService mBinderService;
     private PinnerHandler mPinnerHandler = null;
@@ -214,6 +221,11 @@
         protected void publishBinderService(PinnerService service, Binder binderService) {
             service.publishBinderService("pinner", binderService);
         }
+
+        protected PinnedFile pinFileInternal(String fileToPin,
+                int maxBytesToPin, boolean attemptPinIntrospection) {
+            return PinnerService.pinFileInternal(fileToPin, maxBytesToPin, attemptPinIntrospection);
+        }
     }
 
     public PinnerService(Context context) {
@@ -233,6 +245,8 @@
                 com.android.internal.R.bool.config_pinnerHomeApp);
         mConfiguredToPinAssistant = context.getResources().getBoolean(
                 com.android.internal.R.bool.config_pinnerAssistantApp);
+        mConfiguredWebviewPinBytes = context.getResources().getInteger(
+                com.android.internal.R.integer.config_pinnerWebviewPinBytes);
         mPinKeys = createPinKeys();
         mPinnerHandler = new PinnerHandler(BackgroundThread.get().getLooper());
 
@@ -322,7 +336,7 @@
     public List<PinnedFileStats> dumpDataForStatsd() {
         List<PinnedFileStats> pinnedFileStats = new ArrayList<>();
         synchronized (PinnerService.this) {
-            for (PinnedFile pinnedFile : mPinnedFiles) {
+            for (PinnedFile pinnedFile : mPinnedFiles.values()) {
                 pinnedFileStats.add(new PinnedFileStats(SYSTEM_UID, pinnedFile));
             }
 
@@ -358,39 +372,17 @@
             com.android.internal.R.array.config_defaultPinnerServiceFiles);
         // Continue trying to pin each file even if we fail to pin some of them
         for (String fileToPin : filesToPin) {
-            PinnedFile pf = pinFile(fileToPin,
-                                    Integer.MAX_VALUE,
-                                    /*attemptPinIntrospection=*/false);
+            PinnedFile pf = mInjector.pinFileInternal(fileToPin, Integer.MAX_VALUE,
+                    /*attemptPinIntrospection=*/false);
             if (pf == null) {
                 Slog.e(TAG, "Failed to pin file = " + fileToPin);
                 continue;
             }
             synchronized (this) {
-                mPinnedFiles.add(pf);
+                mPinnedFiles.put(pf.fileName, pf);
             }
-            if (fileToPin.endsWith(".jar") | fileToPin.endsWith(".apk")) {
-                // Check whether the runtime has compilation artifacts to pin.
-                String arch = VMRuntime.getInstructionSet(Build.SUPPORTED_ABIS[0]);
-                String[] files = null;
-                try {
-                    files = DexFile.getDexFileOutputPaths(fileToPin, arch);
-                } catch (IOException ioe) { }
-                if (files == null) {
-                    continue;
-                }
-                for (String file : files) {
-                    PinnedFile df = pinFile(file,
-                                            Integer.MAX_VALUE,
-                                            /*attemptPinIntrospection=*/false);
-                    if (df == null) {
-                        Slog.i(TAG, "Failed to pin ART file = " + file);
-                        continue;
-                    }
-                    synchronized (this) {
-                        mPinnedFiles.add(df);
-                    }
-                }
-            }
+            pf.groupName = "system";
+            pinOptimizedDexDependencies(pf, Integer.MAX_VALUE, null);
         }
 
         refreshPinAnonConfig();
@@ -487,7 +479,7 @@
             pinnedAppFiles = new ArrayList<>(app.mFiles);
         }
         for (PinnedFile pinnedFile : pinnedAppFiles) {
-            pinnedFile.close();
+            unpinFile(pinnedFile.fileName);
         }
     }
 
@@ -495,6 +487,19 @@
         return ResolverActivity.class.getName().equals(info.name);
     }
 
+    public int getWebviewPinQuota() {
+        if (!pinWebview()) {
+            return 0;
+        }
+        int quota = mConfiguredWebviewPinBytes;
+        int overrideQuota = SystemProperties.getInt("pinner.pin_webview_size", -1);
+        if (overrideQuota != -1) {
+            // Quota was overridden
+            quota = overrideQuota;
+        }
+        return quota;
+    }
+
     private ApplicationInfo getCameraInfo(int userHandle) {
         Intent cameraIntent = new Intent(MediaStore.INTENT_ACTION_STILL_IMAGE_CAMERA);
         ApplicationInfo info = getApplicationInfoForIntent(cameraIntent, userHandle,
@@ -728,7 +733,7 @@
             case KEY_ASSISTANT:
                 return "Assistant";
             default:
-                return null;
+                return "";
         }
     }
 
@@ -868,11 +873,12 @@
                 continue;
             }
 
-            PinnedFile pf = pinFile(apk, apkPinSizeLimit, /*attemptPinIntrospection=*/true);
+            PinnedFile pf = mInjector.pinFileInternal(apk, apkPinSizeLimit, /*attemptPinIntrospection=*/true);
             if (pf == null) {
                 Slog.e(TAG, "Failed to pin " + apk);
                 continue;
             }
+            pf.groupName = getNameForKey(key);
 
             if (DEBUG) {
                 Slog.i(TAG, "Pinned " + pf.fileName);
@@ -882,40 +888,118 @@
             }
 
             apkPinSizeLimit -= pf.bytesPinned;
+            if (apk.equals(appInfo.sourceDir)) {
+                pinOptimizedDexDependencies(pf, apkPinSizeLimit, appInfo);
+            }
         }
+    }
 
-        // determine the ABI from either ApplicationInfo or Build
-        String abi = appInfo.primaryCpuAbi != null ? appInfo.primaryCpuAbi :
-                Build.SUPPORTED_ABIS[0];
-        String arch = VMRuntime.getInstructionSet(abi);
-        // get the path to the odex or oat file
-        String baseCodePath = appInfo.getBaseCodePath();
-        String[] files = null;
-        try {
-            files = DexFile.getDexFileOutputPaths(baseCodePath, arch);
-        } catch (IOException ioe) {}
-        if (files == null) {
-            return;
+    /**
+     * Pin file or apk to memory.
+     *
+     * Prefer to use this method instead of {@link #pinFileInternal(String, int, boolean)} as it
+     * takes care of accounting and if pinning an apk, it also pins any extra optimized art files
+     * that related to the file but not within itself.
+     *
+     * @param fileToPin File to pin
+     * @param maxBytesToPin maximum quota allowed for pinning
+     * @return total bytes that were pinned.
+     */
+    public int pinFile(String fileToPin, int maxBytesToPin, @Nullable ApplicationInfo appInfo,
+            @Nullable String groupName) {
+        PinnedFile existingPin;
+        synchronized(this) {
+            existingPin = mPinnedFiles.get(fileToPin);
         }
-
-        //not pinning the oat/odex is not a fatal error
-        for (String file : files) {
-            PinnedFile pf = pinFile(file, pinSizeLimit, /*attemptPinIntrospection=*/false);
-            if (pf != null) {
-                synchronized (this) {
-                    if (PROP_PIN_ODEX) {
-                      pinnedApp.mFiles.add(pf);
-                    }
-                }
+        if (existingPin != null) {
+            if (existingPin.bytesPinned == maxBytesToPin) {
+                // Duplicate pin requesting same amount of bytes, lets just bail out.
+                return 0;
+            } else {
+                // User decided to pin a different amount of bytes than currently pinned
+                // so this is a valid pin request. Unpin the previous version before repining.
                 if (DEBUG) {
-                    if (PROP_PIN_ODEX) {
-                        Slog.i(TAG, "Pinned " + pf.fileName);
-                    } else {
-                        Slog.i(TAG, "Pinned [skip] " + pf.fileName);
-                    }
+                    Slog.d(TAG, "Unpinning file prior to repin: " + fileToPin);
+                }
+                unpinFile(fileToPin);
+            }
+        }
+
+        boolean isApk = fileToPin.endsWith(".apk");
+        int bytesPinned = 0;
+        PinnedFile pf = mInjector.pinFileInternal(fileToPin, maxBytesToPin,
+                /*attemptPinIntrospection=*/isApk);
+        if (pf == null) {
+            Slog.e(TAG, "Failed to pin file = " + fileToPin);
+            return 0;
+        }
+        pf.groupName = groupName != null ? groupName : "";
+
+        maxBytesToPin -= bytesPinned;
+        bytesPinned += pf.bytesPinned;
+
+        synchronized (this) {
+            mPinnedFiles.put(pf.fileName, pf);
+        }
+        if (maxBytesToPin > 0) {
+            pinOptimizedDexDependencies(pf, maxBytesToPin, appInfo);
+        }
+        return bytesPinned;
+    }
+
+    /**
+     * Pin any dependency optimized files generated by ART.
+     * @param pinnedFile An already pinned file whose dependencies we want pinned.
+     * @param maxBytesToPin Maximum amount of bytes to pin.
+     * @param appInfo Used to determine the ABI in case the application has one custom set, when set
+     *                to null it will use the default supported ABI by the device.
+     * @return total bytes pinned.
+     */
+    private int pinOptimizedDexDependencies(
+            PinnedFile pinnedFile, int maxBytesToPin, @Nullable ApplicationInfo appInfo) {
+        if (pinnedFile == null) {
+            return 0;
+        }
+
+        int bytesPinned = 0;
+        if (pinnedFile.fileName.endsWith(".jar") | pinnedFile.fileName.endsWith(".apk")) {
+            String abi = null;
+            if (appInfo != null) {
+                abi = appInfo.primaryCpuAbi;
+            }
+            if (abi == null) {
+                abi = Build.SUPPORTED_ABIS[0];
+            }
+            // Check whether the runtime has compilation artifacts to pin.
+            String arch = VMRuntime.getInstructionSet(abi);
+            String[] files = null;
+            try {
+                files = DexFile.getDexFileOutputPaths(pinnedFile.fileName, arch);
+            } catch (IOException ioe) {
+            }
+            if (files == null) {
+                return bytesPinned;
+            }
+            for (String file : files) {
+                // Unpin if it was already pinned prior to re-pinning.
+                unpinFile(file);
+
+                PinnedFile df = mInjector.pinFileInternal(file, Integer.MAX_VALUE,
+                        /*attemptPinIntrospection=*/false);
+                if (df == null) {
+                    Slog.i(TAG, "Failed to pin ART file = " + file);
+                    return bytesPinned;
+                }
+                df.groupName = pinnedFile.groupName;
+                pinnedFile.pinnedDeps.add(df);
+                maxBytesToPin -= df.bytesPinned;
+                bytesPinned += df.bytesPinned;
+                synchronized (this) {
+                    mPinnedFiles.put(df.fileName, df);
                 }
             }
         }
+        return bytesPinned;
     }
 
     /** mlock length bytes of fileToPin in memory
@@ -955,9 +1039,12 @@
      *   zip in order to extract the
      * @return Pinned memory resource owner thing or null on error
      */
-    private static PinnedFile pinFile(String fileToPin,
-                                      int maxBytesToPin,
-                                      boolean attemptPinIntrospection) {
+    private static PinnedFile pinFileInternal(
+            String fileToPin, int maxBytesToPin, boolean attemptPinIntrospection) {
+        if (DEBUG) {
+            Slog.d(TAG, "pin file: " + fileToPin + " use-pinlist: " + attemptPinIntrospection);
+        }
+        Trace.beginSection("pinFile:" + fileToPin);
         ZipFile fileAsZip = null;
         InputStream pinRangeStream = null;
         try {
@@ -968,16 +1055,19 @@
             if (fileAsZip != null) {
                 pinRangeStream = maybeOpenPinMetaInZip(fileAsZip, fileToPin);
             }
-
-            Slog.d(TAG, "pinRangeStream: " + pinRangeStream);
-
-            PinRangeSource pinRangeSource = (pinRangeStream != null)
-                ? new PinRangeSourceStream(pinRangeStream)
-                : new PinRangeSourceStatic(0, Integer.MAX_VALUE /* will be clipped */);
-            return pinFileRanges(fileToPin, maxBytesToPin, pinRangeSource);
+            boolean use_pinlist = (pinRangeStream != null);
+            PinRangeSource pinRangeSource = use_pinlist
+                    ? new PinRangeSourceStream(pinRangeStream)
+                    : new PinRangeSourceStatic(0, Integer.MAX_VALUE /* will be clipped */);
+            PinnedFile pinnedFile = pinFileRanges(fileToPin, maxBytesToPin, pinRangeSource);
+            if (pinnedFile != null) {
+                pinnedFile.used_pinlist = use_pinlist;
+            }
+            return pinnedFile;
         } finally {
             safeClose(pinRangeStream);
             safeClose(fileAsZip);  // Also closes any streams we've opened
+            Trace.endSection();
         }
     }
 
@@ -1013,9 +1103,23 @@
             return null;
         }
 
+        // Looking at root directory is the old behavior but still some apps rely on it so keeping
+        // for backward compatibility. As doing a single item lookup is cheap in the root.
         ZipEntry pinMetaEntry = zipFile.getEntry(PIN_META_FILENAME);
+
+        if (pinMetaEntry == null) {
+            // It is usually within an apk's control to include files in assets/ directory
+            // so this would be the expected point to have the pinlist.meta coming from.
+            // we explicitly avoid doing an exhaustive search because it may be expensive so
+            // prefer to have a good known location to retrieve the file.
+            pinMetaEntry = zipFile.getEntry("assets/" + PIN_META_FILENAME);
+        }
+
         InputStream pinMetaStream = null;
         if (pinMetaEntry != null) {
+            if (DEBUG) {
+                Slog.d(TAG, "Found pinlist.meta for " + fileName);
+            }
             try {
                 pinMetaStream = zipFile.getInputStream(pinMetaEntry);
             } catch (IOException ex) {
@@ -1024,6 +1128,10 @@
                                      fileName),
                        ex);
             }
+        } else {
+            Slog.w(TAG,
+                    String.format(
+                            "Could not find pinlist.meta for \"%s\": pinning as blob", fileName));
         }
         return pinMetaStream;
     }
@@ -1160,6 +1268,49 @@
             }
         }
     }
+    private List<PinnedFile> getAllPinsForGroup(String group) {
+        List<PinnedFile> filesInGroup;
+        synchronized (this) {
+            filesInGroup = mPinnedFiles.values()
+                                   .stream()
+                                   .filter(pf -> pf.groupName.equals(group))
+                                   .toList();
+        }
+        return filesInGroup;
+    }
+    public void unpinGroup(String group) {
+        List<PinnedFile> pinnedFiles = getAllPinsForGroup(group);
+        for (PinnedFile pf : pinnedFiles) {
+            unpinFile(pf.fileName);
+        }
+    }
+
+    public void unpinFile(String filename) {
+        PinnedFile pinnedFile;
+        synchronized (this) {
+            pinnedFile = mPinnedFiles.get(filename);
+        }
+        if (pinnedFile == null) {
+            // File not pinned, nothing to do.
+            return;
+        }
+        pinnedFile.close();
+        synchronized (this) {
+            if (DEBUG) {
+                Slog.d(TAG, "Unpinned file: " + filename);
+            }
+            mPinnedFiles.remove(pinnedFile.fileName);
+            for (PinnedFile dep : pinnedFile.pinnedDeps) {
+                if (dep == null) {
+                    continue;
+                }
+                mPinnedFiles.remove(dep.fileName);
+                if (DEBUG) {
+                    Slog.d(TAG, "Unpinned dependency: " + dep.fileName);
+                }
+            }
+        }
+    }
 
     private static int clamp(int min, int value, int max) {
         return Math.max(min, Math.min(value, max));
@@ -1205,17 +1356,44 @@
         }
     }
 
-    private final class BinderService extends Binder {
+    public List<PinnedFileStat> getPinnerStats() {
+        ArrayList<PinnedFileStat> stats = new ArrayList<>();
+        Collection<PinnedApp> pinnedApps;
+        synchronized(this) {
+            pinnedApps = mPinnedApps.values();
+        }
+        for (PinnedApp pinnedApp : pinnedApps) {
+            for (PinnedFile pf : pinnedApp.mFiles) {
+                PinnedFileStat stat =
+                        new PinnedFileStat(pf.fileName, pf.bytesPinned, pf.groupName);
+                stats.add(stat);
+            }
+        }
+
+        Collection<PinnedFile> pinnedFiles;
+        synchronized(this) {
+            pinnedFiles = mPinnedFiles.values();
+        }
+        for (PinnedFile pf : pinnedFiles) {
+            PinnedFileStat stat = new PinnedFileStat(pf.fileName, pf.bytesPinned, pf.groupName);
+            stats.add(stat);
+        }
+        if (mCurrentlyPinnedAnonSize > 0) {
+            stats.add(new PinnedFileStat(ANON_REGION_STAT_NAME,
+                        mCurrentlyPinnedAnonSize, ANON_REGION_STAT_NAME));
+        }
+        return stats;
+    }
+
+    public final class BinderService extends IPinnerService.Stub {
         @Override
         protected void dump(FileDescriptor fd, PrintWriter pw, String[] args) {
             if (!DumpUtils.checkDumpPermission(mContext, TAG, pw)) return;
+            HashSet<PinnedFile> shownPins = new HashSet<>();
+            HashSet<String> groups = new HashSet<>();
+            final int bytesPerMB = 1024 * 1024;
             synchronized (PinnerService.this) {
                 long totalSize = 0;
-                for (PinnedFile pinnedFile : mPinnedFiles) {
-                    pw.format("%s %s\n", pinnedFile.fileName, pinnedFile.bytesPinned);
-                    totalSize += pinnedFile.bytesPinned;
-                }
-                pw.println();
                 for (int key : mPinnedApps.keySet()) {
                     PinnedApp app = mPinnedApps.get(key);
                     pw.print(getNameForKey(key));
@@ -1223,14 +1401,53 @@
                     pw.print(" active="); pw.print(app.active);
                     pw.println();
                     for (PinnedFile pf : mPinnedApps.get(key).mFiles) {
-                        pw.print("  "); pw.format("%s %s\n", pf.fileName, pf.bytesPinned);
+                        pw.print("  ");
+                        pw.format("%s pinned:%d bytes (%d MB) pinlist:%b\n", pf.fileName,
+                                pf.bytesPinned, pf.bytesPinned / bytesPerMB, pf.used_pinlist);
                         totalSize += pf.bytesPinned;
+                        shownPins.add(pf);
+                        for (PinnedFile dep : pf.pinnedDeps) {
+                            pw.print("  ");
+                            pw.format("%s pinned:%d bytes (%d MB) pinlist:%b (Dependency)\n", dep.fileName,
+                                    dep.bytesPinned, dep.bytesPinned / bytesPerMB, dep.used_pinlist);
+                            totalSize += dep.bytesPinned;
+                            shownPins.add(dep);
+                        }
                     }
                 }
-                if (mPinAnonAddress != 0) {
-                    pw.format("Pinned anon region: %s\n", mCurrentlyPinnedAnonSize);
+                pw.println();
+                for (PinnedFile pinnedFile : mPinnedFiles.values()) {
+                    if (!groups.contains(pinnedFile.groupName)) {
+                        groups.add(pinnedFile.groupName);
+                    }
                 }
-                pw.format("Total size: %s\n", totalSize);
+                boolean firstPinInGroup = true;
+                for (String group : groups) {
+                    List<PinnedFile> groupPins = getAllPinsForGroup(group);
+                    for (PinnedFile pinnedFile : groupPins) {
+                        if (shownPins.contains(pinnedFile)) {
+                            // Already showed in the dump and accounted for, skip.
+                            continue;
+                        }
+                        if (firstPinInGroup) {
+                            firstPinInGroup = false;
+                            // Ensure we only print when there are pins for groups not yet shown
+                            // in the pinned app section.
+                            pw.print("Group:" + group);
+                            pw.println();
+                        }
+                        pw.format("  %s pinned:%d bytes (%d MB) pinlist:%b\n", pinnedFile.fileName,
+                                pinnedFile.bytesPinned, pinnedFile.bytesPinned / bytesPerMB,
+                                pinnedFile.used_pinlist);
+                        totalSize += pinnedFile.bytesPinned;
+                    }
+                }
+                pw.println();
+                if (mPinAnonAddress != 0) {
+                    pw.format("Pinned anon region: %d (%d MB)\n", mCurrentlyPinnedAnonSize, mCurrentlyPinnedAnonSize / bytesPerMB);
+                    totalSize += mCurrentlyPinnedAnonSize;
+                }
+                pw.format("Total pinned: %s bytes (%s MB)\n", totalSize, totalSize / bytesPerMB);
                 pw.println();
                 if (!mPendingRepin.isEmpty()) {
                     pw.print("Pending repin: ");
@@ -1277,14 +1494,29 @@
 
             resultReceiver.send(0, null);
         }
+
+        @EnforcePermission(android.Manifest.permission.DUMP)
+        @Override
+        public List<PinnedFileStat> getPinnerStats() {
+            getPinnerStats_enforcePermission();
+            return PinnerService.this.getPinnerStats();
+        }
     }
 
-    private static final class PinnedFile implements AutoCloseable {
+    @VisibleForTesting
+    public static final class PinnedFile implements AutoCloseable {
         private long mAddress;
         final int mapSize;
         final String fileName;
         final int bytesPinned;
 
+        // Whether this file was pinned using a pinlist
+        boolean used_pinlist;
+
+        // User defined group name for pinner accounting
+        String groupName = "";
+        ArrayList<PinnedFile> pinnedDeps = new ArrayList<>();
+
         PinnedFile(long address, int mapSize, String fileName, int bytesPinned) {
              mAddress = address;
              this.mapSize = mapSize;
@@ -1298,6 +1530,11 @@
                 safeMunmap(mAddress, mapSize);
                 mAddress = -1;
             }
+            for (PinnedFile dep : pinnedDeps) {
+                if (dep != null) {
+                    dep.close();
+                }
+            }
         }
 
         @Override
@@ -1355,5 +1592,4 @@
             }
         }
     }
-
 }
diff --git a/services/core/java/com/android/server/StorageManagerService.java b/services/core/java/com/android/server/StorageManagerService.java
index f6835fe..39b8643 100644
--- a/services/core/java/com/android/server/StorageManagerService.java
+++ b/services/core/java/com/android/server/StorageManagerService.java
@@ -215,7 +215,7 @@
     public static final int FAILED_MOUNT_RESET_TIMEOUT_SECONDS = 10;
 
     /** Extended timeout for the system server watchdog. */
-    private static final int SLOW_OPERATION_WATCHDOG_TIMEOUT_MS = 60 * 1000;
+    private static final int SLOW_OPERATION_WATCHDOG_TIMEOUT_MS = 20 * 1000;
 
     /** Extended timeout for the system server watchdog for vold#partition operation. */
     private static final int PARTITION_OPERATION_WATCHDOG_TIMEOUT_MS = 3 * 60 * 1000;
@@ -1235,11 +1235,16 @@
         }
     }
 
+    private void extendWatchdogTimeout(String reason) {
+        Watchdog w = Watchdog.getInstance();
+        w.pauseWatchingMonitorsFor(SLOW_OPERATION_WATCHDOG_TIMEOUT_MS, reason);
+        w.pauseWatchingCurrentThreadFor(SLOW_OPERATION_WATCHDOG_TIMEOUT_MS, reason);
+    }
+
     private void onUserStopped(int userId) {
         Slog.d(TAG, "onUserStopped " + userId);
 
-        Watchdog.getInstance().pauseWatchingMonitorsFor(
-                SLOW_OPERATION_WATCHDOG_TIMEOUT_MS, "#onUserStopped might be slow");
+        extendWatchdogTimeout("#onUserStopped might be slow");
         try {
             mVold.onUserStopped(userId);
             mStoraged.onUserStopped(userId);
@@ -1322,8 +1327,7 @@
                 unlockedUsers.add(userId);
             }
         }
-        Watchdog.getInstance().pauseWatchingMonitorsFor(
-                SLOW_OPERATION_WATCHDOG_TIMEOUT_MS, "#onUserStopped might be slow");
+        extendWatchdogTimeout("#onUserStopped might be slow");
         for (Integer userId : unlockedUsers) {
             try {
                 mVold.onUserStopped(userId);
@@ -2343,8 +2347,7 @@
         try {
             // TODO(b/135341433): Remove cautious logging when FUSE is stable
             Slog.i(TAG, "Mounting volume " + vol);
-            Watchdog.getInstance().pauseWatchingMonitorsFor(
-                    SLOW_OPERATION_WATCHDOG_TIMEOUT_MS, "#mount might be slow");
+            extendWatchdogTimeout("#mount might be slow");
             mVold.mount(vol.id, vol.mountFlags, vol.mountUserId, new IVoldMountCallback.Stub() {
                 @Override
                 public boolean onVolumeChecking(FileDescriptor fd, String path,
@@ -2474,8 +2477,7 @@
 
         final CountDownLatch latch = findOrCreateDiskScanLatch(diskId);
 
-        Watchdog.getInstance().pauseWatchingMonitorsFor(
-                PARTITION_OPERATION_WATCHDOG_TIMEOUT_MS, "#partition might be very slow");
+        extendWatchdogTimeout("#partition might be slow");
         try {
             mVold.partition(diskId, IVold.PARTITION_TYPE_PUBLIC, -1);
             waitForLatch(latch, "partitionPublic", 3 * DateUtils.MINUTE_IN_MILLIS);
@@ -2493,8 +2495,7 @@
 
         final CountDownLatch latch = findOrCreateDiskScanLatch(diskId);
 
-        Watchdog.getInstance().pauseWatchingMonitorsFor(
-                PARTITION_OPERATION_WATCHDOG_TIMEOUT_MS, "#partition might be very slow");
+        extendWatchdogTimeout("#partition might be slow");
         try {
             mVold.partition(diskId, IVold.PARTITION_TYPE_PRIVATE, -1);
             waitForLatch(latch, "partitionPrivate", 3 * DateUtils.MINUTE_IN_MILLIS);
@@ -2512,8 +2513,7 @@
 
         final CountDownLatch latch = findOrCreateDiskScanLatch(diskId);
 
-        Watchdog.getInstance().pauseWatchingMonitorsFor(
-                PARTITION_OPERATION_WATCHDOG_TIMEOUT_MS, "#partition might be very slow");
+        extendWatchdogTimeout("#partition might be slow");
         try {
             mVold.partition(diskId, IVold.PARTITION_TYPE_MIXED, ratio);
             waitForLatch(latch, "partitionMixed", 3 * DateUtils.MINUTE_IN_MILLIS);
@@ -3622,8 +3622,7 @@
 
         @Override
         public ParcelFileDescriptor open() throws AppFuseMountException {
-            Watchdog.getInstance().pauseWatchingMonitorsFor(
-                SLOW_OPERATION_WATCHDOG_TIMEOUT_MS, "#open might be slow");
+            extendWatchdogTimeout("#open might be slow");
             try {
                 final FileDescriptor fd = mVold.mountAppFuse(uid, mountId);
                 mMounted = true;
@@ -3636,8 +3635,7 @@
         @Override
         public ParcelFileDescriptor openFile(int mountId, int fileId, int flags)
                 throws AppFuseMountException {
-            Watchdog.getInstance().pauseWatchingMonitorsFor(
-                SLOW_OPERATION_WATCHDOG_TIMEOUT_MS, "#openFile might be slow");
+            extendWatchdogTimeout("#openFile might be slow");
             try {
                 return new ParcelFileDescriptor(
                         mVold.openAppFuseFile(uid, mountId, fileId, flags));
@@ -3648,8 +3646,7 @@
 
         @Override
         public void close() throws Exception {
-            Watchdog.getInstance().pauseWatchingMonitorsFor(
-                SLOW_OPERATION_WATCHDOG_TIMEOUT_MS, "#close might be slow");
+            extendWatchdogTimeout("#close might be slow");
             if (mMounted) {
                 BackgroundThread.getHandler().post(() -> {
                     try {
diff --git a/services/core/java/com/android/server/TelephonyRegistry.java b/services/core/java/com/android/server/TelephonyRegistry.java
index c718d39..9eb35fd 100644
--- a/services/core/java/com/android/server/TelephonyRegistry.java
+++ b/services/core/java/com/android/server/TelephonyRegistry.java
@@ -2096,14 +2096,48 @@
     /**
      * Send a notification to registrants about the data activity state.
      *
+     * @param subId the subscriptionId for the data connection
+     * @param state indicates the latest data activity type
+     * e.g.,{@link TelephonyManager#DATA_ACTIVITY_IN}
+     *
+     */
+
+    public void notifyDataActivityForSubscriber(int subId, int state) {
+        if (!checkNotifyPermission("notifyDataActivity()")) {
+            return;
+        }
+        int phoneId = getPhoneIdFromSubId(subId);
+        synchronized (mRecords) {
+            if (validatePhoneId(phoneId)) {
+                mDataActivity[phoneId] = state;
+                for (Record r : mRecords) {
+                    // Notify by correct subId.
+                    if (r.matchTelephonyCallbackEvent(
+                            TelephonyCallback.EVENT_DATA_ACTIVITY_CHANGED)
+                            && idMatch(r, subId, phoneId)) {
+                        try {
+                            r.callback.onDataActivity(state);
+                        } catch (RemoteException ex) {
+                            mRemoveList.add(r.binder);
+                        }
+                    }
+                }
+            }
+            handleRemoveListLocked();
+        }
+    }
+
+    /**
+     * Send a notification to registrants about the data activity state.
+     *
      * @param phoneId the phoneId carrying the data connection
      * @param subId the subscriptionId for the data connection
      * @param state indicates the latest data activity type
      * e.g.,{@link TelephonyManager#DATA_ACTIVITY_IN}
      *
      */
-    public void notifyDataActivityForSubscriber(int phoneId, int subId, int state) {
-        if (!checkNotifyPermission("notifyDataActivity()" )) {
+    public void notifyDataActivityForSubscriberWithSlot(int phoneId, int subId, int state) {
+        if (!checkNotifyPermission("notifyDataActivityWithSlot()")) {
             return;
         }
 
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 36356bd..e656030 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -120,12 +120,15 @@
 import com.android.server.power.optimization.Flags;
 import com.android.server.power.stats.AggregatedPowerStatsConfig;
 import com.android.server.power.stats.BatteryExternalStatsWorker;
+import com.android.server.power.stats.BatteryStatsDumpHelperImpl;
 import com.android.server.power.stats.BatteryStatsImpl;
 import com.android.server.power.stats.BatteryUsageStatsProvider;
 import com.android.server.power.stats.CpuAggregatedPowerStatsProcessor;
 import com.android.server.power.stats.PowerStatsAggregator;
+import com.android.server.power.stats.PowerStatsExporter;
 import com.android.server.power.stats.PowerStatsScheduler;
 import com.android.server.power.stats.PowerStatsStore;
+import com.android.server.power.stats.PowerStatsUidResolver;
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
 import com.android.server.power.stats.wakeups.CpuWakeupStats;
 
@@ -181,6 +184,8 @@
     private final BatteryExternalStatsWorker mWorker;
     private final BatteryUsageStatsProvider mBatteryUsageStatsProvider;
     private final AtomicFile mConfigFile;
+    private final BatteryStats.BatteryStatsDumpHelper mDumpHelper;
+    private final PowerStatsUidResolver mPowerStatsUidResolver;
 
     private volatile boolean mMonitorEnabled = true;
 
@@ -408,9 +413,10 @@
                         .setResetOnUnplugAfterSignificantCharge(resetOnUnplugAfterSignificantCharge)
                         .setPowerStatsThrottlePeriodCpu(powerStatsThrottlePeriodCpu)
                         .build();
+        mPowerStatsUidResolver = new PowerStatsUidResolver();
         mStats = new BatteryStatsImpl(mBatteryStatsConfig, Clock.SYSTEM_CLOCK, mMonotonicClock,
                 systemDir, handler, this, this, mUserManagerUserInfoProvider, mPowerProfile,
-                mCpuScalingPolicies);
+                mCpuScalingPolicies, mPowerStatsUidResolver);
         mWorker = new BatteryExternalStatsWorker(context, mStats);
         mStats.setExternalStatsSyncLocked(mWorker);
         mStats.setRadioScanningTimeoutLocked(mContext.getResources().getInteger(
@@ -419,8 +425,6 @@
 
         AggregatedPowerStatsConfig aggregatedPowerStatsConfig = getAggregatedPowerStatsConfig();
         mPowerStatsStore = new PowerStatsStore(systemDir, mHandler, aggregatedPowerStatsConfig);
-        mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context, mStats,
-                mPowerStatsStore);
         mPowerStatsAggregator = new PowerStatsAggregator(aggregatedPowerStatsConfig,
                 mStats.getHistory());
         final long aggregatedPowerStatsSpanDuration = context.getResources().getInteger(
@@ -429,7 +433,14 @@
                 com.android.internal.R.integer.config_powerStatsAggregationPeriod);
         mPowerStatsScheduler = new PowerStatsScheduler(context, mPowerStatsAggregator,
                 aggregatedPowerStatsSpanDuration, powerStatsAggregationPeriod, mPowerStatsStore,
-                Clock.SYSTEM_CLOCK, mMonotonicClock, mHandler, mStats, mBatteryUsageStatsProvider);
+                Clock.SYSTEM_CLOCK, mMonotonicClock, mHandler, mStats);
+        PowerStatsExporter powerStatsExporter =
+                new PowerStatsExporter(mPowerStatsStore, mPowerStatsAggregator);
+        mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context,
+                powerStatsExporter, mPowerProfile, mCpuScalingPolicies,
+                mPowerStatsStore, Clock.SYSTEM_CLOCK);
+        mStats.saveBatteryUsageStatsOnReset(mBatteryUsageStatsProvider, mPowerStatsStore);
+        mDumpHelper = new BatteryStatsDumpHelperImpl(mBatteryUsageStatsProvider);
         mCpuWakeupStats = new CpuWakeupStats(context, R.xml.irq_device_map, mHandler);
         mConfigFile = new AtomicFile(new File(systemDir, "battery_usage_stats_config"));
     }
@@ -469,9 +480,10 @@
     }
 
     public void systemServicesReady() {
+        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(Flags.streamlinedBatteryStats());
+        mWorker.systemServicesReady();
         mStats.systemServicesReady(mContext);
         mCpuWakeupStats.systemServicesReady();
-        mWorker.systemServicesReady();
         final INetworkManagementService nms = INetworkManagementService.Stub.asInterface(
                 ServiceManager.getService(Context.NETWORKMANAGEMENT_SERVICE));
         final ConnectivityManager cm = mContext.getSystemService(ConnectivityManager.class);
@@ -775,25 +787,15 @@
     }
 
     void addIsolatedUid(final int isolatedUid, final int appUid) {
-        synchronized (mLock) {
-            final long elapsedRealtime = SystemClock.elapsedRealtime();
-            final long uptime = SystemClock.uptimeMillis();
-            mHandler.post(() -> {
-                synchronized (mStats) {
-                    mStats.addIsolatedUidLocked(isolatedUid, appUid, elapsedRealtime, uptime);
-                }
-            });
-        }
+        mPowerStatsUidResolver.noteIsolatedUidAdded(isolatedUid, appUid);
+        FrameworkStatsLog.write(FrameworkStatsLog.ISOLATED_UID_CHANGED, appUid, isolatedUid,
+                FrameworkStatsLog.ISOLATED_UID_CHANGED__EVENT__CREATED);
     }
 
     void removeIsolatedUid(final int isolatedUid, final int appUid) {
-        synchronized (mLock) {
-            mHandler.post(() -> {
-                synchronized (mStats) {
-                    mStats.scheduleRemoveIsolatedUidLocked(isolatedUid, appUid);
-                }
-            });
-        }
+        mPowerStatsUidResolver.noteIsolatedUidRemoved(isolatedUid, appUid);
+        FrameworkStatsLog.write(FrameworkStatsLog.ISOLATED_UID_CHANGED, -1, isolatedUid,
+                FrameworkStatsLog.ISOLATED_UID_CHANGED__EVENT__REMOVED);
     }
 
     void noteProcessStart(final String name, final int uid) {
@@ -877,12 +879,15 @@
 
         awaitCompletion();
 
-        if (mBatteryUsageStatsProvider.shouldUpdateStats(queries,
+        if (BatteryUsageStatsProvider.shouldUpdateStats(queries,
+                SystemClock.elapsedRealtime(),
                 mWorker.getLastCollectionTimeStamp())) {
             syncStats("get-stats", BatteryExternalStatsWorker.UPDATE_ALL);
         }
 
-        return mBatteryUsageStatsProvider.getBatteryUsageStats(queries);
+        synchronized (mStats) {
+            return mBatteryUsageStatsProvider.getBatteryUsageStats(mStats, queries);
+        }
     }
 
     /** Register callbacks for statsd pulled atoms. */
@@ -2723,7 +2728,7 @@
         synchronized (mStats) {
             mStats.prepareForDumpLocked();
             BatteryUsageStats batteryUsageStats =
-                    mBatteryUsageStatsProvider.getBatteryUsageStats(query);
+                    mBatteryUsageStatsProvider.getBatteryUsageStats(mStats, query);
             if (proto) {
                 batteryUsageStats.dumpToProto(fd);
             } else {
@@ -3008,11 +3013,11 @@
                                         mBatteryStatsConfig, Clock.SYSTEM_CLOCK, mMonotonicClock,
                                         null, mStats.mHandler, null, null,
                                         mUserManagerUserInfoProvider, mPowerProfile,
-                                        mCpuScalingPolicies);
+                                        mCpuScalingPolicies, null);
                                 checkinStats.readSummaryFromParcel(in);
                                 in.recycle();
-                                checkinStats.dumpProtoLocked(
-                                        mContext, fd, apps, flags, historyStart);
+                                checkinStats.dumpProtoLocked(mContext, fd, apps, flags,
+                                        historyStart, mDumpHelper);
                                 mStats.mCheckinFile.delete();
                                 return;
                             }
@@ -3026,7 +3031,7 @@
             if (DBG) Slog.d(TAG, "begin dumpProtoLocked from UID " + Binder.getCallingUid());
             awaitCompletion();
             synchronized (mStats) {
-                mStats.dumpProtoLocked(mContext, fd, apps, flags, historyStart);
+                mStats.dumpProtoLocked(mContext, fd, apps, flags, historyStart, mDumpHelper);
                 if (writeData) {
                     mStats.writeAsyncLocked();
                 }
@@ -3050,11 +3055,11 @@
                                         mBatteryStatsConfig, Clock.SYSTEM_CLOCK, mMonotonicClock,
                                         null, mStats.mHandler, null, null,
                                         mUserManagerUserInfoProvider, mPowerProfile,
-                                        mCpuScalingPolicies);
+                                        mCpuScalingPolicies, null);
                                 checkinStats.readSummaryFromParcel(in);
                                 in.recycle();
                                 checkinStats.dumpCheckin(mContext, pw, apps, flags,
-                                        historyStart);
+                                        historyStart, mDumpHelper);
                                 mStats.mCheckinFile.delete();
                                 return;
                             }
@@ -3067,7 +3072,7 @@
             }
             if (DBG) Slog.d(TAG, "begin dumpCheckin from UID " + Binder.getCallingUid());
             awaitCompletion();
-            mStats.dumpCheckin(mContext, pw, apps, flags, historyStart);
+            mStats.dumpCheckin(mContext, pw, apps, flags, historyStart, mDumpHelper);
             if (writeData) {
                 mStats.writeAsyncLocked();
             }
@@ -3076,7 +3081,7 @@
             if (DBG) Slog.d(TAG, "begin dump from UID " + Binder.getCallingUid());
             awaitCompletion();
 
-            mStats.dump(mContext, pw, flags, reqUid, historyStart);
+            mStats.dump(mContext, pw, flags, reqUid, historyStart, mDumpHelper);
             if (writeData) {
                 mStats.writeAsyncLocked();
             }
diff --git a/services/core/java/com/android/server/am/ProcessList.java b/services/core/java/com/android/server/am/ProcessList.java
index cb2b5fb..4ff34b1 100644
--- a/services/core/java/com/android/server/am/ProcessList.java
+++ b/services/core/java/com/android/server/am/ProcessList.java
@@ -34,7 +34,7 @@
 import static android.os.Process.getTotalMemory;
 import static android.os.Process.killProcessQuiet;
 import static android.os.Process.startWebView;
-import static android.system.OsConstants.*;
+import static android.system.OsConstants.EAGAIN;
 
 import static com.android.sdksandbox.flags.Flags.selinuxSdkSandboxAudit;
 import static com.android.server.am.ActivityManagerDebugConfig.DEBUG_LRU;
@@ -133,7 +133,6 @@
 import com.android.internal.app.ProcessMap;
 import com.android.internal.os.Zygote;
 import com.android.internal.util.ArrayUtils;
-import com.android.internal.util.FrameworkStatsLog;
 import com.android.internal.util.MemInfoReader;
 import com.android.server.AppStateTracker;
 import com.android.server.LocalServices;
@@ -3299,8 +3298,6 @@
             // about the process state of the isolated UID *before* it is registered with the
             // owning application.
             mService.mBatteryStatsService.addIsolatedUid(uid, info.uid);
-            FrameworkStatsLog.write(FrameworkStatsLog.ISOLATED_UID_CHANGED, info.uid, uid,
-                    FrameworkStatsLog.ISOLATED_UID_CHANGED__EVENT__CREATED);
         }
         final ProcessRecord r = new ProcessRecord(mService, info, proc, uid,
                 sdkSandboxClientAppPackage,
diff --git a/services/core/java/com/android/server/appop/AppOpsCheckingServiceImpl.java b/services/core/java/com/android/server/appop/AppOpsCheckingServiceImpl.java
index 0ee7d9c..0916967 100644
--- a/services/core/java/com/android/server/appop/AppOpsCheckingServiceImpl.java
+++ b/services/core/java/com/android/server/appop/AppOpsCheckingServiceImpl.java
@@ -20,6 +20,7 @@
 import static android.app.AppOpsManager.MODE_FOREGROUND;
 import static android.app.AppOpsManager.OP_SCHEDULE_EXACT_ALARM;
 import static android.app.AppOpsManager.OP_USE_FULL_SCREEN_INTENT;
+import static android.companion.virtual.VirtualDeviceManager.PERSISTENT_DEVICE_ID_DEFAULT;
 
 import android.annotation.NonNull;
 import android.annotation.UserIdInt;
@@ -150,7 +151,7 @@
     }
 
     @Override
-    public SparseIntArray getNonDefaultUidModes(int uid) {
+    public SparseIntArray getNonDefaultUidModes(int uid, String persistentDeviceId) {
         synchronized (mLock) {
             SparseIntArray opModes = mUidModes.get(uid, null);
             if (opModes == null) {
@@ -176,7 +177,7 @@
     }
 
     @Override
-    public int getUidMode(int uid, int op) {
+    public int getUidMode(int uid, String persistentDeviceId, int op) {
         synchronized (mLock) {
             SparseIntArray opModes = mUidModes.get(uid, null);
             if (opModes == null) {
@@ -187,7 +188,7 @@
     }
 
     @Override
-    public boolean setUidMode(int uid, int op, int mode) {
+    public boolean setUidMode(int uid, String persistentDeviceId, int op, int mode) {
         final int defaultMode = AppOpsManager.opToDefaultMode(op);
         List<AppOpsModeChangedListener> listenersCopy;
         synchronized (mLock) {
@@ -329,7 +330,7 @@
     }
 
     @Override
-    public SparseBooleanArray getForegroundOps(int uid) {
+    public SparseBooleanArray getForegroundOps(int uid, String persistentDeviceId) {
         SparseBooleanArray result = new SparseBooleanArray();
         synchronized (mLock) {
             SparseIntArray modes = mUidModes.get(uid);
@@ -606,9 +607,17 @@
         for (final String pkg : packagesDeclaringPermission) {
             for (int userId : userIds) {
                 final int uid = pmi.getPackageUid(pkg, 0, userId);
-                final int oldMode = getUidMode(uid, OP_SCHEDULE_EXACT_ALARM);
+                final int oldMode =
+                        getUidMode(
+                                uid,
+                                PERSISTENT_DEVICE_ID_DEFAULT,
+                                OP_SCHEDULE_EXACT_ALARM);
                 if (oldMode == AppOpsManager.opToDefaultMode(OP_SCHEDULE_EXACT_ALARM)) {
-                    setUidMode(uid, OP_SCHEDULE_EXACT_ALARM, MODE_ALLOWED);
+                    setUidMode(
+                            uid,
+                            PERSISTENT_DEVICE_ID_DEFAULT,
+                            OP_SCHEDULE_EXACT_ALARM,
+                            MODE_ALLOWED);
                 }
             }
             // This appop is meant to be controlled at a uid level. So we leave package modes as
@@ -641,7 +650,10 @@
                 final int flags = permissionManager.getPermissionFlags(pkg, permissionName,
                         UserHandle.of(userId));
                 if ((flags & PackageManager.FLAG_PERMISSION_USER_SET) == 0) {
-                    setUidMode(uid, OP_USE_FULL_SCREEN_INTENT,
+                    setUidMode(
+                            uid,
+                            PERSISTENT_DEVICE_ID_DEFAULT,
+                            OP_USE_FULL_SCREEN_INTENT,
                             AppOpsManager.opToDefaultMode(OP_USE_FULL_SCREEN_INTENT));
                 }
             }
diff --git a/services/core/java/com/android/server/appop/AppOpsCheckingServiceInterface.java b/services/core/java/com/android/server/appop/AppOpsCheckingServiceInterface.java
index f6e6bc0..f056f6b 100644
--- a/services/core/java/com/android/server/appop/AppOpsCheckingServiceInterface.java
+++ b/services/core/java/com/android/server/appop/AppOpsCheckingServiceInterface.java
@@ -59,8 +59,9 @@
      * Returns a copy of non-default app-ops with op as keys and their modes as values for a uid.
      * Returns an empty SparseIntArray if nothing is set.
      * @param uid for which we need the app-ops and their modes.
+     * @param persistentDeviceId device for which we need the app-ops and their modes
      */
-    SparseIntArray getNonDefaultUidModes(int uid);
+    SparseIntArray getNonDefaultUidModes(int uid, String persistentDeviceId);
 
     /**
      * Returns a copy of non-default app-ops with op as keys and their modes as values for a package
@@ -75,20 +76,22 @@
      * Returns the app-op mode for a particular app-op of a uid.
      * Returns default op mode if the op mode for particular uid and op is not set.
      * @param uid user id for which we need the mode.
+     * @param persistentDeviceId device for which we need the mode
      * @param op app-op for which we need the mode.
      * @return mode of the app-op.
      */
-    int getUidMode(int uid, int op);
+    int getUidMode(int uid, String persistentDeviceId, int op);
 
     /**
      * Set the app-op mode for a particular uid and op.
      * The mode is not set if the mode is the same as the default mode for the op.
      * @param uid user id for which we want to set the mode.
+     * @param persistentDeviceId device for which we want to set the mode.
      * @param op app-op for which we want to set the mode.
      * @param mode mode for the app-op.
      * @return true if op mode is changed.
      */
-    boolean setUidMode(int uid, int op, @Mode int mode);
+    boolean setUidMode(int uid, String persistentDeviceId, int op, @Mode int mode);
 
     /**
      * Gets the app-op mode for a particular package.
@@ -130,10 +133,11 @@
 
     /**
      * @param uid UID to query foreground ops for.
+     * @param persistentDeviceId device to query foreground ops for
      * @return SparseBooleanArray where the keys are the op codes for which their modes are
      * MODE_FOREGROUND for the passed UID.
      */
-    SparseBooleanArray getForegroundOps(int uid);
+    SparseBooleanArray getForegroundOps(int uid, String persistentDeviceId);
 
     /**
      *
diff --git a/services/core/java/com/android/server/appop/AppOpsCheckingServiceLoggingDecorator.java b/services/core/java/com/android/server/appop/AppOpsCheckingServiceLoggingDecorator.java
index ccdf3a5..f6da166 100644
--- a/services/core/java/com/android/server/appop/AppOpsCheckingServiceLoggingDecorator.java
+++ b/services/core/java/com/android/server/appop/AppOpsCheckingServiceLoggingDecorator.java
@@ -60,9 +60,9 @@
     }
 
     @Override
-    public SparseIntArray getNonDefaultUidModes(int uid) {
+    public SparseIntArray getNonDefaultUidModes(int uid, String persistentDeviceId) {
         Log.i(LOG_TAG, "getNonDefaultUidModes(uid = " + uid + ")");
-        return mService.getNonDefaultUidModes(uid);
+        return mService.getNonDefaultUidModes(uid, persistentDeviceId);
     }
 
     @Override
@@ -73,15 +73,15 @@
     }
 
     @Override
-    public int getUidMode(int uid, int op) {
+    public int getUidMode(int uid, String persistentDeviceId, int op) {
         Log.i(LOG_TAG, "getUidMode(uid = " + uid + ", op = " + op + ")");
-        return mService.getUidMode(uid, op);
+        return mService.getUidMode(uid, persistentDeviceId, op);
     }
 
     @Override
-    public boolean setUidMode(int uid, int op, int mode) {
+    public boolean setUidMode(int uid, String persistentDeviceId, int op, int mode) {
         Log.i(LOG_TAG, "setUidMode(uid = " + uid + ", op = " + op + ", mode = " + mode + ")");
-        return mService.setUidMode(uid, op, mode);
+        return mService.setUidMode(uid, persistentDeviceId, op, mode);
     }
 
     @Override
@@ -117,9 +117,9 @@
     }
 
     @Override
-    public SparseBooleanArray getForegroundOps(int uid) {
+    public SparseBooleanArray getForegroundOps(int uid, String persistentDeviceId) {
         Log.i(LOG_TAG, "getForegroundOps(uid = " + uid + ")");
-        return mService.getForegroundOps(uid);
+        return mService.getForegroundOps(uid, persistentDeviceId);
     }
 
     @Override
diff --git a/services/core/java/com/android/server/appop/AppOpsCheckingServiceTracingDecorator.java b/services/core/java/com/android/server/appop/AppOpsCheckingServiceTracingDecorator.java
index c3a02a8..55cf7ed 100644
--- a/services/core/java/com/android/server/appop/AppOpsCheckingServiceTracingDecorator.java
+++ b/services/core/java/com/android/server/appop/AppOpsCheckingServiceTracingDecorator.java
@@ -81,11 +81,11 @@
     }
 
     @Override
-    public SparseIntArray getNonDefaultUidModes(int uid) {
+    public SparseIntArray getNonDefaultUidModes(int uid, String persistentDeviceId) {
         Trace.traceBegin(TRACE_TAG,
                 "TaggedTracingAppOpsCheckingServiceInterfaceImpl#getNonDefaultUidModes");
         try {
-            return mService.getNonDefaultUidModes(uid);
+            return mService.getNonDefaultUidModes(uid, persistentDeviceId);
         } finally {
             Trace.traceEnd(TRACE_TAG);
         }
@@ -103,20 +103,21 @@
     }
 
     @Override
-    public int getUidMode(int uid, int op) {
+    public int getUidMode(int uid, String persistentDeviceId, int op) {
         Trace.traceBegin(TRACE_TAG, "TaggedTracingAppOpsCheckingServiceInterfaceImpl#getUidMode");
         try {
-            return mService.getUidMode(uid, op);
+            return mService.getUidMode(uid, persistentDeviceId, op);
         } finally {
             Trace.traceEnd(TRACE_TAG);
         }
     }
 
     @Override
-    public boolean setUidMode(int uid, int op, @AppOpsManager.Mode int mode) {
+    public boolean setUidMode(
+            int uid, String persistentDeviceId, int op, @AppOpsManager.Mode int mode) {
         Trace.traceBegin(TRACE_TAG, "TaggedTracingAppOpsCheckingServiceInterfaceImpl#setUidMode");
         try {
-            return mService.setUidMode(uid, op, mode);
+            return mService.setUidMode(uid, persistentDeviceId, op, mode);
         } finally {
             Trace.traceEnd(TRACE_TAG);
         }
@@ -179,11 +180,11 @@
     }
 
     @Override
-    public SparseBooleanArray getForegroundOps(int uid) {
+    public SparseBooleanArray getForegroundOps(int uid, String persistentDeviceId) {
         Trace.traceBegin(TRACE_TAG,
                 "TaggedTracingAppOpsCheckingServiceInterfaceImpl#getForegroundOps");
         try {
-            return mService.getForegroundOps(uid);
+            return mService.getForegroundOps(uid, persistentDeviceId);
         } finally {
             Trace.traceEnd(TRACE_TAG);
         }
diff --git a/services/core/java/com/android/server/appop/AppOpsService.java b/services/core/java/com/android/server/appop/AppOpsService.java
index 14aab13..3446737 100644
--- a/services/core/java/com/android/server/appop/AppOpsService.java
+++ b/services/core/java/com/android/server/appop/AppOpsService.java
@@ -63,6 +63,7 @@
 import static android.app.AppOpsManager.opRestrictsRead;
 import static android.app.AppOpsManager.opToName;
 import static android.app.AppOpsManager.opToPublicName;
+import static android.companion.virtual.VirtualDeviceManager.PERSISTENT_DEVICE_ID_DEFAULT;
 import static android.content.pm.PermissionInfo.PROTECTION_DANGEROUS;
 import static android.content.pm.PermissionInfo.PROTECTION_FLAG_APPOP;
 
@@ -1349,7 +1350,10 @@
 
                 SparseBooleanArray foregroundOps = new SparseBooleanArray();
 
-                SparseBooleanArray uidForegroundOps = mAppOpsCheckingService.getForegroundOps(uid);
+                // TODO(b/299330771): Check uidForegroundOps for all devices.
+                SparseBooleanArray uidForegroundOps =
+                        mAppOpsCheckingService.getForegroundOps(
+                                uid, PERSISTENT_DEVICE_ID_DEFAULT);
                 for (int i = 0; i < uidForegroundOps.size(); i++) {
                     foregroundOps.put(uidForegroundOps.keyAt(i), true);
                 }
@@ -1369,10 +1373,16 @@
                         continue;
                     }
                     final int code = foregroundOps.keyAt(fgi);
-
-                    if (mAppOpsCheckingService.getUidMode(uidState.uid, code)
+                    // TODO(b/299330771): Notify op changes for all relevant devices.
+                    if (mAppOpsCheckingService.getUidMode(
+                                            uidState.uid,
+                                            PERSISTENT_DEVICE_ID_DEFAULT,
+                                            code)
                                     != AppOpsManager.opToDefaultMode(code)
-                            && mAppOpsCheckingService.getUidMode(uidState.uid, code)
+                            && mAppOpsCheckingService.getUidMode(
+                                            uidState.uid,
+                                            PERSISTENT_DEVICE_ID_DEFAULT,
+                                            code)
                                     == AppOpsManager.MODE_FOREGROUND) {
                         mHandler.sendMessage(PooledLambda.obtainMessage(
                                 AppOpsService::notifyOpChangedForAllPkgsInUid,
@@ -1489,7 +1499,11 @@
     @Nullable
     private ArrayList<AppOpsManager.OpEntry> collectUidOps(@NonNull UidState uidState,
             @Nullable int[] ops) {
-        final SparseIntArray opModes = mAppOpsCheckingService.getNonDefaultUidModes(uidState.uid);
+        // TODO(b/299330771): Make this methods device-aware, currently it represents only the
+        // primary device.
+        final SparseIntArray opModes =
+                mAppOpsCheckingService.getNonDefaultUidModes(
+                        uidState.uid, PERSISTENT_DEVICE_ID_DEFAULT);
         if (opModes == null) {
             return null;
         }
@@ -1844,16 +1858,22 @@
                 uidState = new UidState(uid);
                 mUidStates.put(uid, uidState);
             }
-            if (mAppOpsCheckingService.getUidMode(uidState.uid, code)
+            // TODO(b/266164193): Ensure this behavior is device-aware after uid op mode for runtime
+            //  permissions is deprecated.
+            if (mAppOpsCheckingService.getUidMode(
+                            uidState.uid, PERSISTENT_DEVICE_ID_DEFAULT, code)
                     != AppOpsManager.opToDefaultMode(code)) {
-                previousMode = mAppOpsCheckingService.getUidMode(uidState.uid, code);
+                previousMode =
+                        mAppOpsCheckingService.getUidMode(
+                                uidState.uid, PERSISTENT_DEVICE_ID_DEFAULT, code);
             } else {
                 // doesn't look right but is legacy behavior.
                 previousMode = MODE_DEFAULT;
             }
 
             mIgnoredCallback = permissionPolicyCallback;
-            if (!mAppOpsCheckingService.setUidMode(uidState.uid, code, mode)) {
+            if (!mAppOpsCheckingService.setUidMode(
+                    uidState.uid, PERSISTENT_DEVICE_ID_DEFAULT, code, mode)) {
                 return;
             }
             if (mode != MODE_ERRORED && mode != previousMode) {
@@ -2275,8 +2295,10 @@
             boolean changed = false;
             for (int i = mUidStates.size() - 1; i >= 0; i--) {
                 UidState uidState = mUidStates.valueAt(i);
-
-                SparseIntArray opModes = mAppOpsCheckingService.getNonDefaultUidModes(uidState.uid);
+                // TODO(b/299330771): Check non default modes for all devices.
+                SparseIntArray opModes =
+                        mAppOpsCheckingService.getNonDefaultUidModes(
+                                uidState.uid, PERSISTENT_DEVICE_ID_DEFAULT);
                 if (opModes != null && (uidState.uid == reqUid || reqUid == -1)) {
                     final int uidOpCount = opModes.size();
                     for (int j = uidOpCount - 1; j >= 0; j--) {
@@ -2285,7 +2307,12 @@
                             int previousMode = opModes.valueAt(j);
                             int newMode = isUidOpGrantedByRole(uidState.uid, code) ? MODE_ALLOWED :
                                     AppOpsManager.opToDefaultMode(code);
-                            mAppOpsCheckingService.setUidMode(uidState.uid, code, newMode);
+                            // TODO(b/299330771): Set mode for all necessary devices.
+                            mAppOpsCheckingService.setUidMode(
+                                    uidState.uid,
+                                    PERSISTENT_DEVICE_ID_DEFAULT,
+                                    code,
+                                    newMode);
                             for (String packageName : getPackagesForUid(uidState.uid)) {
                                 callbacks = addCallbacks(callbacks, code, uidState.uid, packageName,
                                         previousMode, mOpModeWatchers.get(code));
@@ -2601,10 +2628,14 @@
             }
             code = AppOpsManager.opToSwitch(code);
             UidState uidState = getUidStateLocked(uid, false);
+            // TODO(b/299330771): Check mode for the relevant device.
             if (uidState != null
-                    && mAppOpsCheckingService.getUidMode(uidState.uid, code)
+                    && mAppOpsCheckingService.getUidMode(
+                                    uidState.uid, PERSISTENT_DEVICE_ID_DEFAULT, code)
                             != AppOpsManager.opToDefaultMode(code)) {
-                final int rawMode = mAppOpsCheckingService.getUidMode(uidState.uid, code);
+                final int rawMode =
+                        mAppOpsCheckingService.getUidMode(
+                                uidState.uid, PERSISTENT_DEVICE_ID_DEFAULT, code);
                 return raw ? rawMode : uidState.evalMode(code, rawMode);
             }
             Op op = getOpLocked(code, uid, packageName, null, false, pvr.bypass, /* edit */ false);
@@ -2851,13 +2882,19 @@
                 return new SyncNotedAppOp(AppOpsManager.MODE_IGNORED, code, attributionTag,
                         packageName);
             }
+            // TODO(b/299330771): Check mode for the relevant device.
             // If there is a non-default per UID policy (we set UID op mode only if
             // non-default) it takes over, otherwise use the per package policy.
-            if (mAppOpsCheckingService.getUidMode(uidState.uid, switchCode)
+            if (mAppOpsCheckingService.getUidMode(
+                            uidState.uid, PERSISTENT_DEVICE_ID_DEFAULT, switchCode)
                     != AppOpsManager.opToDefaultMode(switchCode)) {
                 final int uidMode =
                         uidState.evalMode(
-                                code, mAppOpsCheckingService.getUidMode(uidState.uid, switchCode));
+                                code,
+                                mAppOpsCheckingService.getUidMode(
+                                        uidState.uid,
+                                        PERSISTENT_DEVICE_ID_DEFAULT,
+                                        switchCode));
                 if (uidMode != AppOpsManager.MODE_ALLOWED) {
                     if (DEBUG) Slog.d(TAG, "noteOperation: uid reject #" + uidMode + " for code "
                             + switchCode + " (" + code + ") uid " + uid + " package "
@@ -3396,13 +3433,19 @@
             isRestricted = isOpRestrictedLocked(uid, code, packageName, attributionTag, pvr.bypass,
                     false);
             final int switchCode = AppOpsManager.opToSwitch(code);
+            // TODO(b/299330771): Check mode for the relevant device.
             // If there is a non-default per UID policy (we set UID op mode only if
             // non-default) it takes over, otherwise use the per package policy.
-            if (mAppOpsCheckingService.getUidMode(uidState.uid, switchCode)
+            if (mAppOpsCheckingService.getUidMode(
+                            uidState.uid, PERSISTENT_DEVICE_ID_DEFAULT, switchCode)
                     != AppOpsManager.opToDefaultMode(switchCode)) {
                 final int uidMode =
                         uidState.evalMode(
-                                code, mAppOpsCheckingService.getUidMode(uidState.uid, switchCode));
+                                code,
+                                mAppOpsCheckingService.getUidMode(
+                                        uidState.uid,
+                                        PERSISTENT_DEVICE_ID_DEFAULT,
+                                        switchCode));
                 if (!shouldStartForMode(uidMode, startIfModeDefault)) {
                     if (DEBUG) {
                         Slog.d(TAG, "startOperation: uid reject #" + uidMode + " for code "
@@ -3511,13 +3554,19 @@
             isRestricted = isOpRestrictedLocked(uid, code, packageName, attributionTag, pvr.bypass,
                     false);
             final int switchCode = AppOpsManager.opToSwitch(code);
+            // TODO(b/299330771): Check mode for the relevant device.
             // If there is a non-default mode per UID policy (we set UID op mode only if
             // non-default) it takes over, otherwise use the per package policy.
-            if (mAppOpsCheckingService.getUidMode(uidState.uid, switchCode)
+            if (mAppOpsCheckingService.getUidMode(
+                            uidState.uid, PERSISTENT_DEVICE_ID_DEFAULT, switchCode)
                     != AppOpsManager.opToDefaultMode(switchCode)) {
                 final int uidMode =
                         uidState.evalMode(
-                                code, mAppOpsCheckingService.getUidMode(uidState.uid, switchCode));
+                                code,
+                                mAppOpsCheckingService.getUidMode(
+                                        uidState.uid,
+                                        PERSISTENT_DEVICE_ID_DEFAULT,
+                                        switchCode));
                 if (!shouldStartForMode(uidMode, startIfModeDefault)) {
                     if (DEBUG) {
                         Slog.d(TAG, "startOperation: uid reject #" + uidMode + " for code "
@@ -5664,8 +5713,10 @@
             }
             for (int i=0; i<mUidStates.size(); i++) {
                 UidState uidState = mUidStates.valueAt(i);
+                // TODO(b/299330771): Get modes for all devices.
                 final SparseIntArray opModes =
-                        mAppOpsCheckingService.getNonDefaultUidModes(uidState.uid);
+                        mAppOpsCheckingService.getNonDefaultUidModes(
+                                uidState.uid, PERSISTENT_DEVICE_ID_DEFAULT);
                 final ArrayMap<String, Ops> pkgOps = uidState.pkgOps;
 
                 if (dumpWatchers || dumpHistory) {
diff --git a/services/core/java/com/android/server/appop/AppOpsServiceTestingShim.java b/services/core/java/com/android/server/appop/AppOpsServiceTestingShim.java
index 98e6476..c9fa9e6 100644
--- a/services/core/java/com/android/server/appop/AppOpsServiceTestingShim.java
+++ b/services/core/java/com/android/server/appop/AppOpsServiceTestingShim.java
@@ -65,9 +65,9 @@
     }
 
     @Override
-    public SparseIntArray getNonDefaultUidModes(int uid) {
-        SparseIntArray oldVal = mOldImplementation.getNonDefaultUidModes(uid);
-        SparseIntArray newVal = mNewImplementation.getNonDefaultUidModes(uid);
+    public SparseIntArray getNonDefaultUidModes(int uid, String persistentDeviceId) {
+        SparseIntArray oldVal = mOldImplementation.getNonDefaultUidModes(uid, persistentDeviceId);
+        SparseIntArray newVal = mNewImplementation.getNonDefaultUidModes(uid, persistentDeviceId);
 
         if (!Objects.equals(oldVal, newVal)) {
             signalImplDifference("getNonDefaultUidModes");
@@ -89,9 +89,9 @@
     }
 
     @Override
-    public int getUidMode(int uid, int op) {
-        int oldVal = mOldImplementation.getUidMode(uid, op);
-        int newVal = mNewImplementation.getUidMode(uid, op);
+    public int getUidMode(int uid, String persistentDeviceId, int op) {
+        int oldVal = mOldImplementation.getUidMode(uid, persistentDeviceId, op);
+        int newVal = mNewImplementation.getUidMode(uid, persistentDeviceId, op);
 
         if (oldVal != newVal) {
             signalImplDifference("getUidMode");
@@ -101,9 +101,9 @@
     }
 
     @Override
-    public boolean setUidMode(int uid, int op, int mode) {
-        boolean oldVal = mOldImplementation.setUidMode(uid, op, mode);
-        boolean newVal = mNewImplementation.setUidMode(uid, op, mode);
+    public boolean setUidMode(int uid, String persistentDeviceId, int op, int mode) {
+        boolean oldVal = mOldImplementation.setUidMode(uid, persistentDeviceId, op, mode);
+        boolean newVal = mNewImplementation.setUidMode(uid, persistentDeviceId, op, mode);
 
         if (oldVal != newVal) {
             signalImplDifference("setUidMode");
@@ -155,9 +155,9 @@
     }
 
     @Override
-    public SparseBooleanArray getForegroundOps(int uid) {
-        SparseBooleanArray oldVal = mOldImplementation.getForegroundOps(uid);
-        SparseBooleanArray newVal = mNewImplementation.getForegroundOps(uid);
+    public SparseBooleanArray getForegroundOps(int uid, String persistentDeviceId) {
+        SparseBooleanArray oldVal = mOldImplementation.getForegroundOps(uid, persistentDeviceId);
+        SparseBooleanArray newVal = mNewImplementation.getForegroundOps(uid, persistentDeviceId);
 
         if (!Objects.equals(oldVal, newVal)) {
             signalImplDifference("getForegroundOps");
diff --git a/services/core/java/com/android/server/flags/OWNERS b/services/core/java/com/android/server/flags/OWNERS
new file mode 100644
index 0000000..535a750
--- /dev/null
+++ b/services/core/java/com/android/server/flags/OWNERS
@@ -0,0 +1 @@
+per-file pinner.aconfig = edgararriaga@google.com
\ No newline at end of file
diff --git a/services/core/java/com/android/server/flags/pinner.aconfig b/services/core/java/com/android/server/flags/pinner.aconfig
new file mode 100644
index 0000000..606a6be
--- /dev/null
+++ b/services/core/java/com/android/server/flags/pinner.aconfig
@@ -0,0 +1,8 @@
+package: "com.android.server.flags"
+
+flag {
+    name: "pin_webview"
+    namespace: "system_performance"
+    description: "This flag controls if webview should be pinned in memory."
+    bug: "307594624"
+}
\ No newline at end of file
diff --git a/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java b/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java
index aadd03b..894226c 100644
--- a/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java
+++ b/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java
@@ -33,6 +33,7 @@
 import org.xmlpull.v1.XmlPullParserException;
 
 import java.io.IOException;
+import java.io.StringWriter;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
@@ -265,4 +266,11 @@
             ipw.decreaseIndent();
         }
     }
+
+    @Override
+    public String toString() {
+        StringWriter sw = new StringWriter();
+        dump(new IndentingPrintWriter(sw));
+        return sw.toString();
+    }
 }
diff --git a/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java b/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
index f9d57e4..a8eda3c 100644
--- a/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
+++ b/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
@@ -44,11 +44,8 @@
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.power.EnergyConsumerStats;
-import com.android.internal.util.FrameworkStatsLog;
 import com.android.server.LocalServices;
 
-import libcore.util.EmptyArray;
-
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
@@ -128,9 +125,6 @@
     private boolean mUseLatestStates = true;
 
     @GuardedBy("this")
-    private final IntArray mUidsToRemove = new IntArray();
-
-    @GuardedBy("this")
     private Future<?> mWakelockChangesUpdate;
 
     @GuardedBy("this")
@@ -260,7 +254,6 @@
 
     @Override
     public synchronized Future<?> scheduleCpuSyncDueToRemovedUid(int uid) {
-        mUidsToRemove.add(uid);
         return scheduleSyncLocked("remove-uid", UPDATE_CPU);
     }
 
@@ -459,7 +452,6 @@
             // Capture a snapshot of the state we are meant to process.
             final int updateFlags;
             final String reason;
-            final int[] uidsToRemove;
             final boolean onBattery;
             final boolean onBatteryScreenOff;
             final int screenState;
@@ -468,7 +460,6 @@
             synchronized (BatteryExternalStatsWorker.this) {
                 updateFlags = mUpdateFlags;
                 reason = mCurrentReason;
-                uidsToRemove = mUidsToRemove.size() > 0 ? mUidsToRemove.toArray() : EmptyArray.INT;
                 onBattery = mOnBattery;
                 onBatteryScreenOff = mOnBatteryScreenOff;
                 screenState = mScreenState;
@@ -476,7 +467,6 @@
                 useLatestStates = mUseLatestStates;
                 mUpdateFlags = 0;
                 mCurrentReason = null;
-                mUidsToRemove.clear();
                 mCurrentFuture = null;
                 mUseLatestStates = true;
                 if ((updateFlags & UPDATE_ALL) == UPDATE_ALL) {
@@ -512,12 +502,6 @@
 
                 // Clean up any UIDs if necessary.
                 synchronized (mStats) {
-                    for (int uid : uidsToRemove) {
-                        FrameworkStatsLog.write(FrameworkStatsLog.ISOLATED_UID_CHANGED, -1, uid,
-                                FrameworkStatsLog.ISOLATED_UID_CHANGED__EVENT__REMOVED);
-                        mStats.maybeRemoveIsolatedUidLocked(uid, SystemClock.elapsedRealtime(),
-                                SystemClock.uptimeMillis());
-                    }
                     mStats.clearPendingRemovedUidsLocked();
                 }
             } catch (Exception e) {
diff --git a/services/core/java/com/android/server/power/stats/BatteryStatsDumpHelperImpl.java b/services/core/java/com/android/server/power/stats/BatteryStatsDumpHelperImpl.java
new file mode 100644
index 0000000..ad146af
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsDumpHelperImpl.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import android.os.BatteryStats;
+import android.os.BatteryUsageStats;
+import android.os.BatteryUsageStatsQuery;
+
+public class BatteryStatsDumpHelperImpl implements BatteryStats.BatteryStatsDumpHelper {
+    private final BatteryUsageStatsProvider mBatteryUsageStatsProvider;
+
+    public BatteryStatsDumpHelperImpl(BatteryUsageStatsProvider batteryUsageStatsProvider) {
+        mBatteryUsageStatsProvider = batteryUsageStatsProvider;
+    }
+
+    @Override
+    public BatteryUsageStats getBatteryUsageStats(BatteryStats batteryStats, boolean detailed) {
+        BatteryUsageStatsQuery.Builder builder = new BatteryUsageStatsQuery.Builder()
+                .setMaxStatsAgeMs(0);
+        if (detailed) {
+            builder.includePowerModels().includeProcessStateData().includeVirtualUids();
+        }
+        return mBatteryUsageStatsProvider.getBatteryUsageStats((BatteryStatsImpl) batteryStats,
+                builder.build());
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
index eb40104..0491c14 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -185,7 +185,7 @@
     // TODO: remove "tcp" from network methods, since we measure total stats.
 
     // Current on-disk Parcel version. Must be updated when the format of the parcelable changes
-    public static final int VERSION = 213;
+    public static final int VERSION = 214;
 
     // The maximum number of names wakelocks we will keep track of
     // per uid; once the limit is reached, we batch the remaining wakelocks
@@ -220,6 +220,8 @@
     public static final int RESET_REASON_FULL_CHARGE = 3;
     public static final int RESET_REASON_ENERGY_CONSUMER_BUCKETS_CHANGE = 4;
     public static final int RESET_REASON_PLUGGED_IN_FOR_LONG_DURATION = 5;
+    @NonNull
+    private final MonotonicClock mMonotonicClock;
 
     protected Clock mClock;
 
@@ -393,19 +395,9 @@
         }
     }
 
-    /**
-     * Listener for the battery stats reset.
-     */
-    public interface BatteryResetListener {
-
-        /**
-         * Callback invoked immediately prior to resetting battery stats.
-         * @param resetReason One of the RESET_REASON_* constants.
-         */
-        void prepareForBatteryStatsReset(int resetReason);
-    }
-
-    private BatteryResetListener mBatteryResetListener;
+    private boolean mSaveBatteryUsageStatsOnReset;
+    private BatteryUsageStatsProvider mBatteryUsageStatsProvider;
+    private PowerStatsStore mPowerStatsStore;
 
     public interface BatteryCallback {
         public void batteryNeedsCpuUpdate();
@@ -787,13 +779,10 @@
     private BatteryCallback mCallback;
 
     /**
-     * Mapping isolated uids to the actual owning app uid.
+     * Mapping child uids to their parent uid.
      */
-    private final SparseIntArray mIsolatedUids = new SparseIntArray();
-    /**
-     * Internal reference count of isolated uids.
-     */
-    private final SparseIntArray mIsolatedUidRefCounts = new SparseIntArray();
+    @VisibleForTesting
+    protected final PowerStatsUidResolver mPowerStatsUidResolver;
 
     /**
      * The statistics we have collected organized by uids.
@@ -874,6 +863,8 @@
     long mUptimeStartUs;
     long mRealtimeUs;
     long mRealtimeStartUs;
+    long mMonotonicStartTime;
+    long mMonotonicEndTime = MonotonicClock.UNDEFINED;
 
     int mWakeLockNesting;
     boolean mWakeLockImportant;
@@ -1724,25 +1715,26 @@
     }
 
     @VisibleForTesting
-    public BatteryStatsImpl(Clock clock, File historyDirectory) {
+    public BatteryStatsImpl(Clock clock, File historyDirectory, @NonNull Handler handler,
+            @NonNull PowerStatsUidResolver powerStatsUidResolver) {
         init(clock);
         mBatteryStatsConfig = new BatteryStatsConfig.Builder().build();
-        mHandler = null;
+        mHandler = handler;
+        mPowerStatsUidResolver = powerStatsUidResolver;
         mConstants = new Constants(mHandler);
         mStartClockTimeMs = clock.currentTimeMillis();
         mDailyFile = null;
+        mMonotonicClock = new MonotonicClock(0, mClock);
         if (historyDirectory == null) {
             mCheckinFile = null;
             mStatsFile = null;
             mHistory = new BatteryStatsHistory(mConstants.MAX_HISTORY_FILES,
-                    mConstants.MAX_HISTORY_BUFFER, mStepDetailsCalculator, mClock,
-                    new MonotonicClock(0, mClock));
+                    mConstants.MAX_HISTORY_BUFFER, mStepDetailsCalculator, mClock, mMonotonicClock);
         } else {
             mCheckinFile = new AtomicFile(new File(historyDirectory, "batterystats-checkin.bin"));
             mStatsFile = new AtomicFile(new File(historyDirectory, "batterystats.bin"));
             mHistory = new BatteryStatsHistory(historyDirectory, mConstants.MAX_HISTORY_FILES,
-                    mConstants.MAX_HISTORY_BUFFER, mStepDetailsCalculator, mClock,
-                    new MonotonicClock(0, mClock));
+                    mConstants.MAX_HISTORY_BUFFER, mStepDetailsCalculator, mClock, mMonotonicClock);
         }
         mPlatformIdleStateCallback = null;
         mEnergyConsumerRetriever = null;
@@ -4278,92 +4270,51 @@
         }
     }
 
-    @GuardedBy("this")
-    public void addIsolatedUidLocked(int isolatedUid, int appUid) {
-        addIsolatedUidLocked(isolatedUid, appUid,
-                mClock.elapsedRealtime(), mClock.uptimeMillis());
+    private void onIsolatedUidAdded(int isolatedUid, int parentUid) {
+        long realtime = mClock.elapsedRealtime();
+        long uptime = mClock.uptimeMillis();
+        synchronized (this) {
+            getUidStatsLocked(parentUid, realtime, uptime).addIsolatedUid(isolatedUid);
+        }
     }
 
-    @GuardedBy("this")
-    @SuppressWarnings("GuardedBy")   // errorprone false positive on u.addIsolatedUid
-    public void addIsolatedUidLocked(int isolatedUid, int appUid,
-            long elapsedRealtimeMs, long uptimeMs) {
-        mIsolatedUids.put(isolatedUid, appUid);
-        mIsolatedUidRefCounts.put(isolatedUid, 1);
-        final Uid u = getUidStatsLocked(appUid, elapsedRealtimeMs, uptimeMs);
-        u.addIsolatedUid(isolatedUid);
+    private void onBeforeIsolatedUidRemoved(int isolatedUid, int parentUid) {
+        long realtime = mClock.elapsedRealtime();
+        mPowerStatsUidResolver.retainIsolatedUid(isolatedUid);
+        synchronized (this) {
+            mPendingRemovedUids.add(new UidToRemove(isolatedUid, realtime));
+        }
+        if (mExternalSync != null) {
+            mExternalSync.scheduleCpuSyncDueToRemovedUid(isolatedUid);
+        }
     }
 
-    /**
-     * Schedules a read of the latest cpu times before removing the isolated UID.
-     * @see #removeIsolatedUidLocked(int, int, int)
-     */
-    public void scheduleRemoveIsolatedUidLocked(int isolatedUid, int appUid) {
-        int curUid = mIsolatedUids.get(isolatedUid, -1);
-        if (curUid == appUid) {
-            if (mExternalSync != null) {
-                mExternalSync.scheduleCpuSyncDueToRemovedUid(isolatedUid);
-            }
+    private void onAfterIsolatedUidRemoved(int isolatedUid, int parentUid) {
+        long realtime = mClock.elapsedRealtime();
+        long uptime = mClock.uptimeMillis();
+        synchronized (this) {
+            getUidStatsLocked(parentUid, realtime, uptime).removeIsolatedUid(isolatedUid);
         }
     }
 
     /**
      * Isolated uid should only be removed after all wakelocks associated with the uid are stopped
      * and the cpu time-in-state has been read one last time for the uid.
-     *
-     * @see #scheduleRemoveIsolatedUidLocked(int, int)
-     *
-     * @return true if the isolated uid is actually removed.
      */
     @GuardedBy("this")
-    public boolean maybeRemoveIsolatedUidLocked(int isolatedUid, long elapsedRealtimeMs,
-            long uptimeMs) {
-        final int refCount = mIsolatedUidRefCounts.get(isolatedUid, 0) - 1;
-        if (refCount > 0) {
-            // Isolated uid is still being tracked
-            mIsolatedUidRefCounts.put(isolatedUid, refCount);
-            return false;
-        }
-
-        final int idx = mIsolatedUids.indexOfKey(isolatedUid);
-        if (idx >= 0) {
-            final int ownerUid = mIsolatedUids.valueAt(idx);
-            final Uid u = getUidStatsLocked(ownerUid, elapsedRealtimeMs, uptimeMs);
-            u.removeIsolatedUid(isolatedUid);
-            mIsolatedUids.removeAt(idx);
-            mIsolatedUidRefCounts.delete(isolatedUid);
-        } else {
-            Slog.w(TAG, "Attempted to remove untracked isolated uid (" + isolatedUid + ")");
-        }
-        mPendingRemovedUids.add(new UidToRemove(isolatedUid, elapsedRealtimeMs));
-
-        return true;
-    }
-
-    /**
-     * Increment the ref count for an isolated uid.
-     * call #maybeRemoveIsolatedUidLocked to decrement.
-     */
-    public void incrementIsolatedUidRefCount(int uid) {
-        final int refCount = mIsolatedUidRefCounts.get(uid, 0);
-        if (refCount <= 0) {
-            // Uid is not mapped or referenced
-            Slog.w(TAG,
-                    "Attempted to increment ref counted of untracked isolated uid (" + uid + ")");
-            return;
-        }
-        mIsolatedUidRefCounts.put(uid, refCount + 1);
+    public void releaseIsolatedUidLocked(int isolatedUid, long elapsedRealtimeMs, long uptimeMs) {
+        mPowerStatsUidResolver.releaseIsolatedUid(isolatedUid);
     }
 
     private int mapUid(int uid) {
         if (Process.isSdkSandboxUid(uid)) {
             return Process.getAppUidForSdkSandboxUid(uid);
         }
-        return mapIsolatedUid(uid);
+        return mPowerStatsUidResolver.mapUid(uid);
     }
 
     private int mapIsolatedUid(int uid) {
-        return mIsolatedUids.get(/*key=*/uid, /*valueIfKeyNotFound=*/uid);
+        return mPowerStatsUidResolver.mapUid(uid);
     }
 
     @GuardedBy("this")
@@ -4745,7 +4696,7 @@
             if (mappedUid != uid) {
                 // Prevent the isolated uid mapping from being removed while the wakelock is
                 // being held.
-                incrementIsolatedUidRefCount(uid);
+                mPowerStatsUidResolver.retainIsolatedUid(uid);
             }
             if (mOnBatteryScreenOffTimeBase.isRunning()) {
                 // We only update the cpu time when a wake lock is acquired if the screen is off.
@@ -4825,7 +4776,7 @@
 
             if (mappedUid != uid) {
                 // Decrement the ref count for the isolated uid and delete the mapping if uneeded.
-                maybeRemoveIsolatedUidLocked(uid, elapsedRealtimeMs, uptimeMs);
+                releaseIsolatedUidLocked(uid, elapsedRealtimeMs, uptimeMs);
             }
         }
     }
@@ -4996,7 +4947,7 @@
         if (mappedUid != uid) {
             // Prevent the isolated uid mapping from being removed while the wakelock is
             // being held.
-            incrementIsolatedUidRefCount(uid);
+            mPowerStatsUidResolver.retainIsolatedUid(uid);
         }
     }
 
@@ -5048,7 +4999,7 @@
                 historyName, mappedUid);
         if (mappedUid != uid) {
             // Decrement the ref count for the isolated uid and delete the mapping if uneeded.
-            maybeRemoveIsolatedUidLocked(uid, elapsedRealtimeMs, uptimeMs);
+            releaseIsolatedUidLocked(uid, elapsedRealtimeMs, uptimeMs);
         }
     }
 
@@ -7642,35 +7593,53 @@
     /**
      * Returns the names of custom power components.
      */
-    @GuardedBy("this")
     @Override
     public @NonNull String[] getCustomEnergyConsumerNames() {
-        if (mEnergyConsumerStatsConfig == null) {
-            return new String[0];
-        }
-        final String[] names = mEnergyConsumerStatsConfig.getCustomBucketNames();
-        for (int i = 0; i < names.length; i++) {
-            if (TextUtils.isEmpty(names[i])) {
-                names[i] = "CUSTOM_" + BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID + i;
+        synchronized (this) {
+            if (mEnergyConsumerStatsConfig == null) {
+                return new String[0];
             }
+            final String[] names = mEnergyConsumerStatsConfig.getCustomBucketNames();
+            for (int i = 0; i < names.length; i++) {
+                if (TextUtils.isEmpty(names[i])) {
+                    names[i] = "CUSTOM_" + BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID + i;
+                }
+            }
+            return names;
         }
-        return names;
     }
 
-    @GuardedBy("this")
-    @Override public long getStartClockTime() {
-        final long currentTimeMs = mClock.currentTimeMillis();
-        if ((currentTimeMs > MILLISECONDS_IN_YEAR
-                && mStartClockTimeMs < (currentTimeMs - MILLISECONDS_IN_YEAR))
+    @Override
+    public long getStartClockTime() {
+        synchronized (this) {
+            final long currentTimeMs = mClock.currentTimeMillis();
+            if ((currentTimeMs > MILLISECONDS_IN_YEAR
+                    && mStartClockTimeMs < (currentTimeMs - MILLISECONDS_IN_YEAR))
                 || (mStartClockTimeMs > currentTimeMs)) {
-            // If the start clock time has changed by more than a year, then presumably
-            // the previous time was completely bogus.  So we are going to figure out a
-            // new time based on how much time has elapsed since we started counting.
-            mHistory.recordCurrentTimeChange(mClock.elapsedRealtime(), mClock.uptimeMillis(),
-                    currentTimeMs);
-            adjustStartClockTime(currentTimeMs);
+                // If the start clock time has changed by more than a year, then presumably
+                // the previous time was completely bogus.  So we are going to figure out a
+                // new time based on how much time has elapsed since we started counting.
+                mHistory.recordCurrentTimeChange(mClock.elapsedRealtime(), mClock.uptimeMillis(),
+                        currentTimeMs);
+                adjustStartClockTime(currentTimeMs);
+            }
+            return mStartClockTimeMs;
         }
-        return mStartClockTimeMs;
+    }
+
+    /**
+     * Returns the monotonic time when the BatteryStats session started.
+     */
+    public long getMonotonicStartTime() {
+        return mMonotonicStartTime;
+    }
+
+    /**
+     * Returns the monotonic time when the BatteryStats session ended, or
+     * {@link MonotonicClock#UNDEFINED} if the session is still ongoing.
+     */
+    public long getMonotonicEndTime() {
+        return mMonotonicEndTime;
     }
 
     @Override public String getStartPlatformVersion() {
@@ -8197,7 +8166,9 @@
             return mProportionalSystemServiceUsage;
         }
 
-        @GuardedBy("mBsi")
+        /**
+         * Adds isolated UID to the list of children.
+         */
         public void addIsolatedUid(int isolatedUid) {
             if (mChildUids == null) {
                 mChildUids = new SparseArray<>();
@@ -8207,6 +8178,9 @@
             mChildUids.put(isolatedUid, new ChildUid());
         }
 
+        /**
+         * Removes isolated UID from the list of children.
+         */
         public void removeIsolatedUid(int isolatedUid) {
             final int idx = mChildUids == null ? -1 : mChildUids.indexOfKey(isolatedUid);
             if (idx < 0) {
@@ -10910,15 +10884,18 @@
             @NonNull Handler handler, @Nullable PlatformIdleStateCallback cb,
             @Nullable EnergyStatsRetriever energyStatsCb,
             @NonNull UserInfoProvider userInfoProvider, @NonNull PowerProfile powerProfile,
-            @NonNull CpuScalingPolicies cpuScalingPolicies) {
+            @NonNull CpuScalingPolicies cpuScalingPolicies,
+            @NonNull PowerStatsUidResolver powerStatsUidResolver) {
         init(clock);
 
         mBatteryStatsConfig = config;
+        mMonotonicClock = monotonicClock;
         mHandler = new MyHandler(handler.getLooper());
         mConstants = new Constants(mHandler);
 
         mPowerProfile = powerProfile;
         mCpuScalingPolicies = cpuScalingPolicies;
+        mPowerStatsUidResolver = powerStatsUidResolver;
 
         initPowerProfile();
 
@@ -10927,17 +10904,17 @@
             mCheckinFile = null;
             mDailyFile = null;
             mHistory = new BatteryStatsHistory(mConstants.MAX_HISTORY_FILES,
-                    mConstants.MAX_HISTORY_BUFFER, mStepDetailsCalculator, mClock, monotonicClock);
+                    mConstants.MAX_HISTORY_BUFFER, mStepDetailsCalculator, mClock, mMonotonicClock);
         } else {
             mStatsFile = new AtomicFile(new File(systemDir, "batterystats.bin"));
             mCheckinFile = new AtomicFile(new File(systemDir, "batterystats-checkin.bin"));
             mDailyFile = new AtomicFile(new File(systemDir, "batterystats-daily.xml"));
             mHistory = new BatteryStatsHistory(systemDir, mConstants.MAX_HISTORY_FILES,
-                    mConstants.MAX_HISTORY_BUFFER, mStepDetailsCalculator, mClock, monotonicClock);
+                    mConstants.MAX_HISTORY_BUFFER, mStepDetailsCalculator, mClock, mMonotonicClock);
         }
 
         mCpuPowerStatsCollector = new CpuPowerStatsCollector(mCpuScalingPolicies, mPowerProfile,
-                () -> mBatteryVoltageMv, mHandler,
+                mPowerStatsUidResolver, () -> mBatteryVoltageMv, mHandler,
                 mBatteryStatsConfig.getPowerStatsThrottlePeriodCpu());
         mCpuPowerStatsCollector.addConsumer(this::recordPowerStats);
 
@@ -10954,6 +10931,23 @@
         mEnergyConsumerRetriever = energyStatsCb;
         mUserInfoProvider = userInfoProvider;
 
+        mPowerStatsUidResolver.addListener(new PowerStatsUidResolver.Listener() {
+            @Override
+            public void onIsolatedUidAdded(int isolatedUid, int parentUid) {
+                BatteryStatsImpl.this.onIsolatedUidAdded(isolatedUid, parentUid);
+            }
+
+            @Override
+            public void onBeforeIsolatedUidRemoved(int isolatedUid, int parentUid) {
+                BatteryStatsImpl.this.onBeforeIsolatedUidRemoved(isolatedUid, parentUid);
+            }
+
+            @Override
+            public void onAfterIsolatedUidRemoved(int isolatedUid, int parentUid) {
+                BatteryStatsImpl.this.onAfterIsolatedUidRemoved(isolatedUid, parentUid);
+            }
+        });
+
         // Notify statsd that the system is initially not in doze.
         mDeviceIdleMode = DEVICE_IDLE_MODE_OFF;
         FrameworkStatsLog.write(FrameworkStatsLog.DEVICE_IDLE_MODE_STATE_CHANGED, mDeviceIdleMode);
@@ -11497,6 +11491,7 @@
         mUptimeUs = 0;
         mRealtimeStartUs = realtimeUs;
         mUptimeStartUs = uptimeUs;
+        mMonotonicStartTime = mMonotonicClock.monotonicTime();
     }
 
     void initDischarge(long elapsedRealtimeUs) {
@@ -11517,8 +11512,17 @@
         mDischargeCounter.reset(false, elapsedRealtimeUs);
     }
 
-    public void setBatteryResetListener(BatteryResetListener batteryResetListener) {
-        mBatteryResetListener = batteryResetListener;
+    /**
+     * Associates the BatteryStatsImpl object with a BatteryUsageStatsProvider and PowerStatsStore
+     * to allow for a snapshot of battery usage stats to be taken and stored just before battery
+     * reset.
+     */
+    public void saveBatteryUsageStatsOnReset(
+            @NonNull BatteryUsageStatsProvider batteryUsageStatsProvider,
+            @NonNull PowerStatsStore powerStatsStore) {
+        mSaveBatteryUsageStatsOnReset = true;
+        mBatteryUsageStatsProvider = batteryUsageStatsProvider;
+        mPowerStatsStore = powerStatsStore;
     }
 
     @GuardedBy("this")
@@ -11557,9 +11561,7 @@
     @GuardedBy("this")
     private void resetAllStatsLocked(long uptimeMillis, long elapsedRealtimeMillis,
             int resetReason) {
-        if (mBatteryResetListener != null) {
-            mBatteryResetListener.prepareForBatteryStatsReset(resetReason);
-        }
+        saveBatteryUsageStatsOnReset(resetReason);
 
         final long uptimeUs = uptimeMillis * 1000;
         final long elapsedRealtimeUs = elapsedRealtimeMillis * 1000;
@@ -11707,6 +11709,31 @@
         mHandler.sendEmptyMessage(MSG_REPORT_RESET_STATS);
     }
 
+    private void saveBatteryUsageStatsOnReset(int resetReason) {
+        if (!mSaveBatteryUsageStatsOnReset
+                || resetReason == BatteryStatsImpl.RESET_REASON_CORRUPT_FILE) {
+            return;
+        }
+
+        final BatteryUsageStats batteryUsageStats;
+        synchronized (this) {
+            batteryUsageStats = mBatteryUsageStatsProvider.getBatteryUsageStats(this,
+                    new BatteryUsageStatsQuery.Builder()
+                            .setMaxStatsAgeMs(0)
+                            .includePowerModels()
+                            .includeProcessStateData()
+                            .build());
+        }
+
+        // TODO(b/188068523): BatteryUsageStats should use monotonic time for start and end
+        // Once that change is made, we will be able to use the BatteryUsageStats' monotonic
+        // start time
+        long monotonicStartTime =
+                mMonotonicClock.monotonicTime() - batteryUsageStats.getStatsDuration();
+        mHandler.post(() ->
+                mPowerStatsStore.storeBatteryUsageStats(monotonicStartTime, batteryUsageStats));
+    }
+
     @GuardedBy("this")
     private void initActiveHistoryEventsLocked(long elapsedRealtimeMs, long uptimeMs) {
         for (int i=0; i<HistoryItem.EVENT_COUNT; i++) {
@@ -15137,6 +15164,7 @@
             if (mKernelSingleUidTimeReader != null) {
                 mKernelSingleUidTimeReader.removeUidsInRange(startUid, endUid);
             }
+            mPowerStatsUidResolver.releaseUidsInRange(startUid, endUid);
             // Treat as one. We don't know how many uids there are in between.
             mNumUidsRemoved++;
         } else {
@@ -15192,10 +15220,11 @@
         mShuttingDown = true;
     }
 
-    @GuardedBy("this")
     @Override
     public boolean isProcessStateDataAvailable() {
-        return trackPerProcStateCpuTimes();
+        synchronized (this) {
+            return trackPerProcStateCpuTimes();
+        }
     }
 
     @GuardedBy("this")
@@ -15862,6 +15891,8 @@
         mUptimeUs = in.readLong();
         mRealtimeUs = in.readLong();
         mStartClockTimeMs = in.readLong();
+        mMonotonicStartTime = in.readLong();
+        mMonotonicEndTime = in.readLong();
         mStartPlatformVersion = in.readString();
         mEndPlatformVersion = in.readString();
         mOnBatteryTimeBase.readSummaryFromParcel(in);
@@ -16382,6 +16413,8 @@
         out.writeLong(computeUptime(nowUptime, STATS_SINCE_CHARGED));
         out.writeLong(computeRealtime(nowRealtime, STATS_SINCE_CHARGED));
         out.writeLong(mStartClockTimeMs);
+        out.writeLong(mMonotonicStartTime);
+        out.writeLong(mMonotonicClock.monotonicTime());
         out.writeString(mStartPlatformVersion);
         out.writeString(mEndPlatformVersion);
         mOnBatteryTimeBase.writeSummaryToParcel(out, nowUptime, nowRealtime);
@@ -16912,7 +16945,8 @@
     }
 
     @GuardedBy("this")
-    public void dump(Context context, PrintWriter pw, int flags, int reqUid, long histStart) {
+    public void dump(Context context, PrintWriter pw, int flags, int reqUid, long histStart,
+            BatteryStatsDumpHelper dumpHelper) {
         if (DEBUG) {
             pw.println("mOnBatteryTimeBase:");
             mOnBatteryTimeBase.dump(pw, "  ");
@@ -16984,7 +17018,7 @@
             pr.println("*** Camera timer:");
             mCameraOnTimer.logState(pr, "  ");
         }
-        super.dump(context, pw, flags, reqUid, histStart);
+        super.dump(context, pw, flags, reqUid, histStart, dumpHelper);
 
         synchronized (this) {
             pw.print("Per process state tracking available: ");
@@ -16998,15 +17032,7 @@
             pw.print("UIDs removed since the later of device start or stats reset: ");
             pw.println(mNumUidsRemoved);
 
-            pw.println("Currently mapped isolated uids:");
-            final int numIsolatedUids = mIsolatedUids.size();
-            for (int i = 0; i < numIsolatedUids; i++) {
-                final int isolatedUid = mIsolatedUids.keyAt(i);
-                final int ownerUid = mIsolatedUids.valueAt(i);
-                final int refCount = mIsolatedUidRefCounts.get(isolatedUid);
-                pw.println(
-                        "  " + isolatedUid + "->" + ownerUid + " (ref count = " + refCount + ")");
-            }
+            mPowerStatsUidResolver.dump(pw);
 
             pw.println();
             dumpConstantsLocked(pw);
@@ -17020,15 +17046,4 @@
             dumpEnergyConsumerStatsLocked(pw);
         }
     }
-
-    @Override
-    protected BatteryUsageStats getBatteryUsageStats(Context context, boolean detailed) {
-        final BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context, this);
-        BatteryUsageStatsQuery.Builder builder = new BatteryUsageStatsQuery.Builder()
-                .setMaxStatsAgeMs(0);
-        if (detailed) {
-            builder.includePowerModels().includeProcessStateData().includeVirtualUids();
-        }
-        return provider.getBatteryUsageStats(builder.build());
-    }
 }
diff --git a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
index 83d7d72..303c245 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
@@ -23,14 +23,12 @@
 import android.os.BatteryUsageStats;
 import android.os.BatteryUsageStatsQuery;
 import android.os.Process;
-import android.os.SystemClock;
 import android.os.UidBatteryConsumer;
 import android.util.Log;
 import android.util.Slog;
 import android.util.SparseArray;
 
-import com.android.internal.annotations.GuardedBy;
-import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.os.Clock;
 import com.android.internal.os.CpuScalingPolicies;
 import com.android.internal.os.PowerProfile;
 
@@ -45,27 +43,25 @@
 public class BatteryUsageStatsProvider {
     private static final String TAG = "BatteryUsageStatsProv";
     private final Context mContext;
-    private final BatteryStats mStats;
+    private boolean mPowerStatsExporterEnabled;
+    private final PowerStatsExporter mPowerStatsExporter;
     private final PowerStatsStore mPowerStatsStore;
     private final PowerProfile mPowerProfile;
     private final CpuScalingPolicies mCpuScalingPolicies;
+    private final Clock mClock;
     private final Object mLock = new Object();
     private List<PowerCalculator> mPowerCalculators;
 
-    public BatteryUsageStatsProvider(Context context, BatteryStats stats) {
-        this(context, stats, null);
-    }
-
-    @VisibleForTesting
-    public BatteryUsageStatsProvider(Context context, BatteryStats stats,
-            PowerStatsStore powerStatsStore) {
+    public BatteryUsageStatsProvider(Context context,
+            PowerStatsExporter powerStatsExporter,
+            PowerProfile powerProfile, CpuScalingPolicies cpuScalingPolicies,
+            PowerStatsStore powerStatsStore, Clock clock) {
         mContext = context;
-        mStats = stats;
+        mPowerStatsExporter = powerStatsExporter;
         mPowerStatsStore = powerStatsStore;
-        mPowerProfile = stats instanceof BatteryStatsImpl
-                ? ((BatteryStatsImpl) stats).getPowerProfile()
-                : new PowerProfile(context);
-        mCpuScalingPolicies = stats.getCpuScalingPolicies();
+        mPowerProfile = powerProfile;
+        mCpuScalingPolicies = cpuScalingPolicies;
+        mClock = clock;
     }
 
     private List<PowerCalculator> getPowerCalculators() {
@@ -75,7 +71,10 @@
 
                 // Power calculators are applied in the order of registration
                 mPowerCalculators.add(new BatteryChargeCalculator());
-                mPowerCalculators.add(new CpuPowerCalculator(mCpuScalingPolicies, mPowerProfile));
+                if (mPowerStatsExporterEnabled) {
+                    mPowerCalculators.add(
+                            new CpuPowerCalculator(mCpuScalingPolicies, mPowerProfile));
+                }
                 mPowerCalculators.add(new MemoryPowerCalculator(mPowerProfile));
                 mPowerCalculators.add(new WakelockPowerCalculator(mPowerProfile));
                 if (!BatteryStats.checkWifiOnly(mContext)) {
@@ -111,27 +110,28 @@
      * Returns true if the last update was too long ago for the tolerances specified
      * by the supplied queries.
      */
-    public boolean shouldUpdateStats(List<BatteryUsageStatsQuery> queries,
-            long lastUpdateTimeStampMs) {
+    public static boolean shouldUpdateStats(List<BatteryUsageStatsQuery> queries,
+            long elapsedRealtime, long lastUpdateTimeStampMs) {
         long allowableStatsAge = Long.MAX_VALUE;
         for (int i = queries.size() - 1; i >= 0; i--) {
             BatteryUsageStatsQuery query = queries.get(i);
             allowableStatsAge = Math.min(allowableStatsAge, query.getMaxStatsAge());
         }
 
-        return elapsedRealtime() - lastUpdateTimeStampMs > allowableStatsAge;
+        return elapsedRealtime - lastUpdateTimeStampMs > allowableStatsAge;
     }
 
     /**
      * Returns snapshots of battery attribution data, one per supplied query.
      */
-    public List<BatteryUsageStats> getBatteryUsageStats(List<BatteryUsageStatsQuery> queries) {
+    public List<BatteryUsageStats> getBatteryUsageStats(BatteryStatsImpl stats,
+            List<BatteryUsageStatsQuery> queries) {
         ArrayList<BatteryUsageStats> results = new ArrayList<>(queries.size());
-        synchronized (mStats) {
-            mStats.prepareForDumpLocked();
-            final long currentTimeMillis = currentTimeMillis();
+        synchronized (stats) {
+            stats.prepareForDumpLocked();
+            final long currentTimeMillis = mClock.currentTimeMillis();
             for (int i = 0; i < queries.size(); i++) {
-                results.add(getBatteryUsageStats(queries.get(i), currentTimeMillis));
+                results.add(getBatteryUsageStats(stats, queries.get(i), currentTimeMillis));
             }
         }
         return results;
@@ -140,60 +140,59 @@
     /**
      * Returns a snapshot of battery attribution data.
      */
-    @VisibleForTesting
-    public BatteryUsageStats getBatteryUsageStats(BatteryUsageStatsQuery query) {
-        synchronized (mStats) {
-            return getBatteryUsageStats(query, currentTimeMillis());
-        }
+    public BatteryUsageStats getBatteryUsageStats(BatteryStatsImpl stats,
+            BatteryUsageStatsQuery query) {
+        return getBatteryUsageStats(stats, query, mClock.currentTimeMillis());
     }
 
-    @GuardedBy("mStats")
-    private BatteryUsageStats getBatteryUsageStats(BatteryUsageStatsQuery query,
-            long currentTimeMs) {
+    private BatteryUsageStats getBatteryUsageStats(BatteryStatsImpl stats,
+            BatteryUsageStatsQuery query, long currentTimeMs) {
         if (query.getToTimestamp() == 0) {
-            return getCurrentBatteryUsageStats(query, currentTimeMs);
+            return getCurrentBatteryUsageStats(stats, query, currentTimeMs);
         } else {
-            return getAggregatedBatteryUsageStats(query);
+            return getAggregatedBatteryUsageStats(stats, query);
         }
     }
 
-    @GuardedBy("mStats")
-    private BatteryUsageStats getCurrentBatteryUsageStats(BatteryUsageStatsQuery query,
-            long currentTimeMs) {
-        final long realtimeUs = elapsedRealtime() * 1000;
-        final long uptimeUs = uptimeMillis() * 1000;
+    private BatteryUsageStats getCurrentBatteryUsageStats(BatteryStatsImpl stats,
+            BatteryUsageStatsQuery query, long currentTimeMs) {
+        final long realtimeUs = mClock.elapsedRealtime() * 1000;
+        final long uptimeUs = mClock.uptimeMillis() * 1000;
 
         final boolean includePowerModels = (query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_POWER_MODELS) != 0;
         final boolean includeProcessStateData = ((query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_PROCESS_STATE_DATA) != 0)
-                && mStats.isProcessStateDataAvailable();
+                && stats.isProcessStateDataAvailable();
         final boolean includeVirtualUids =  ((query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_VIRTUAL_UIDS) != 0);
         final double minConsumedPowerThreshold = query.getMinConsumedPowerThreshold();
 
         final BatteryUsageStats.Builder batteryUsageStatsBuilder = new BatteryUsageStats.Builder(
-                mStats.getCustomEnergyConsumerNames(), includePowerModels,
+                stats.getCustomEnergyConsumerNames(), includePowerModels,
                 includeProcessStateData, minConsumedPowerThreshold);
         // TODO(b/188068523): use a monotonic clock to ensure resilience of order and duration
-        // of stats sessions to wall-clock adjustments
-        batteryUsageStatsBuilder.setStatsStartTimestamp(mStats.getStartClockTime());
+        // of batteryUsageStats sessions to wall-clock adjustments
+        batteryUsageStatsBuilder.setStatsStartTimestamp(stats.getStartClockTime());
         batteryUsageStatsBuilder.setStatsEndTimestamp(currentTimeMs);
 
-        SparseArray<? extends BatteryStats.Uid> uidStats = mStats.getUidStats();
-        for (int i = uidStats.size() - 1; i >= 0; i--) {
-            final BatteryStats.Uid uid = uidStats.valueAt(i);
-            if (!includeVirtualUids && uid.getUid() == Process.SDK_SANDBOX_VIRTUAL_UID) {
-                continue;
-            }
+        synchronized (stats) {
+            SparseArray<? extends BatteryStats.Uid> uidStats = stats.getUidStats();
+            for (int i = uidStats.size() - 1; i >= 0; i--) {
+                final BatteryStats.Uid uid = uidStats.valueAt(i);
+                if (!includeVirtualUids && uid.getUid() == Process.SDK_SANDBOX_VIRTUAL_UID) {
+                    continue;
+                }
 
-            batteryUsageStatsBuilder.getOrCreateUidBatteryConsumerBuilder(uid)
-                    .setTimeInProcessStateMs(UidBatteryConsumer.PROCESS_STATE_BACKGROUND,
-                            getProcessBackgroundTimeMs(uid, realtimeUs))
-                    .setTimeInProcessStateMs(UidBatteryConsumer.PROCESS_STATE_FOREGROUND,
-                            getProcessForegroundTimeMs(uid, realtimeUs))
-                    .setTimeInProcessStateMs(UidBatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE,
-                            getProcessForegroundServiceTimeMs(uid, realtimeUs));
+                batteryUsageStatsBuilder.getOrCreateUidBatteryConsumerBuilder(uid)
+                        .setTimeInProcessStateMs(UidBatteryConsumer.PROCESS_STATE_BACKGROUND,
+                                getProcessBackgroundTimeMs(uid, realtimeUs))
+                        .setTimeInProcessStateMs(UidBatteryConsumer.PROCESS_STATE_FOREGROUND,
+                                getProcessForegroundTimeMs(uid, realtimeUs))
+                        .setTimeInProcessStateMs(
+                                UidBatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE,
+                                getProcessForegroundServiceTimeMs(uid, realtimeUs));
+            }
         }
 
         final int[] powerComponents = query.getPowerComponents();
@@ -202,8 +201,8 @@
             PowerCalculator powerCalculator = powerCalculators.get(i);
             if (powerComponents != null) {
                 boolean include = false;
-                for (int j = 0; j < powerComponents.length; j++) {
-                    if (powerCalculator.isPowerComponentSupported(powerComponents[j])) {
+                for (int powerComponent : powerComponents) {
+                    if (powerCalculator.isPowerComponentSupported(powerComponent)) {
                         include = true;
                         break;
                     }
@@ -212,26 +211,24 @@
                     continue;
                 }
             }
-            powerCalculator.calculate(batteryUsageStatsBuilder, mStats, realtimeUs, uptimeUs,
-                    query);
+            powerCalculator.calculate(batteryUsageStatsBuilder, stats, realtimeUs, uptimeUs, query);
+        }
+
+        if (mPowerStatsExporterEnabled) {
+            mPowerStatsExporter.exportAggregatedPowerStats(batteryUsageStatsBuilder,
+                    stats.getMonotonicStartTime(), stats.getMonotonicEndTime());
         }
 
         if ((query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_HISTORY) != 0) {
-            if (!(mStats instanceof BatteryStatsImpl)) {
-                throw new UnsupportedOperationException(
-                        "History cannot be included for " + getClass().getName());
-            }
-
-            BatteryStatsImpl batteryStatsImpl = (BatteryStatsImpl) mStats;
-            batteryUsageStatsBuilder.setBatteryHistory(batteryStatsImpl.copyHistory());
+            batteryUsageStatsBuilder.setBatteryHistory(stats.copyHistory());
         }
 
-        BatteryUsageStats stats = batteryUsageStatsBuilder.build();
+        BatteryUsageStats batteryUsageStats = batteryUsageStatsBuilder.build();
         if (includeProcessStateData) {
-            verify(stats);
+            verify(batteryUsageStats);
         }
-        return stats;
+        return batteryUsageStats;
     }
 
     // STOPSHIP(b/229906525): remove verification before shipping
@@ -308,15 +305,16 @@
                 / 1000;
     }
 
-    private BatteryUsageStats getAggregatedBatteryUsageStats(BatteryUsageStatsQuery query) {
+    private BatteryUsageStats getAggregatedBatteryUsageStats(BatteryStatsImpl stats,
+            BatteryUsageStatsQuery query) {
         final boolean includePowerModels = (query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_POWER_MODELS) != 0;
         final boolean includeProcessStateData = ((query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_PROCESS_STATE_DATA) != 0)
-                && mStats.isProcessStateDataAvailable();
+                && stats.isProcessStateDataAvailable();
         final double minConsumedPowerThreshold = query.getMinConsumedPowerThreshold();
 
-        final String[] customEnergyConsumerNames = mStats.getCustomEnergyConsumerNames();
+        final String[] customEnergyConsumerNames = stats.getCustomEnergyConsumerNames();
         final BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(
                 customEnergyConsumerNames, includePowerModels, includeProcessStateData,
                 minConsumedPowerThreshold);
@@ -386,27 +384,8 @@
         return builder.build();
     }
 
-    private long elapsedRealtime() {
-        if (mStats instanceof BatteryStatsImpl) {
-            return ((BatteryStatsImpl) mStats).mClock.elapsedRealtime();
-        } else {
-            return SystemClock.elapsedRealtime();
-        }
-    }
+    public void setPowerStatsExporterEnabled(boolean enabled) {
 
-    private long uptimeMillis() {
-        if (mStats instanceof BatteryStatsImpl) {
-            return ((BatteryStatsImpl) mStats).mClock.uptimeMillis();
-        } else {
-            return SystemClock.uptimeMillis();
-        }
-    }
-
-    private long currentTimeMillis() {
-        if (mStats instanceof BatteryStatsImpl) {
-            return ((BatteryStatsImpl) mStats).mClock.currentTimeMillis();
-        } else {
-            return System.currentTimeMillis();
-        }
+        mPowerStatsExporterEnabled = enabled;
     }
 }
diff --git a/services/core/java/com/android/server/power/stats/CpuAggregatedPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/CpuAggregatedPowerStatsProcessor.java
index f40eef2..ed9414f 100644
--- a/services/core/java/com/android/server/power/stats/CpuAggregatedPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/CpuAggregatedPowerStatsProcessor.java
@@ -64,7 +64,7 @@
     private PowerStats.Descriptor mLastUsedDescriptor;
     // Cached results of parsing of current PowerStats.Descriptor. Only refreshed when
     // mLastUsedDescriptor changes
-    private CpuPowerStatsCollector.StatsArrayLayout mStatsLayout;
+    private CpuPowerStatsCollector.CpuStatsArrayLayout mStatsLayout;
     // Sequence of steps for power estimation and intermediate results.
     private PowerEstimationPlan mPlan;
 
@@ -106,7 +106,7 @@
         }
 
         mLastUsedDescriptor = descriptor;
-        mStatsLayout = new CpuPowerStatsCollector.StatsArrayLayout();
+        mStatsLayout = new CpuPowerStatsCollector.CpuStatsArrayLayout();
         mStatsLayout.fromExtras(descriptor.extras);
 
         mTmpDeviceStatsArray = new long[descriptor.statsArrayLength];
@@ -149,7 +149,7 @@
 
         if (mPlan == null) {
             mPlan = new PowerEstimationPlan(stats.getConfig());
-            if (mStatsLayout.getCpuClusterEnergyConsumerCount() != 0) {
+            if (mStatsLayout.getEnergyConsumerCount() != 0) {
                 initEnergyConsumerToPowerBracketMaps();
             }
         }
@@ -212,7 +212,7 @@
      *          CL_2: [bracket3, bracket4]
      */
     private void initEnergyConsumerToPowerBracketMaps() {
-        int energyConsumerCount = mStatsLayout.getCpuClusterEnergyConsumerCount();
+        int energyConsumerCount = mStatsLayout.getEnergyConsumerCount();
         int powerBracketCount = mStatsLayout.getCpuPowerBracketCount();
 
         mEnergyConsumerToCombinedEnergyConsumerMap = new int[energyConsumerCount];
@@ -294,7 +294,7 @@
                 continue;
             }
 
-            intermediates.uptime += mStatsLayout.getUptime(mTmpDeviceStatsArray);
+            intermediates.uptime += mStatsLayout.getUsageDuration(mTmpDeviceStatsArray);
 
             for (int cluster = 0; cluster < mCpuClusterCount; cluster++) {
                 intermediates.timeByCluster[cluster] +=
@@ -351,7 +351,7 @@
         int cpuScalingStepCount = mStatsLayout.getCpuScalingStepCount();
         int powerBracketCount = mStatsLayout.getCpuPowerBracketCount();
         int[] scalingStepToBracketMap = mStatsLayout.getScalingStepToPowerBracketMap();
-        int energyConsumerCount = mStatsLayout.getCpuClusterEnergyConsumerCount();
+        int energyConsumerCount = mStatsLayout.getEnergyConsumerCount();
         List<DeviceStateEstimation> deviceStateEstimations = mPlan.deviceStateEstimations;
         for (int dse = deviceStateEstimations.size() - 1; dse >= 0; dse--) {
             DeviceStateEstimation deviceStateEstimation = deviceStateEstimations.get(dse);
@@ -392,7 +392,7 @@
 
     private void adjustEstimatesUsingEnergyConsumers(
             Intermediates intermediates, DeviceStatsIntermediates deviceStatsIntermediates) {
-        int energyConsumerCount = mStatsLayout.getCpuClusterEnergyConsumerCount();
+        int energyConsumerCount = mStatsLayout.getEnergyConsumerCount();
         if (energyConsumerCount == 0) {
             return;
         }
@@ -509,8 +509,8 @@
             }
             sb.append(mStatsLayout.getTimeByCluster(stats, cluster));
         }
-        sb.append("] uptime: ").append(mStatsLayout.getUptime(stats));
-        int energyConsumerCount = mStatsLayout.getCpuClusterEnergyConsumerCount();
+        sb.append("] uptime: ").append(mStatsLayout.getUsageDuration(stats));
+        int energyConsumerCount = mStatsLayout.getEnergyConsumerCount();
         if (energyConsumerCount > 0) {
             sb.append(" energy: [");
             for (int i = 0; i < energyConsumerCount; i++) {
diff --git a/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java
index b8e581f..c05407c 100644
--- a/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java
@@ -22,6 +22,7 @@
 import android.os.BatteryConsumer;
 import android.os.Handler;
 import android.os.PersistableBundle;
+import android.os.Process;
 import android.power.PowerStatsInternal;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -67,6 +68,7 @@
     private final CpuScalingPolicies mCpuScalingPolicies;
     private final PowerProfile mPowerProfile;
     private final KernelCpuStatsReader mKernelCpuStatsReader;
+    private final PowerStatsUidResolver mUidResolver;
     private final Supplier<PowerStatsInternal> mPowerStatsSupplier;
     private final IntSupplier mVoltageSupplier;
     private final int mDefaultCpuPowerBrackets;
@@ -81,7 +83,7 @@
     private PowerStats.Descriptor mPowerStatsDescriptor;
     // Reusable instance
     private PowerStats mCpuPowerStats;
-    private StatsArrayLayout mLayout;
+    private CpuStatsArrayLayout mLayout;
     private long mLastUpdateTimestampNanos;
     private long mLastUpdateUptimeMillis;
     private int mLastVoltageMv;
@@ -91,55 +93,30 @@
      * Captures the positions and lengths of sections of the stats array, such as time-in-state,
      * power usage estimates etc.
      */
-    public static class StatsArrayLayout {
+    public static class CpuStatsArrayLayout extends StatsArrayLayout {
         private static final String EXTRA_DEVICE_TIME_BY_SCALING_STEP_POSITION = "dt";
         private static final String EXTRA_DEVICE_TIME_BY_SCALING_STEP_COUNT = "dtc";
         private static final String EXTRA_DEVICE_TIME_BY_CLUSTER_POSITION = "dc";
         private static final String EXTRA_DEVICE_TIME_BY_CLUSTER_COUNT = "dcc";
-        private static final String EXTRA_DEVICE_POWER_POSITION = "dp";
-        private static final String EXTRA_DEVICE_UPTIME_POSITION = "du";
-        private static final String EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION = "de";
-        private static final String EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT = "dec";
         private static final String EXTRA_UID_BRACKETS_POSITION = "ub";
         private static final String EXTRA_UID_STATS_SCALING_STEP_TO_POWER_BRACKET = "us";
-        private static final String EXTRA_UID_POWER_POSITION = "up";
-
-        private static final double MILLI_TO_NANO_MULTIPLIER = 1000000.0;
-
-        private int mDeviceStatsArrayLength;
-        private int mUidStatsArrayLength;
 
         private int mDeviceCpuTimeByScalingStepPosition;
         private int mDeviceCpuTimeByScalingStepCount;
         private int mDeviceCpuTimeByClusterPosition;
         private int mDeviceCpuTimeByClusterCount;
-        private int mDeviceCpuUptimePosition;
-        private int mDeviceEnergyConsumerPosition;
-        private int mDeviceEnergyConsumerCount;
-        private int mDevicePowerEstimatePosition;
 
         private int mUidPowerBracketsPosition;
         private int mUidPowerBracketCount;
-        private int[][] mEnergyConsumerToPowerBucketMaps;
-        private int mUidPowerEstimatePosition;
 
         private int[] mScalingStepToPowerBracketMap;
 
-        public int getDeviceStatsArrayLength() {
-            return mDeviceStatsArrayLength;
-        }
-
-        public int getUidStatsArrayLength() {
-            return mUidStatsArrayLength;
-        }
-
         /**
          * Declare that the stats array has a section capturing CPU time per scaling step
          */
         public void addDeviceSectionCpuTimeByScalingStep(int scalingStepCount) {
-            mDeviceCpuTimeByScalingStepPosition = mDeviceStatsArrayLength;
+            mDeviceCpuTimeByScalingStepPosition = addDeviceSection(scalingStepCount);
             mDeviceCpuTimeByScalingStepCount = scalingStepCount;
-            mDeviceStatsArrayLength += scalingStepCount;
         }
 
         public int getCpuScalingStepCount() {
@@ -166,9 +143,8 @@
          * Declare that the stats array has a section capturing CPU time in each cluster
          */
         public void addDeviceSectionCpuTimeByCluster(int clusterCount) {
+            mDeviceCpuTimeByClusterPosition = addDeviceSection(clusterCount);
             mDeviceCpuTimeByClusterCount = clusterCount;
-            mDeviceCpuTimeByClusterPosition = mDeviceStatsArrayLength;
-            mDeviceStatsArrayLength += clusterCount;
         }
 
         public int getCpuClusterCount() {
@@ -192,86 +168,12 @@
         }
 
         /**
-         * Declare that the stats array has a section capturing CPU uptime
-         */
-        public void addDeviceSectionUptime() {
-            mDeviceCpuUptimePosition = mDeviceStatsArrayLength++;
-        }
-
-        /**
-         * Saves the CPU uptime duration in the corresponding <code>stats</code> element.
-         */
-        public void setUptime(long[] stats, long value) {
-            stats[mDeviceCpuUptimePosition] = value;
-        }
-
-        /**
-         * Extracts the CPU uptime duration from the corresponding <code>stats</code> element.
-         */
-        public long getUptime(long[] stats) {
-            return stats[mDeviceCpuUptimePosition];
-        }
-
-        /**
-         * Declares that the stats array has a section capturing EnergyConsumer data from
-         * PowerStatsService.
-         */
-        public void addDeviceSectionEnergyConsumers(int energyConsumerCount) {
-            mDeviceEnergyConsumerPosition = mDeviceStatsArrayLength;
-            mDeviceEnergyConsumerCount = energyConsumerCount;
-            mDeviceStatsArrayLength += energyConsumerCount;
-        }
-
-        public int getCpuClusterEnergyConsumerCount() {
-            return mDeviceEnergyConsumerCount;
-        }
-
-        /**
-         * Saves the accumulated energy for the specified rail the corresponding
-         * <code>stats</code> element.
-         */
-        public void setConsumedEnergy(long[] stats, int index, long energy) {
-            stats[mDeviceEnergyConsumerPosition + index] = energy;
-        }
-
-        /**
-         * Extracts the EnergyConsumer data from a device stats array for the specified
-         * EnergyConsumer.
-         */
-        public long getConsumedEnergy(long[] stats, int index) {
-            return stats[mDeviceEnergyConsumerPosition + index];
-        }
-
-        /**
-         * Declare that the stats array has a section capturing a power estimate
-         */
-        public void addDeviceSectionPowerEstimate() {
-            mDevicePowerEstimatePosition = mDeviceStatsArrayLength++;
-        }
-
-        /**
-         * Converts the supplied mAh power estimate to a long and saves it in the corresponding
-         * element of <code>stats</code>.
-         */
-        public void setDevicePowerEstimate(long[] stats, double power) {
-            stats[mDevicePowerEstimatePosition] = (long) (power * MILLI_TO_NANO_MULTIPLIER);
-        }
-
-        /**
-         * Extracts the power estimate from a device stats array and converts it to mAh.
-         */
-        public double getDevicePowerEstimate(long[] stats) {
-            return stats[mDevicePowerEstimatePosition] / MILLI_TO_NANO_MULTIPLIER;
-        }
-
-        /**
          * Declare that the UID stats array has a section capturing CPU time per power bracket.
          */
         public void addUidSectionCpuTimeByPowerBracket(int[] scalingStepToPowerBracketMap) {
             mScalingStepToPowerBracketMap = scalingStepToPowerBracketMap;
-            mUidPowerBracketsPosition = mUidStatsArrayLength;
             updatePowerBracketCount();
-            mUidStatsArrayLength += mUidPowerBracketCount;
+            mUidPowerBracketsPosition = addUidSection(mUidPowerBracketCount);
         }
 
         private void updatePowerBracketCount() {
@@ -306,31 +208,10 @@
         }
 
         /**
-         * Declare that the UID stats array has a section capturing a power estimate
-         */
-        public void addUidSectionPowerEstimate() {
-            mUidPowerEstimatePosition = mUidStatsArrayLength++;
-        }
-
-        /**
-         * Converts the supplied mAh power estimate to a long and saves it in the corresponding
-         * element of <code>stats</code>.
-         */
-        public void setUidPowerEstimate(long[] stats, double power) {
-            stats[mUidPowerEstimatePosition] = (long) (power * MILLI_TO_NANO_MULTIPLIER);
-        }
-
-        /**
-         * Extracts the power estimate from a UID stats array and converts it to mAh.
-         */
-        public double getUidPowerEstimate(long[] stats) {
-            return stats[mUidPowerEstimatePosition] / MILLI_TO_NANO_MULTIPLIER;
-        }
-
-        /**
          * Copies the elements of the stats array layout into <code>extras</code>
          */
         public void toExtras(PersistableBundle extras) {
+            super.toExtras(extras);
             extras.putInt(EXTRA_DEVICE_TIME_BY_SCALING_STEP_POSITION,
                     mDeviceCpuTimeByScalingStepPosition);
             extras.putInt(EXTRA_DEVICE_TIME_BY_SCALING_STEP_COUNT,
@@ -339,22 +220,16 @@
                     mDeviceCpuTimeByClusterPosition);
             extras.putInt(EXTRA_DEVICE_TIME_BY_CLUSTER_COUNT,
                     mDeviceCpuTimeByClusterCount);
-            extras.putInt(EXTRA_DEVICE_UPTIME_POSITION, mDeviceCpuUptimePosition);
-            extras.putInt(EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION,
-                    mDeviceEnergyConsumerPosition);
-            extras.putInt(EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT,
-                    mDeviceEnergyConsumerCount);
-            extras.putInt(EXTRA_DEVICE_POWER_POSITION, mDevicePowerEstimatePosition);
             extras.putInt(EXTRA_UID_BRACKETS_POSITION, mUidPowerBracketsPosition);
-            extras.putIntArray(EXTRA_UID_STATS_SCALING_STEP_TO_POWER_BRACKET,
+            putIntArray(extras, EXTRA_UID_STATS_SCALING_STEP_TO_POWER_BRACKET,
                     mScalingStepToPowerBracketMap);
-            extras.putInt(EXTRA_UID_POWER_POSITION, mUidPowerEstimatePosition);
         }
 
         /**
          * Retrieves elements of the stats array layout from <code>extras</code>
          */
         public void fromExtras(PersistableBundle extras) {
+            super.fromExtras(extras);
             mDeviceCpuTimeByScalingStepPosition =
                     extras.getInt(EXTRA_DEVICE_TIME_BY_SCALING_STEP_POSITION);
             mDeviceCpuTimeByScalingStepCount =
@@ -363,43 +238,39 @@
                     extras.getInt(EXTRA_DEVICE_TIME_BY_CLUSTER_POSITION);
             mDeviceCpuTimeByClusterCount =
                     extras.getInt(EXTRA_DEVICE_TIME_BY_CLUSTER_COUNT);
-            mDeviceCpuUptimePosition = extras.getInt(EXTRA_DEVICE_UPTIME_POSITION);
-            mDeviceEnergyConsumerPosition = extras.getInt(EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION);
-            mDeviceEnergyConsumerCount = extras.getInt(EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT);
-            mDevicePowerEstimatePosition = extras.getInt(EXTRA_DEVICE_POWER_POSITION);
             mUidPowerBracketsPosition = extras.getInt(EXTRA_UID_BRACKETS_POSITION);
             mScalingStepToPowerBracketMap =
-                    extras.getIntArray(EXTRA_UID_STATS_SCALING_STEP_TO_POWER_BRACKET);
+                    getIntArray(extras, EXTRA_UID_STATS_SCALING_STEP_TO_POWER_BRACKET);
             if (mScalingStepToPowerBracketMap == null) {
                 mScalingStepToPowerBracketMap = new int[mDeviceCpuTimeByScalingStepCount];
             }
             updatePowerBracketCount();
-            mUidPowerEstimatePosition = extras.getInt(EXTRA_UID_POWER_POSITION);
         }
     }
 
     public CpuPowerStatsCollector(CpuScalingPolicies cpuScalingPolicies, PowerProfile powerProfile,
-            IntSupplier voltageSupplier, Handler handler, long throttlePeriodMs) {
-        this(cpuScalingPolicies, powerProfile, handler, new KernelCpuStatsReader(),
+            PowerStatsUidResolver uidResolver, IntSupplier voltageSupplier, Handler handler,
+            long throttlePeriodMs) {
+        this(cpuScalingPolicies, powerProfile, handler, new KernelCpuStatsReader(), uidResolver,
                 () -> LocalServices.getService(PowerStatsInternal.class), voltageSupplier,
                 throttlePeriodMs, Clock.SYSTEM_CLOCK, DEFAULT_CPU_POWER_BRACKETS,
                 DEFAULT_CPU_POWER_BRACKETS_PER_ENERGY_CONSUMER);
     }
 
     public CpuPowerStatsCollector(CpuScalingPolicies cpuScalingPolicies, PowerProfile powerProfile,
-                                  Handler handler, KernelCpuStatsReader kernelCpuStatsReader,
-            Supplier<PowerStatsInternal> powerStatsSupplier,
+            Handler handler, KernelCpuStatsReader kernelCpuStatsReader,
+            PowerStatsUidResolver uidResolver, Supplier<PowerStatsInternal> powerStatsSupplier,
             IntSupplier voltageSupplier, long throttlePeriodMs, Clock clock,
             int defaultCpuPowerBrackets, int defaultCpuPowerBracketsPerEnergyConsumer) {
         super(handler, throttlePeriodMs, clock);
         mCpuScalingPolicies = cpuScalingPolicies;
         mPowerProfile = powerProfile;
         mKernelCpuStatsReader = kernelCpuStatsReader;
+        mUidResolver = uidResolver;
         mPowerStatsSupplier = powerStatsSupplier;
         mVoltageSupplier = voltageSupplier;
         mDefaultCpuPowerBrackets = defaultCpuPowerBrackets;
         mDefaultCpuPowerBracketsPerEnergyConsumer = defaultCpuPowerBracketsPerEnergyConsumer;
-
     }
 
     /**
@@ -409,13 +280,13 @@
         setEnabled(Flags.streamlinedBatteryStats());
     }
 
-    private void ensureInitialized() {
+    private boolean ensureInitialized() {
         if (mIsInitialized) {
-            return;
+            return true;
         }
 
         if (!isEnabled()) {
-            return;
+            return false;
         }
 
         mIsPerUidTimeInStateSupported = mKernelCpuStatsReader.nativeIsSupportedFeature();
@@ -432,10 +303,10 @@
         mTempCpuTimeByScalingStep = new long[cpuScalingStepCount];
         int[] scalingStepToPowerBracketMap = initPowerBrackets();
 
-        mLayout = new StatsArrayLayout();
+        mLayout = new CpuStatsArrayLayout();
         mLayout.addDeviceSectionCpuTimeByScalingStep(cpuScalingStepCount);
         mLayout.addDeviceSectionCpuTimeByCluster(mCpuScalingPolicies.getPolicies().length);
-        mLayout.addDeviceSectionUptime();
+        mLayout.addDeviceSectionUsageDuration();
         mLayout.addDeviceSectionEnergyConsumers(mCpuEnergyConsumerIds.length);
         mLayout.addDeviceSectionPowerEstimate();
         mLayout.addUidSectionCpuTimeByPowerBracket(scalingStepToPowerBracketMap);
@@ -451,6 +322,7 @@
         mTempUidStats = new long[mLayout.getCpuPowerBracketCount()];
 
         mIsInitialized = true;
+        return true;
     }
 
     private void readCpuEnergyConsumerIds() {
@@ -590,7 +462,9 @@
      * Prints the definitions of power brackets.
      */
     public void dumpCpuPowerBracketsLocked(PrintWriter pw) {
-        ensureInitialized();
+        if (!ensureInitialized()) {
+            return;
+        }
 
         if (mLayout == null) {
             return;
@@ -610,7 +484,9 @@
      */
     @VisibleForTesting
     public String getCpuPowerBracketDescription(int powerBracket) {
-        ensureInitialized();
+        if (!ensureInitialized()) {
+            return "";
+        }
 
         int[] stepToPowerBracketMap = mLayout.getScalingStepToPowerBracketMap();
         StringBuilder sb = new StringBuilder();
@@ -647,14 +523,18 @@
      */
     @VisibleForTesting
     public PowerStats.Descriptor getPowerStatsDescriptor() {
-        ensureInitialized();
+        if (!ensureInitialized()) {
+            return null;
+        }
 
         return mPowerStatsDescriptor;
     }
 
     @Override
     protected PowerStats collectStats() {
-        ensureInitialized();
+        if (!ensureInitialized()) {
+            return null;
+        }
 
         if (!mIsPerUidTimeInStateSupported) {
             return null;
@@ -682,7 +562,7 @@
         if (uptimeDelta > mCpuPowerStats.durationMs) {
             uptimeDelta = mCpuPowerStats.durationMs;
         }
-        mLayout.setUptime(mCpuPowerStats.stats, uptimeDelta);
+        mLayout.setUsageDuration(mCpuPowerStats.stats, uptimeDelta);
 
         if (mCpuEnergyConsumerIds.length != 0) {
             collectEnergyConsumers();
@@ -761,7 +641,21 @@
             uidStats.timeByPowerBracket[bracket] = timeByPowerBracket[bracket];
         }
         if (nonzero) {
-            mCpuPowerStats.uidStats.put(uid, uidStats.stats);
+            int ownerUid;
+            if (Process.isSdkSandboxUid(uid)) {
+                ownerUid = Process.getAppUidForSdkSandboxUid(uid);
+            } else {
+                ownerUid = mUidResolver.mapUid(uid);
+            }
+
+            long[] ownerStats = mCpuPowerStats.uidStats.get(ownerUid);
+            if (ownerStats == null) {
+                mCpuPowerStats.uidStats.put(ownerUid, uidStats.stats);
+            } else {
+                for (int i = 0; i < ownerStats.length; i++) {
+                    ownerStats[i] += uidStats.stats[i];
+                }
+            }
         }
     }
 
diff --git a/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java b/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
index 2c7843e..0facb9c 100644
--- a/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
+++ b/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
@@ -44,6 +44,7 @@
     private static final String XML_TAG_DEVICE_STATS = "device-stats";
     private static final String XML_TAG_UID_STATS = "uid-stats";
     private static final String XML_ATTR_UID = "uid";
+    private static final long UNKNOWN = -1;
 
     public final int powerComponentId;
     private final MultiStateStats.States[] mDeviceStateConfig;
@@ -51,17 +52,16 @@
     @NonNull
     private final AggregatedPowerStatsConfig.PowerComponent mConfig;
     private final int[] mDeviceStates;
-    private final long[] mDeviceStateTimestamps;
 
     private MultiStateStats.Factory mStatsFactory;
     private MultiStateStats.Factory mUidStatsFactory;
     private PowerStats.Descriptor mPowerStatsDescriptor;
+    private long mPowerStatsTimestamp;
     private MultiStateStats mDeviceStats;
     private final SparseArray<UidStats> mUidStats = new SparseArray<>();
 
     private static class UidStats {
         public int[] states;
-        public long[] stateTimestampMs;
         public MultiStateStats stats;
     }
 
@@ -71,7 +71,7 @@
         mDeviceStateConfig = config.getDeviceStateConfig();
         mUidStateConfig = config.getUidStateConfig();
         mDeviceStates = new int[mDeviceStateConfig.length];
-        mDeviceStateTimestamps = new long[mDeviceStateConfig.length];
+        mPowerStatsTimestamp = UNKNOWN;
     }
 
     @NonNull
@@ -85,8 +85,11 @@
     }
 
     void setState(@AggregatedPowerStatsConfig.TrackedState int stateId, int state, long time) {
+        if (mDeviceStats == null) {
+            createDeviceStats();
+        }
+
         mDeviceStates[stateId] = state;
-        mDeviceStateTimestamps[stateId] = time;
 
         if (mDeviceStateConfig[stateId].isTracked()) {
             if (mDeviceStats != null) {
@@ -97,6 +100,11 @@
         if (mUidStateConfig[stateId].isTracked()) {
             for (int i = mUidStats.size() - 1; i >= 0; i--) {
                 PowerComponentAggregatedPowerStats.UidStats uidStats = mUidStats.valueAt(i);
+                if (uidStats.stats == null) {
+                    createUidStats(uidStats);
+                }
+
+                uidStats.states[stateId] = state;
                 if (uidStats.stats != null) {
                     uidStats.stats.setState(stateId, state, time);
                 }
@@ -111,8 +119,11 @@
         }
 
         UidStats uidStats = getUidStats(uid);
+        if (uidStats.stats == null) {
+            createUidStats(uidStats);
+        }
+
         uidStats.states[stateId] = state;
-        uidStats.stateTimestampMs[stateId] = time;
 
         if (uidStats.stats != null) {
             uidStats.stats.setState(stateId, state, time);
@@ -150,10 +161,11 @@
             }
             uidStats.stats.increment(powerStats.uidStats.valueAt(i), timestampMs);
         }
+
+        mPowerStatsTimestamp = timestampMs;
     }
 
     void reset() {
-        mPowerStatsDescriptor = null;
         mStatsFactory = null;
         mUidStatsFactory = null;
         mDeviceStats = null;
@@ -163,12 +175,10 @@
     }
 
     private UidStats getUidStats(int uid) {
-        // TODO(b/292247660): map isolated and sandbox UIDs
         UidStats uidStats = mUidStats.get(uid);
         if (uidStats == null) {
             uidStats = new UidStats();
             uidStats.states = new int[mUidStateConfig.length];
-            uidStats.stateTimestampMs = new long[mUidStateConfig.length];
             mUidStats.put(uid, uidStats);
         }
         return uidStats;
@@ -209,42 +219,38 @@
         return false;
     }
 
-    private boolean createDeviceStats() {
+    private void createDeviceStats() {
         if (mStatsFactory == null) {
             if (mPowerStatsDescriptor == null) {
-                return false;
+                return;
             }
             mStatsFactory = new MultiStateStats.Factory(
                     mPowerStatsDescriptor.statsArrayLength, mDeviceStateConfig);
         }
 
         mDeviceStats = mStatsFactory.create();
-        for (int stateId = 0; stateId < mDeviceStateConfig.length; stateId++) {
-            mDeviceStats.setState(stateId, mDeviceStates[stateId],
-                    mDeviceStateTimestamps[stateId]);
+        if (mPowerStatsTimestamp != UNKNOWN) {
+            for (int stateId = 0; stateId < mDeviceStateConfig.length; stateId++) {
+                mDeviceStats.setState(stateId, mDeviceStates[stateId], mPowerStatsTimestamp);
+            }
         }
-        return true;
     }
 
-    private boolean createUidStats(UidStats uidStats) {
+    private void createUidStats(UidStats uidStats) {
         if (mUidStatsFactory == null) {
             if (mPowerStatsDescriptor == null) {
-                return false;
+                return;
             }
             mUidStatsFactory = new MultiStateStats.Factory(
                     mPowerStatsDescriptor.uidStatsArrayLength, mUidStateConfig);
         }
 
         uidStats.stats = mUidStatsFactory.create();
-        for (int stateId = 0; stateId < mDeviceStateConfig.length; stateId++) {
-            uidStats.stats.setState(stateId, mDeviceStates[stateId],
-                    mDeviceStateTimestamps[stateId]);
+        for (int stateId = 0; stateId < mUidStateConfig.length; stateId++) {
+            if (mPowerStatsTimestamp != UNKNOWN) {
+                uidStats.stats.setState(stateId, uidStats.states[stateId], mPowerStatsTimestamp);
+            }
         }
-        for (int stateId = mDeviceStateConfig.length; stateId < mUidStateConfig.length; stateId++) {
-            uidStats.stats.setState(stateId, uidStats.states[stateId],
-                    uidStats.stateTimestampMs[stateId]);
-        }
-        return true;
     }
 
     public void writeXml(TypedXmlSerializer serializer) throws IOException {
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java b/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java
index 2f9d567..3f88a2d 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java
@@ -29,13 +29,18 @@
  * {@link AggregatedPowerStats} that adds up power stats from the samples found in battery history.
  */
 public class PowerStatsAggregator {
+    private static final long UNINITIALIZED = -1;
     private final AggregatedPowerStats mStats;
+    private final AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
     private final BatteryStatsHistory mHistory;
     private final SparseArray<AggregatedPowerStatsProcessor> mProcessors = new SparseArray<>();
+    private int mCurrentBatteryState = AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
+    private int mCurrentScreenState = AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
 
     public PowerStatsAggregator(AggregatedPowerStatsConfig aggregatedPowerStatsConfig,
             BatteryStatsHistory history) {
         mStats = new AggregatedPowerStats(aggregatedPowerStatsConfig);
+        mAggregatedPowerStatsConfig = aggregatedPowerStatsConfig;
         mHistory = history;
         for (AggregatedPowerStatsConfig.PowerComponent powerComponentsConfig :
                 aggregatedPowerStatsConfig.getPowerComponentsAggregatedStatsConfigs()) {
@@ -44,6 +49,10 @@
         }
     }
 
+    AggregatedPowerStatsConfig getConfig() {
+        return mAggregatedPowerStatsConfig;
+    }
+
     /**
      * Iterates of the battery history and aggregates power stats between the specified times.
      * The start and end are specified in the battery-stats monotonic time, which is the
@@ -58,18 +67,20 @@
      */
     public void aggregatePowerStats(long startTimeMs, long endTimeMs,
             Consumer<AggregatedPowerStats> consumer) {
-        int currentBatteryState = AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
-        int currentScreenState = AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
-        long baseTime = -1;
+        boolean clockUpdateAdded = false;
+        long baseTime = startTimeMs > 0 ? startTimeMs : UNINITIALIZED;
         long lastTime = 0;
         try (BatteryStatsHistoryIterator iterator =
                      mHistory.copy().iterate(startTimeMs, endTimeMs)) {
             while (iterator.hasNext()) {
                 BatteryStats.HistoryItem item = iterator.next();
 
-                if (baseTime < 0) {
+                if (!clockUpdateAdded) {
                     mStats.addClockUpdate(item.time, item.currentTime);
-                    baseTime = item.time;
+                    if (baseTime == UNINITIALIZED) {
+                        baseTime = item.time;
+                    }
+                    clockUpdateAdded = true;
                 } else if (item.cmd == BatteryStats.HistoryItem.CMD_CURRENT_TIME
                            || item.cmd == BatteryStats.HistoryItem.CMD_RESET) {
                     mStats.addClockUpdate(item.time, item.currentTime);
@@ -81,20 +92,20 @@
                         (item.states & BatteryStats.HistoryItem.STATE_BATTERY_PLUGGED_FLAG) != 0
                                 ? AggregatedPowerStatsConfig.POWER_STATE_OTHER
                                 : AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
-                if (batteryState != currentBatteryState) {
+                if (batteryState != mCurrentBatteryState) {
                     mStats.setDeviceState(AggregatedPowerStatsConfig.STATE_POWER, batteryState,
                             item.time);
-                    currentBatteryState = batteryState;
+                    mCurrentBatteryState = batteryState;
                 }
 
                 int screenState =
                         (item.states & BatteryStats.HistoryItem.STATE_SCREEN_ON_FLAG) != 0
                                 ? AggregatedPowerStatsConfig.SCREEN_STATE_ON
                                 : AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
-                if (screenState != currentScreenState) {
+                if (screenState != mCurrentScreenState) {
                     mStats.setDeviceState(AggregatedPowerStatsConfig.STATE_SCREEN, screenState,
                             item.time);
-                    currentScreenState = screenState;
+                    mCurrentScreenState = screenState;
                 }
 
                 if (item.processStateChange != null) {
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsCollector.java b/services/core/java/com/android/server/power/stats/PowerStatsCollector.java
index 84cc21e..abfe9de 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsCollector.java
@@ -16,10 +16,13 @@
 
 package com.android.server.power.stats;
 
+import android.annotation.Nullable;
 import android.os.ConditionVariable;
 import android.os.Handler;
+import android.os.PersistableBundle;
 import android.util.FastImmutableArraySet;
 import android.util.IndentingPrintWriter;
+import android.util.Slog;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.os.Clock;
@@ -38,6 +41,7 @@
  * except where noted.
  */
 public abstract class PowerStatsCollector {
+    private static final String TAG = "PowerStatsCollector";
     private static final int MILLIVOLTS_PER_VOLT = 1000;
     private final Handler mHandler;
     protected final Clock mClock;
@@ -46,6 +50,200 @@
     private boolean mEnabled;
     private long mLastScheduledUpdateMs = -1;
 
+    /**
+     * Captures the positions and lengths of sections of the stats array, such as usage duration,
+     * power usage estimates etc.
+     */
+    public static class StatsArrayLayout {
+        private static final String EXTRA_DEVICE_POWER_POSITION = "dp";
+        private static final String EXTRA_DEVICE_DURATION_POSITION = "dd";
+        private static final String EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION = "de";
+        private static final String EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT = "dec";
+        private static final String EXTRA_UID_POWER_POSITION = "up";
+
+        protected static final double MILLI_TO_NANO_MULTIPLIER = 1000000.0;
+
+        private int mDeviceStatsArrayLength;
+        private int mUidStatsArrayLength;
+
+        protected int mDeviceDurationPosition;
+        private int mDeviceEnergyConsumerPosition;
+        private int mDeviceEnergyConsumerCount;
+        private int mDevicePowerEstimatePosition;
+        private int mUidPowerEstimatePosition;
+
+        public int getDeviceStatsArrayLength() {
+            return mDeviceStatsArrayLength;
+        }
+
+        public int getUidStatsArrayLength() {
+            return mUidStatsArrayLength;
+        }
+
+        protected int addDeviceSection(int length) {
+            int position = mDeviceStatsArrayLength;
+            mDeviceStatsArrayLength += length;
+            return position;
+        }
+
+        protected int addUidSection(int length) {
+            int position = mUidStatsArrayLength;
+            mUidStatsArrayLength += length;
+            return position;
+        }
+
+        /**
+         * Declare that the stats array has a section capturing usage duration
+         */
+        public void addDeviceSectionUsageDuration() {
+            mDeviceDurationPosition = addDeviceSection(1);
+        }
+
+        /**
+         * Saves the usage duration in the corresponding <code>stats</code> element.
+         */
+        public void setUsageDuration(long[] stats, long value) {
+            stats[mDeviceDurationPosition] = value;
+        }
+
+        /**
+         * Extracts the usage duration from the corresponding <code>stats</code> element.
+         */
+        public long getUsageDuration(long[] stats) {
+            return stats[mDeviceDurationPosition];
+        }
+
+        /**
+         * Declares that the stats array has a section capturing EnergyConsumer data from
+         * PowerStatsService.
+         */
+        public void addDeviceSectionEnergyConsumers(int energyConsumerCount) {
+            mDeviceEnergyConsumerPosition = addDeviceSection(energyConsumerCount);
+            mDeviceEnergyConsumerCount = energyConsumerCount;
+        }
+
+        public int getEnergyConsumerCount() {
+            return mDeviceEnergyConsumerCount;
+        }
+
+        /**
+         * Saves the accumulated energy for the specified rail the corresponding
+         * <code>stats</code> element.
+         */
+        public void setConsumedEnergy(long[] stats, int index, long energy) {
+            stats[mDeviceEnergyConsumerPosition + index] = energy;
+        }
+
+        /**
+         * Extracts the EnergyConsumer data from a device stats array for the specified
+         * EnergyConsumer.
+         */
+        public long getConsumedEnergy(long[] stats, int index) {
+            return stats[mDeviceEnergyConsumerPosition + index];
+        }
+
+        /**
+         * Declare that the stats array has a section capturing a power estimate
+         */
+        public void addDeviceSectionPowerEstimate() {
+            mDevicePowerEstimatePosition = addDeviceSection(1);
+        }
+
+        /**
+         * Converts the supplied mAh power estimate to a long and saves it in the corresponding
+         * element of <code>stats</code>.
+         */
+        public void setDevicePowerEstimate(long[] stats, double power) {
+            stats[mDevicePowerEstimatePosition] = (long) (power * MILLI_TO_NANO_MULTIPLIER);
+        }
+
+        /**
+         * Extracts the power estimate from a device stats array and converts it to mAh.
+         */
+        public double getDevicePowerEstimate(long[] stats) {
+            return stats[mDevicePowerEstimatePosition] / MILLI_TO_NANO_MULTIPLIER;
+        }
+
+        /**
+         * Declare that the UID stats array has a section capturing a power estimate
+         */
+        public void addUidSectionPowerEstimate() {
+            mUidPowerEstimatePosition = addUidSection(1);
+        }
+
+        /**
+         * Converts the supplied mAh power estimate to a long and saves it in the corresponding
+         * element of <code>stats</code>.
+         */
+        public void setUidPowerEstimate(long[] stats, double power) {
+            stats[mUidPowerEstimatePosition] = (long) (power * MILLI_TO_NANO_MULTIPLIER);
+        }
+
+        /**
+         * Extracts the power estimate from a UID stats array and converts it to mAh.
+         */
+        public double getUidPowerEstimate(long[] stats) {
+            return stats[mUidPowerEstimatePosition] / MILLI_TO_NANO_MULTIPLIER;
+        }
+
+        /**
+         * Copies the elements of the stats array layout into <code>extras</code>
+         */
+        public void toExtras(PersistableBundle extras) {
+            extras.putInt(EXTRA_DEVICE_DURATION_POSITION, mDeviceDurationPosition);
+            extras.putInt(EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION,
+                    mDeviceEnergyConsumerPosition);
+            extras.putInt(EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT,
+                    mDeviceEnergyConsumerCount);
+            extras.putInt(EXTRA_DEVICE_POWER_POSITION, mDevicePowerEstimatePosition);
+            extras.putInt(EXTRA_UID_POWER_POSITION, mUidPowerEstimatePosition);
+        }
+
+        /**
+         * Retrieves elements of the stats array layout from <code>extras</code>
+         */
+        public void fromExtras(PersistableBundle extras) {
+            mDeviceDurationPosition = extras.getInt(EXTRA_DEVICE_DURATION_POSITION);
+            mDeviceEnergyConsumerPosition = extras.getInt(EXTRA_DEVICE_ENERGY_CONSUMERS_POSITION);
+            mDeviceEnergyConsumerCount = extras.getInt(EXTRA_DEVICE_ENERGY_CONSUMERS_COUNT);
+            mDevicePowerEstimatePosition = extras.getInt(EXTRA_DEVICE_POWER_POSITION);
+            mUidPowerEstimatePosition = extras.getInt(EXTRA_UID_POWER_POSITION);
+        }
+
+        protected void putIntArray(PersistableBundle extras, String key, int[] array) {
+            if (array == null) {
+                return;
+            }
+
+            StringBuilder sb = new StringBuilder();
+            for (int value : array) {
+                if (!sb.isEmpty()) {
+                    sb.append(',');
+                }
+                sb.append(value);
+            }
+            extras.putString(key, sb.toString());
+        }
+
+        protected int[] getIntArray(PersistableBundle extras, String key) {
+            String string = extras.getString(key);
+            if (string == null) {
+                return null;
+            }
+            String[] values = string.trim().split(",");
+            int[] result = new int[values.length];
+            for (int i = 0; i < values.length; i++) {
+                try {
+                    result[i] = Integer.parseInt(values[i]);
+                } catch (NumberFormatException e) {
+                    Slog.wtf(TAG, "Invalid CSV format: " + string);
+                    return null;
+                }
+            }
+            return result;
+        }
+    }
+
     @GuardedBy("this")
     @SuppressWarnings("unchecked")
     private volatile FastImmutableArraySet<Consumer<PowerStats>> mConsumerList =
@@ -141,6 +339,7 @@
         return true;
     }
 
+    @Nullable
     protected abstract PowerStats collectStats();
 
     /**
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsExporter.java b/services/core/java/com/android/server/power/stats/PowerStatsExporter.java
new file mode 100644
index 0000000..70c24d5
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/PowerStatsExporter.java
@@ -0,0 +1,227 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import android.os.AggregateBatteryConsumer;
+import android.os.BatteryConsumer;
+import android.os.BatteryUsageStats;
+import android.os.UidBatteryConsumer;
+import android.util.Slog;
+
+import com.android.internal.os.MultiStateStats;
+import com.android.internal.os.PowerStats;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Given a time range, converts accumulated PowerStats to BatteryUsageStats.  Combines
+ * stores spans of PowerStats with the yet-unprocessed tail of battery history.
+ */
+public class PowerStatsExporter {
+    private static final String TAG = "PowerStatsExporter";
+    private final PowerStatsStore mPowerStatsStore;
+    private final PowerStatsAggregator mPowerStatsAggregator;
+    private final long mBatterySessionTimeSpanSlackMillis;
+    private static final long BATTERY_SESSION_TIME_SPAN_SLACK_MILLIS = TimeUnit.MINUTES.toMillis(2);
+
+    public PowerStatsExporter(PowerStatsStore powerStatsStore,
+            PowerStatsAggregator powerStatsAggregator) {
+        this(powerStatsStore, powerStatsAggregator, BATTERY_SESSION_TIME_SPAN_SLACK_MILLIS);
+    }
+
+    public PowerStatsExporter(PowerStatsStore powerStatsStore,
+            PowerStatsAggregator powerStatsAggregator,
+            long batterySessionTimeSpanSlackMillis) {
+        mPowerStatsStore = powerStatsStore;
+        mPowerStatsAggregator = powerStatsAggregator;
+        mBatterySessionTimeSpanSlackMillis = batterySessionTimeSpanSlackMillis;
+    }
+
+    /**
+     * Populates the provided BatteryUsageStats.Builder with power estimates from the accumulated
+     * PowerStats, both stored in PowerStatsStore and not-yet processed.
+     */
+    public void exportAggregatedPowerStats(BatteryUsageStats.Builder batteryUsageStatsBuilder,
+            long monotonicStartTime, long monotonicEndTime) {
+        long maxEndTime = monotonicStartTime;
+        List<PowerStatsSpan.Metadata> spans = mPowerStatsStore.getTableOfContents();
+        for (int i = spans.size() - 1; i >= 0; i--) {
+            PowerStatsSpan.Metadata metadata = spans.get(i);
+            if (!metadata.getSections().contains(AggregatedPowerStatsSection.TYPE)) {
+                continue;
+            }
+
+            List<PowerStatsSpan.TimeFrame> timeFrames = metadata.getTimeFrames();
+            long spanMinTime = Long.MAX_VALUE;
+            long spanMaxTime = Long.MIN_VALUE;
+            for (int j = 0; j < timeFrames.size(); j++) {
+                PowerStatsSpan.TimeFrame timeFrame = timeFrames.get(j);
+                long startMonotonicTime = timeFrame.startMonotonicTime;
+                long endMonotonicTime = startMonotonicTime + timeFrame.duration;
+                if (startMonotonicTime < spanMinTime) {
+                    spanMinTime = startMonotonicTime;
+                }
+                if (endMonotonicTime > spanMaxTime) {
+                    spanMaxTime = endMonotonicTime;
+                }
+            }
+
+            if (!(spanMinTime >= monotonicStartTime && spanMaxTime < monotonicEndTime)) {
+                continue;
+            }
+
+            if (spanMaxTime > maxEndTime) {
+                maxEndTime = spanMaxTime;
+            }
+
+            PowerStatsSpan span = mPowerStatsStore.loadPowerStatsSpan(metadata.getId(),
+                    AggregatedPowerStatsSection.TYPE);
+            if (span == null) {
+                Slog.e(TAG, "Could not read PowerStatsStore section " + metadata);
+                continue;
+            }
+            List<PowerStatsSpan.Section> sections = span.getSections();
+            for (int k = 0; k < sections.size(); k++) {
+                PowerStatsSpan.Section section = sections.get(k);
+                populateBatteryUsageStatsBuilder(batteryUsageStatsBuilder,
+                        ((AggregatedPowerStatsSection) section).getAggregatedPowerStats());
+            }
+        }
+
+        if (maxEndTime < monotonicEndTime - mBatterySessionTimeSpanSlackMillis) {
+            mPowerStatsAggregator.aggregatePowerStats(maxEndTime, monotonicEndTime,
+                    stats -> populateBatteryUsageStatsBuilder(batteryUsageStatsBuilder, stats));
+        }
+    }
+
+    private void populateBatteryUsageStatsBuilder(
+            BatteryUsageStats.Builder batteryUsageStatsBuilder, AggregatedPowerStats stats) {
+        AggregatedPowerStatsConfig config = mPowerStatsAggregator.getConfig();
+        List<AggregatedPowerStatsConfig.PowerComponent> powerComponents =
+                config.getPowerComponentsAggregatedStatsConfigs();
+        for (int i = powerComponents.size() - 1; i >= 0; i--) {
+            populateBatteryUsageStatsBuilder(batteryUsageStatsBuilder, stats,
+                    powerComponents.get(i));
+        }
+    }
+
+    private void populateBatteryUsageStatsBuilder(
+            BatteryUsageStats.Builder batteryUsageStatsBuilder, AggregatedPowerStats stats,
+            AggregatedPowerStatsConfig.PowerComponent powerComponent) {
+        int powerComponentId = powerComponent.getPowerComponentId();
+        PowerComponentAggregatedPowerStats powerComponentStats = stats.getPowerComponentStats(
+                powerComponentId);
+        if (powerComponentStats == null) {
+            return;
+        }
+
+        PowerStats.Descriptor descriptor = powerComponentStats.getPowerStatsDescriptor();
+        if (descriptor == null) {
+            return;
+        }
+
+        PowerStatsCollector.StatsArrayLayout layout = new PowerStatsCollector.StatsArrayLayout();
+        layout.fromExtras(descriptor.extras);
+
+        long[] deviceStats = new long[descriptor.statsArrayLength];
+        double[] totalPower = new double[1];
+        MultiStateStats.States.forEachTrackedStateCombination(powerComponent.getDeviceStateConfig(),
+                states -> {
+                    if (states[AggregatedPowerStatsConfig.STATE_POWER]
+                            != AggregatedPowerStatsConfig.POWER_STATE_BATTERY) {
+                        return;
+                    }
+
+                    if (!powerComponentStats.getDeviceStats(deviceStats, states)) {
+                        return;
+                    }
+
+                    totalPower[0] += layout.getDevicePowerEstimate(deviceStats);
+                });
+
+        AggregateBatteryConsumer.Builder deviceScope =
+                batteryUsageStatsBuilder.getAggregateBatteryConsumerBuilder(
+                        BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE);
+        deviceScope.addConsumedPower(powerComponentId,
+                totalPower[0], BatteryConsumer.POWER_MODEL_UNDEFINED);
+
+        long[] uidStats = new long[descriptor.uidStatsArrayLength];
+        ArrayList<Integer> uids = new ArrayList<>();
+        powerComponentStats.collectUids(uids);
+
+        boolean breakDownByProcState =
+                batteryUsageStatsBuilder.isProcessStateDataNeeded()
+                && powerComponent
+                        .getUidStateConfig()[AggregatedPowerStatsConfig.STATE_PROCESS_STATE]
+                        .isTracked();
+
+        double[] powerByProcState =
+                new double[breakDownByProcState ? BatteryConsumer.PROCESS_STATE_COUNT : 1];
+        double powerAllApps = 0;
+        for (int uid : uids) {
+            UidBatteryConsumer.Builder builder =
+                    batteryUsageStatsBuilder.getOrCreateUidBatteryConsumerBuilder(uid);
+
+            Arrays.fill(powerByProcState, 0);
+
+            MultiStateStats.States.forEachTrackedStateCombination(
+                    powerComponent.getUidStateConfig(),
+                    states -> {
+                        if (states[AggregatedPowerStatsConfig.STATE_POWER]
+                                != AggregatedPowerStatsConfig.POWER_STATE_BATTERY) {
+                            return;
+                        }
+
+                        if (!powerComponentStats.getUidStats(uidStats, uid, states)) {
+                            return;
+                        }
+
+                        double power = layout.getUidPowerEstimate(uidStats);
+                        int procState = breakDownByProcState
+                                ? states[AggregatedPowerStatsConfig.STATE_PROCESS_STATE]
+                                : BatteryConsumer.PROCESS_STATE_UNSPECIFIED;
+                        powerByProcState[procState] += power;
+                    });
+
+            double powerAllProcStates = 0;
+            for (int procState = 0; procState < powerByProcState.length; procState++) {
+                double power = powerByProcState[procState];
+                if (power == 0) {
+                    continue;
+                }
+                powerAllProcStates += power;
+                if (breakDownByProcState
+                        && procState != BatteryConsumer.PROCESS_STATE_UNSPECIFIED) {
+                    builder.addConsumedPower(builder.getKey(powerComponentId, procState),
+                            power, BatteryConsumer.POWER_MODEL_UNDEFINED);
+                }
+            }
+            builder.addConsumedPower(powerComponentId, powerAllProcStates,
+                    BatteryConsumer.POWER_MODEL_UNDEFINED);
+            powerAllApps += powerAllProcStates;
+        }
+
+        AggregateBatteryConsumer.Builder allAppsScope =
+                batteryUsageStatsBuilder.getAggregateBatteryConsumerBuilder(
+                        BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS);
+        allAppsScope.addConsumedPower(powerComponentId, powerAllApps,
+                BatteryConsumer.POWER_MODEL_UNDEFINED);
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsScheduler.java b/services/core/java/com/android/server/power/stats/PowerStatsScheduler.java
index 551302e..97d872a 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsScheduler.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsScheduler.java
@@ -19,8 +19,6 @@
 import android.annotation.DurationMillisLong;
 import android.app.AlarmManager;
 import android.content.Context;
-import android.os.BatteryUsageStats;
-import android.os.BatteryUsageStatsQuery;
 import android.os.ConditionVariable;
 import android.os.Handler;
 import android.util.IndentingPrintWriter;
@@ -52,7 +50,6 @@
     private final MonotonicClock mMonotonicClock;
     private final Handler mHandler;
     private final BatteryStatsImpl mBatteryStats;
-    private final BatteryUsageStatsProvider mBatteryUsageStatsProvider;
     private final PowerStatsAggregator mPowerStatsAggregator;
     private long mLastSavedSpanEndMonotonicTime;
 
@@ -60,7 +57,7 @@
             @DurationMillisLong long aggregatedPowerStatsSpanDuration,
             @DurationMillisLong long powerStatsAggregationPeriod, PowerStatsStore powerStatsStore,
             Clock clock, MonotonicClock monotonicClock, Handler handler,
-            BatteryStatsImpl batteryStats, BatteryUsageStatsProvider batteryUsageStatsProvider) {
+            BatteryStatsImpl batteryStats) {
         mContext = context;
         mPowerStatsAggregator = powerStatsAggregator;
         mAggregatedPowerStatsSpanDuration = aggregatedPowerStatsSpanDuration;
@@ -70,16 +67,15 @@
         mMonotonicClock = monotonicClock;
         mHandler = handler;
         mBatteryStats = batteryStats;
-        mBatteryUsageStatsProvider = batteryUsageStatsProvider;
     }
 
     /**
      * Kicks off the scheduling of power stats aggregation spans.
      */
     public void start(boolean enablePeriodicPowerStatsCollection) {
-        mBatteryStats.setBatteryResetListener(this::storeBatteryUsageStatsOnReset);
         mEnablePeriodicPowerStatsCollection = enablePeriodicPowerStatsCollection;
         if (mEnablePeriodicPowerStatsCollection) {
+            schedulePowerStatsAggregation();
             scheduleNextPowerStatsAggregation();
         }
     }
@@ -235,28 +231,6 @@
         mPowerStatsStore.storeAggregatedPowerStats(stats);
     }
 
-    private void storeBatteryUsageStatsOnReset(int resetReason) {
-        if (resetReason == BatteryStatsImpl.RESET_REASON_CORRUPT_FILE) {
-            return;
-        }
-
-        final BatteryUsageStats batteryUsageStats =
-                mBatteryUsageStatsProvider.getBatteryUsageStats(
-                        new BatteryUsageStatsQuery.Builder()
-                                .setMaxStatsAgeMs(0)
-                                .includePowerModels()
-                                .includeProcessStateData()
-                                .build());
-
-        // TODO(b/188068523): BatteryUsageStats should use monotonic time for start and end
-        // Once that change is made, we will be able to use the BatteryUsageStats' monotonic
-        // start time
-        long monotonicStartTime =
-                mMonotonicClock.monotonicTime() - batteryUsageStats.getStatsDuration();
-        mHandler.post(() ->
-                mPowerStatsStore.storeBatteryUsageStats(monotonicStartTime, batteryUsageStats));
-    }
-
     private void awaitCompletion() {
         ConditionVariable done = new ConditionVariable();
         mHandler.post(done::open);
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsUidResolver.java b/services/core/java/com/android/server/power/stats/PowerStatsUidResolver.java
new file mode 100644
index 0000000..8dc3609
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/PowerStatsUidResolver.java
@@ -0,0 +1,241 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import android.util.IntArray;
+import android.util.Slog;
+import android.util.SparseIntArray;
+
+import java.io.PrintWriter;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Maintains a map of isolated UIDs to their respective owner UIDs, to support combining
+ * power stats for isolated UIDs, which are typically short-lived, into the corresponding app UID.
+ */
+public class PowerStatsUidResolver {
+    private static final String TAG = "PowerStatsUidResolver";
+
+    /**
+     * Listener notified when isolated UIDs are created and removed.
+     */
+    public interface Listener {
+
+        /**
+         * Callback invoked when a new isolated UID is registered.
+         */
+        void onIsolatedUidAdded(int isolatedUid, int parentUid);
+
+        /**
+         * Callback invoked before an isolated UID is evicted from the resolver.
+         * If the listener calls {@link PowerStatsUidResolver#retainIsolatedUid}, the mapping
+         * will be retained until {@link PowerStatsUidResolver#releaseIsolatedUid} is called.
+         */
+        void onBeforeIsolatedUidRemoved(int isolatedUid, int parentUid);
+
+        /**
+         * Callback invoked when an isolated UID to owner UID mapping is removed.
+         */
+        void onAfterIsolatedUidRemoved(int isolatedUid, int parentUid);
+    }
+
+    /**
+     * Mapping isolated uids to the actual owning app uid.
+     */
+    private final SparseIntArray mIsolatedUids = new SparseIntArray();
+
+    /**
+     * Internal reference count of isolated uids.
+     */
+    private final SparseIntArray mIsolatedUidRefCounts = new SparseIntArray();
+
+    // Keep the list read-only in order to avoid locking during the delivery of listener calls.
+    private volatile List<Listener> mListeners = Collections.emptyList();
+
+    /**
+     * Adds a listener.
+     */
+    public void addListener(Listener listener) {
+        synchronized (this) {
+            List<Listener> newList = new ArrayList<>(mListeners);
+            newList.add(listener);
+            mListeners = Collections.unmodifiableList(newList);
+        }
+    }
+
+    /**
+     * Removes a listener.
+     */
+    public void removeListener(Listener listener) {
+        synchronized (this) {
+            List<Listener> newList = new ArrayList<>(mListeners);
+            newList.remove(listener);
+            mListeners = Collections.unmodifiableList(newList);
+        }
+    }
+
+    /**
+     * Remembers the connection between a newly created isolated UID and its owner app UID.
+     * Calls {@link Listener#onIsolatedUidAdded} on each registered listener.
+     */
+    public void noteIsolatedUidAdded(int isolatedUid, int parentUid) {
+        synchronized (this) {
+            mIsolatedUids.put(isolatedUid, parentUid);
+            mIsolatedUidRefCounts.put(isolatedUid, 1);
+        }
+
+        List<Listener> listeners = mListeners;
+        for (int i = listeners.size() - 1; i >= 0; i--) {
+            listeners.get(i).onIsolatedUidAdded(isolatedUid, parentUid);
+        }
+    }
+
+    /**
+     * Handles the removal of an isolated UID by invoking
+     * {@link Listener#onBeforeIsolatedUidRemoved} on each registered listener and the releases
+     * the UID, see {@link #releaseIsolatedUid}.
+     */
+    public void noteIsolatedUidRemoved(int isolatedUid, int parentUid) {
+        synchronized (this) {
+            int curUid = mIsolatedUids.get(isolatedUid, -1);
+            if (curUid != parentUid) {
+                Slog.wtf(TAG, "Attempt to remove an isolated UID " + isolatedUid
+                              + " with the parent UID " + parentUid
+                              + ". The registered parent UID is " + curUid);
+                return;
+            }
+        }
+
+        List<Listener> listeners = mListeners;
+        for (int i = listeners.size() - 1; i >= 0; i--) {
+            listeners.get(i).onBeforeIsolatedUidRemoved(isolatedUid, parentUid);
+        }
+
+        releaseIsolatedUid(isolatedUid);
+    }
+
+    /**
+     * Increments the ref count for an isolated uid.
+     * Call #releaseIsolatedUid to decrement.
+     */
+    public void retainIsolatedUid(int uid) {
+        synchronized (this) {
+            final int refCount = mIsolatedUidRefCounts.get(uid, 0);
+            if (refCount <= 0) {
+                // Uid is not mapped or referenced
+                Slog.w(TAG,
+                        "Attempted to increment ref counted of untracked isolated uid (" + uid
+                        + ")");
+                return;
+            }
+            mIsolatedUidRefCounts.put(uid, refCount + 1);
+        }
+    }
+
+    /**
+     * Decrements the ref count for the given isolated UID.  If the ref count drops to zero,
+     * removes the mapping and calls {@link Listener#onAfterIsolatedUidRemoved} on each registered
+     * listener.
+     */
+    public void releaseIsolatedUid(int isolatedUid) {
+        int parentUid;
+        synchronized (this) {
+            final int refCount = mIsolatedUidRefCounts.get(isolatedUid, 0) - 1;
+            if (refCount > 0) {
+                // Isolated uid is still being tracked
+                mIsolatedUidRefCounts.put(isolatedUid, refCount);
+                return;
+            }
+
+            final int idx = mIsolatedUids.indexOfKey(isolatedUid);
+            if (idx >= 0) {
+                parentUid = mIsolatedUids.valueAt(idx);
+                mIsolatedUids.removeAt(idx);
+                mIsolatedUidRefCounts.delete(isolatedUid);
+            } else {
+                Slog.w(TAG, "Attempted to remove untracked child uid (" + isolatedUid + ")");
+                return;
+            }
+        }
+
+        List<Listener> listeners = mListeners;
+        for (int i = listeners.size() - 1; i >= 0; i--) {
+            listeners.get(i).onAfterIsolatedUidRemoved(isolatedUid, parentUid);
+        }
+    }
+
+    /**
+     * Releases all isolated UIDs in the specified range, both ends inclusive.
+     */
+    public void releaseUidsInRange(int startUid, int endUid) {
+        IntArray toRelease;
+        synchronized (this) {
+            int startIndex = mIsolatedUids.indexOfKey(startUid);
+            int endIndex = mIsolatedUids.indexOfKey(endUid);
+
+            if (startIndex < 0) {
+                startIndex = ~startIndex;
+            }
+
+            if (endIndex < 0) {
+                // In this ~endIndex is pointing just past where endUid would be, so we must -1.
+                endIndex = ~endIndex - 1;
+            }
+
+            if (startIndex > endIndex) {
+                return;
+            }
+
+            toRelease = new IntArray(endIndex - startIndex);
+            for (int i = startIndex; i <= endIndex; i++) {
+                toRelease.add(mIsolatedUids.keyAt(i));
+            }
+        }
+
+        for (int i = toRelease.size() - 1; i >= 0; i--) {
+            releaseIsolatedUid(toRelease.get(i));
+        }
+    }
+
+    /**
+     * Given an isolated UID, returns the corresponding owner UID.  For a non-isolated
+     * UID, returns the UID itself.
+     */
+    public int mapUid(int uid) {
+        synchronized (this) {
+            return mIsolatedUids.get(/*key=*/uid, /*valueIfKeyNotFound=*/uid);
+        }
+    }
+
+    /**
+     * Dumps the current contents of the resolver for the sake of dumpsys.
+     */
+    public void dump(PrintWriter pw) {
+        pw.println("Currently mapped isolated uids:");
+        synchronized (this) {
+            final int numIsolatedUids = mIsolatedUids.size();
+            for (int i = 0; i < numIsolatedUids; i++) {
+                final int isolatedUid = mIsolatedUids.keyAt(i);
+                final int ownerUid = mIsolatedUids.valueAt(i);
+                final int refs = mIsolatedUidRefCounts.get(isolatedUid);
+                pw.println("  " + isolatedUid + "->" + ownerUid + " (ref count = " + refs + ")");
+            }
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/webkit/WebViewUpdateServiceImpl.java b/services/core/java/com/android/server/webkit/WebViewUpdateServiceImpl.java
index cfdef14..26c3fba 100644
--- a/services/core/java/com/android/server/webkit/WebViewUpdateServiceImpl.java
+++ b/services/core/java/com/android/server/webkit/WebViewUpdateServiceImpl.java
@@ -17,6 +17,7 @@
 
 import android.annotation.Nullable;
 import android.content.Context;
+import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager.NameNotFoundException;
 import android.content.pm.Signature;
@@ -29,8 +30,12 @@
 import android.webkit.WebViewProviderInfo;
 import android.webkit.WebViewProviderResponse;
 
+import com.android.server.LocalServices;
+import com.android.server.PinnerService;
+
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 /**
@@ -88,6 +93,8 @@
     private static final int MULTIPROCESS_SETTING_ON_VALUE = Integer.MAX_VALUE;
     private static final int MULTIPROCESS_SETTING_OFF_VALUE = Integer.MIN_VALUE;
 
+    private static final String PIN_GROUP = "webview";
+
     private final SystemInterface mSystemInterface;
     private final Context mContext;
 
@@ -339,6 +346,34 @@
         return newPackage;
     }
 
+    private void pinWebviewIfRequired(ApplicationInfo appInfo) {
+        PinnerService pinnerService = LocalServices.getService(PinnerService.class);
+        int webviewPinQuota = pinnerService.getWebviewPinQuota();
+        if (webviewPinQuota <= 0) {
+            return;
+        }
+
+        pinnerService.unpinGroup(PIN_GROUP);
+
+        ArrayList<String> apksToPin = new ArrayList<>();
+        boolean pinSharedFirst = appInfo.metaData.getBoolean("PIN_SHARED_LIBS_FIRST", true);
+        for (String sharedLib : appInfo.sharedLibraryFiles) {
+            apksToPin.add(sharedLib);
+        }
+        apksToPin.add(appInfo.sourceDir);
+        if (!pinSharedFirst) {
+            // We want to prioritize pinning of the native library that is most likely used by apps
+            // which in some build flavors live in the main apk and as a shared library for others.
+            Collections.reverse(apksToPin);
+        }
+        for (String apk : apksToPin) {
+            if (webviewPinQuota <= 0) {
+                break;
+            }
+            int bytesPinned = pinnerService.pinFile(apk, webviewPinQuota, appInfo, PIN_GROUP);
+            webviewPinQuota -= bytesPinned;
+        }
+    }
     /**
      * This is called when we change WebView provider, either when the current provider is
      * updated or a new provider is chosen / takes precedence.
@@ -347,6 +382,7 @@
         synchronized (mLock) {
             mAnyWebViewInstalled = true;
             if (mNumRelroCreationsStarted == mNumRelroCreationsFinished) {
+                pinWebviewIfRequired(newPackage.applicationInfo);
                 mCurrentWebViewPackage = newPackage;
 
                 // The relro creations might 'finish' (not start at all) before
diff --git a/services/core/java/com/android/server/wm/BackgroundActivityStartController.java b/services/core/java/com/android/server/wm/BackgroundActivityStartController.java
index 39e900a..eafaf8c 100644
--- a/services/core/java/com/android/server/wm/BackgroundActivityStartController.java
+++ b/services/core/java/com/android/server/wm/BackgroundActivityStartController.java
@@ -716,17 +716,18 @@
         // is allowed, or apps like live wallpaper with non app visible window will be allowed.
         final boolean appSwitchAllowedOrFg =
                 appSwitchState == APP_SWITCH_ALLOW || appSwitchState == APP_SWITCH_FG_ONLY;
-        final boolean allowCallingUidStartActivity =
-                ((appSwitchAllowedOrFg || mService.mActiveUids.hasNonAppVisibleWindow(callingUid))
-                        && callingUidHasAnyVisibleWindow)
-                        || isCallingUidPersistentSystemProcess;
-        if (allowCallingUidStartActivity) {
+        if (appSwitchAllowedOrFg && callingUidHasAnyVisibleWindow) {
             return new BalVerdict(BAL_ALLOW_VISIBLE_WINDOW,
-                    /*background*/ false,
-                    "callingUidHasAnyVisibleWindow = "
-                            + callingUid
-                            + ", isCallingUidPersistentSystemProcess = "
-                            + isCallingUidPersistentSystemProcess);
+                    /*background*/ false, "callingUid has visible window");
+        }
+        if (mService.mActiveUids.hasNonAppVisibleWindow(callingUid)) {
+            return new BalVerdict(BAL_ALLOW_VISIBLE_WINDOW,
+                    /*background*/ false, "callingUid has non-app visible window");
+        }
+
+        if (isCallingUidPersistentSystemProcess) {
+            return new BalVerdict(BAL_ALLOW_ALLOWLISTED_COMPONENT,
+                    /*background*/ false, "callingUid is persistent system process");
         }
 
         // don't abort if the callingUid has START_ACTIVITIES_FROM_BACKGROUND permission
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 07cbd58..bb59936 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -1153,12 +1153,12 @@
         mDisplay = display;
         mDisplayId = display.getDisplayId();
         mCurrentUniqueDisplayId = display.getUniqueId();
-        mDisplayUpdater = new ImmediateDisplayUpdater(this);
         mOffTokenAcquirer = mRootWindowContainer.mDisplayOffTokenAcquirer;
         mWallpaperController = new WallpaperController(mWmService, this);
         mWallpaperController.resetLargestDisplay(display);
         display.getDisplayInfo(mDisplayInfo);
         display.getMetrics(mDisplayMetrics);
+        mDisplayUpdater = new ImmediateDisplayUpdater(this);
         mSystemGestureExclusionLimit = mWmService.mConstants.mSystemGestureExclusionLimitDp
                 * mDisplayMetrics.densityDpi / DENSITY_DEFAULT;
         isDefaultDisplay = mDisplayId == DEFAULT_DISPLAY;
diff --git a/services/core/java/com/android/server/wm/ImmediateDisplayUpdater.java b/services/core/java/com/android/server/wm/ImmediateDisplayUpdater.java
index 72e8fcb..4af9013 100644
--- a/services/core/java/com/android/server/wm/ImmediateDisplayUpdater.java
+++ b/services/core/java/com/android/server/wm/ImmediateDisplayUpdater.java
@@ -30,6 +30,7 @@
 
     public ImmediateDisplayUpdater(@NonNull DisplayContent displayContent) {
         mDisplayContent = displayContent;
+        mDisplayInfo.copyFrom(mDisplayContent.getDisplayInfo());
     }
 
     @Override
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 2bd7327..522e7d2 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -627,7 +627,7 @@
     void refreshSecureSurfaceState() {
         forAllWindows((w) -> {
             if (w.mHasSurface) {
-                w.mWinAnimator.setSecureLocked(w.isSecureLocked());
+                w.setSecureLocked(w.isSecureLocked());
             }
         }, true /* traverseTopToBottom */);
     }
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index 7375512..de802b9 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -6177,6 +6177,7 @@
                 // Avoid resuming activities on secondary displays since we don't want bubble
                 // activities to be resumed while bubble is still collapsed.
                 // TODO(b/113840485): Having keyguard going away state for secondary displays.
+                && display != null
                 && display.isDefaultDisplay) {
             return false;
         }
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 575ae69b..dd2b48b 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -2327,8 +2327,8 @@
             boolean wallpaperMayMove = win.mViewVisibility != viewVisibility
                     && win.hasWallpaper();
             wallpaperMayMove |= (flagChanges & FLAG_SHOW_WALLPAPER) != 0;
-            if ((flagChanges & FLAG_SECURE) != 0 && winAnimator.mSurfaceController != null) {
-                winAnimator.mSurfaceController.setSecure(win.isSecureLocked());
+            if ((flagChanges & FLAG_SECURE) != 0) {
+                win.setSecureLocked(win.isSecureLocked());
             }
 
             final boolean wasVisible = win.isVisible();
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 6d6bcc8..e1f1f66 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -109,6 +109,7 @@
 import static com.android.internal.protolog.ProtoLogGroup.WM_DEBUG_STARTING_WINDOW;
 import static com.android.internal.protolog.ProtoLogGroup.WM_DEBUG_SYNC_ENGINE;
 import static com.android.internal.protolog.ProtoLogGroup.WM_DEBUG_WINDOW_INSETS;
+import static com.android.internal.protolog.ProtoLogGroup.WM_SHOW_TRANSACTIONS;
 import static com.android.server.policy.WindowManagerPolicy.FINISH_LAYOUT_REDO_WALLPAPER;
 import static com.android.server.policy.WindowManagerPolicy.TRANSIT_ENTER;
 import static com.android.server.policy.WindowManagerPolicy.TRANSIT_EXIT;
@@ -177,6 +178,7 @@
 import static com.android.server.wm.WindowStateProto.VIEW_VISIBILITY;
 import static com.android.server.wm.WindowStateProto.WINDOW_CONTAINER;
 import static com.android.server.wm.WindowStateProto.WINDOW_FRAMES;
+import static com.android.window.flags.Flags.secureWindowState;
 import static com.android.window.flags.Flags.surfaceTrustedOverlay;
 
 import android.annotation.CallSuper;
@@ -1195,6 +1197,9 @@
         if (surfaceTrustedOverlay() && isWindowTrustedOverlay()) {
             getPendingTransaction().setTrustedOverlay(mSurfaceControl, true);
         }
+        if (secureWindowState()) {
+            getPendingTransaction().setSecure(mSurfaceControl, isSecureLocked());
+        }
     }
 
     void updateTrustedOverlay() {
@@ -6042,4 +6047,25 @@
         // Cancel any draw requests during a sync.
         return mPrepareSyncSeqId > 0;
     }
+
+    void setSecureLocked(boolean isSecure) {
+        ProtoLog.i(WM_SHOW_TRANSACTIONS, "SURFACE isSecure=%b: %s", isSecure, getName());
+        if (secureWindowState()) {
+            if (mSurfaceControl == null) {
+                return;
+            }
+            getPendingTransaction().setSecure(mSurfaceControl, isSecure);
+        } else {
+            if (mWinAnimator.mSurfaceController == null
+                    || mWinAnimator.mSurfaceController.mSurfaceControl == null) {
+                return;
+            }
+            getPendingTransaction().setSecure(mWinAnimator.mSurfaceController.mSurfaceControl,
+                    isSecure);
+        }
+        if (mDisplayContent != null) {
+            mDisplayContent.refreshImeSecureFlag(getSyncTransaction());
+        }
+        mWmService.scheduleAnimationLocked();
+    }
 }
diff --git a/services/core/java/com/android/server/wm/WindowStateAnimator.java b/services/core/java/com/android/server/wm/WindowStateAnimator.java
index 3aac816..44cd23d 100644
--- a/services/core/java/com/android/server/wm/WindowStateAnimator.java
+++ b/services/core/java/com/android/server/wm/WindowStateAnimator.java
@@ -44,6 +44,7 @@
 import static com.android.server.wm.WindowStateAnimatorProto.DRAW_STATE;
 import static com.android.server.wm.WindowStateAnimatorProto.SURFACE;
 import static com.android.server.wm.WindowStateAnimatorProto.SYSTEM_DECOR_RECT;
+import static com.android.window.flags.Flags.secureWindowState;
 
 import android.content.Context;
 import android.graphics.PixelFormat;
@@ -286,8 +287,10 @@
         int flags = SurfaceControl.HIDDEN;
         final WindowManager.LayoutParams attrs = w.mAttrs;
 
-        if (w.isSecureLocked()) {
-            flags |= SurfaceControl.SECURE;
+        if (!secureWindowState()) {
+            if (w.isSecureLocked()) {
+                flags |= SurfaceControl.SECURE;
+            }
         }
 
         if ((mWin.mAttrs.privateFlags & PRIVATE_FLAG_IS_ROUNDED_CORNERS_OVERLAY) != 0) {
@@ -488,13 +491,6 @@
         mSurfaceController.setOpaque(isOpaque);
     }
 
-    void setSecureLocked(boolean isSecure) {
-        if (mSurfaceController == null) {
-            return;
-        }
-        mSurfaceController.setSecure(isSecure);
-    }
-
     void setColorSpaceAgnosticLocked(boolean agnostic) {
         if (mSurfaceController == null) {
             return;
diff --git a/services/core/java/com/android/server/wm/WindowSurfaceController.java b/services/core/java/com/android/server/wm/WindowSurfaceController.java
index d348491..4456a94 100644
--- a/services/core/java/com/android/server/wm/WindowSurfaceController.java
+++ b/services/core/java/com/android/server/wm/WindowSurfaceController.java
@@ -24,7 +24,6 @@
 import static com.android.internal.protolog.ProtoLogGroup.WM_SHOW_SURFACE_ALLOC;
 import static com.android.internal.protolog.ProtoLogGroup.WM_SHOW_TRANSACTIONS;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_VISIBILITY;
-import static com.android.server.wm.WindowManagerDebugConfig.SHOW_LIGHT_TRANSACTIONS;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 import static com.android.server.wm.WindowSurfaceControllerProto.SHOWN;
@@ -152,24 +151,6 @@
         mService.scheduleAnimationLocked();
     }
 
-    void setSecure(boolean isSecure) {
-        ProtoLog.i(WM_SHOW_TRANSACTIONS, "SURFACE isSecure=%b: %s", isSecure, title);
-
-        if (mSurfaceControl == null) {
-            return;
-        }
-        if (SHOW_LIGHT_TRANSACTIONS) Slog.i(TAG, ">>> OPEN TRANSACTION setSecureLocked");
-
-        final SurfaceControl.Transaction t = mAnimator.mWin.getPendingTransaction();
-        t.setSecure(mSurfaceControl, isSecure);
-
-        final DisplayContent dc = mAnimator.mWin.mDisplayContent;
-        if (dc != null) {
-            dc.refreshImeSecureFlag(t);
-        }
-        mService.scheduleAnimationLocked();
-    }
-
     void setColorSpaceAgnostic(SurfaceControl.Transaction t, boolean agnostic) {
         ProtoLog.i(WM_SHOW_TRANSACTIONS, "SURFACE isColorSpaceAgnostic=%b: %s", agnostic, title);
 
diff --git a/services/credentials/java/com/android/server/credentials/RequestSession.java b/services/credentials/java/com/android/server/credentials/RequestSession.java
index 1988bb6..da44aac 100644
--- a/services/credentials/java/com/android/server/credentials/RequestSession.java
+++ b/services/credentials/java/com/android/server/credentials/RequestSession.java
@@ -23,12 +23,14 @@
 import android.content.Context;
 import android.content.Intent;
 import android.credentials.CredentialProviderInfo;
+import android.credentials.flags.Flags;
 import android.credentials.ui.ProviderData;
 import android.credentials.ui.UserSelectionDialogResult;
 import android.os.Binder;
 import android.os.CancellationSignal;
 import android.os.Handler;
 import android.os.IBinder;
+import android.os.IInterface;
 import android.os.Looper;
 import android.os.RemoteException;
 import android.os.UserHandle;
@@ -94,6 +96,9 @@
 
     private final Set<ComponentName> mEnabledProviders;
 
+    private final RequestSessionDeathRecipient mDeathRecipient =
+            new RequestSessionDeathRecipient();
+
     protected PendingIntent mPendingIntent;
 
     @NonNull
@@ -141,11 +146,26 @@
         mRequestSessionMetric.collectInitialPhaseMetricInfo(timestampStarted,
                 mCallingUid, ApiName.getMetricCodeFromRequestInfo(mRequestType));
         setCancellationListener();
+        if (Flags.clearSessionEnabled()) {
+            setUpClientCallbackListener();
+        }
+    }
+
+    private void setUpClientCallbackListener() {
+        if (mClientCallback != null && mClientCallback instanceof IInterface) {
+            IInterface callback = (IInterface) mClientCallback;
+            try {
+                callback.asBinder().linkToDeath(mDeathRecipient, 0);
+            } catch (RemoteException e) {
+                Slog.e(TAG, e.getMessage());
+            }
+        }
     }
 
     private void setCancellationListener() {
         mCancellationSignal.setOnCancelListener(
                 () -> {
+                    Slog.d(TAG, "Cancellation invoked from the client - clearing session");
                     boolean isUiActive = maybeCancelUi();
                     finishSession(!isUiActive);
                 }
@@ -168,6 +188,17 @@
         return false;
     }
 
+    private boolean isUiWaitingForData() {
+        // Technically, the status can also be IN_PROGRESS when the user has made a selection
+        // so this an over estimation, but safe to do so as it is used for cancellation
+        // propagation to the provider in a very narrow time frame. If provider has
+        // already responded, cancellation is not an issue as the cancellation listener
+        // is independent of the service binding.
+        // TODO(b/313512500): Do not propagate cancelation if provider has responded in
+        // query phase.
+        return mCredentialManagerUi.getStatus() == CredentialManagerUi.UiStatus.IN_PROGRESS;
+    }
+
     public abstract ProviderSession initiateProviderSession(CredentialProviderInfo providerInfo,
             RemoteCredentialService remoteCredentialService);
 
@@ -373,4 +404,12 @@
         return chosenProviderSession != null && chosenProviderSession.mProviderInfo != null
                 && chosenProviderSession.mProviderInfo.isPrimary();
     }
+
+    private class RequestSessionDeathRecipient implements IBinder.DeathRecipient {
+        @Override
+        public void binderDied() {
+            Slog.d(TAG, "Client binder died - clearing session");
+            finishSession(isUiWaitingForData());
+        }
+    }
 }
diff --git a/services/permission/OWNERS b/services/permission/OWNERS
index e464038..487c992 100644
--- a/services/permission/OWNERS
+++ b/services/permission/OWNERS
@@ -1,5 +1,3 @@
 #Bug component: 137825
 
-joecastro@google.com
-ntmyren@google.com
-zhanghai@google.com
+include platform/frameworks/base:/core/java/android/permission/OWNERS
diff --git a/services/permission/java/com/android/server/permission/access/appop/AppOpService.kt b/services/permission/java/com/android/server/permission/access/appop/AppOpService.kt
index f94a0d6..8f464d4 100644
--- a/services/permission/java/com/android/server/permission/access/appop/AppOpService.kt
+++ b/services/permission/java/com/android/server/permission/access/appop/AppOpService.kt
@@ -71,7 +71,7 @@
         // Not implemented because upgrades are handled automatically.
     }
 
-    override fun getNonDefaultUidModes(uid: Int): SparseIntArray {
+    override fun getNonDefaultUidModes(uid: Int, persistentDeviceId: String): SparseIntArray {
         return opNameMapToOpSparseArray(getUidModes(uid))
     }
 
@@ -79,7 +79,7 @@
         return opNameMapToOpSparseArray(getPackageModes(packageName, userId))
     }
 
-    override fun getUidMode(uid: Int, op: Int): Int {
+    override fun getUidMode(uid: Int, persistentDeviceId: String, op: Int): Int {
         val appId = UserHandle.getAppId(uid)
         val userId = UserHandle.getUserId(uid)
         val opName = AppOpsManager.opToPublicName(op)
@@ -92,7 +92,7 @@
         return service.getState { with(appIdPolicy) { getAppOpModes(appId, userId) } }?.map
     }
 
-    override fun setUidMode(uid: Int, op: Int, mode: Int): Boolean {
+    override fun setUidMode(uid: Int, persistentDeviceId: String, op: Int, mode: Int): Boolean {
         val appId = UserHandle.getAppId(uid)
         val userId = UserHandle.getUserId(uid)
         val opName = AppOpsManager.opToPublicName(op)
@@ -150,7 +150,7 @@
         // and we have our own persistence.
     }
 
-    override fun getForegroundOps(uid: Int): SparseBooleanArray {
+    override fun getForegroundOps(uid: Int, persistentDeviceId: String): SparseBooleanArray {
         return SparseBooleanArray().apply {
             getUidModes(uid)?.forEachIndexed { _, op, mode ->
                 if (mode == AppOpsManager.MODE_FOREGROUND) {
diff --git a/services/tests/mockingservicestests/src/com/android/server/appop/AppOpsUpgradeTest.java b/services/tests/mockingservicestests/src/com/android/server/appop/AppOpsUpgradeTest.java
index 92d1118..4f672f8 100644
--- a/services/tests/mockingservicestests/src/com/android/server/appop/AppOpsUpgradeTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/appop/AppOpsUpgradeTest.java
@@ -19,6 +19,7 @@
 import static android.app.AppOpsManager.OP_SCHEDULE_EXACT_ALARM;
 import static android.app.AppOpsManager.OP_USE_FULL_SCREEN_INTENT;
 import static android.app.AppOpsManager._NUM_OP;
+import static android.companion.virtual.VirtualDeviceManager.PERSISTENT_DEVICE_ID_DEFAULT;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
@@ -208,8 +209,8 @@
     private void assertSameModes(AppOpsCheckingServiceImpl testService, int op1, int op2) {
         for (int uid : testService.getUidsWithNonDefaultModes()) {
             assertEquals(
-                    testService.getUidMode(uid, op1),
-                    testService.getUidMode(uid, op2)
+                    testService.getUidMode(uid, PERSISTENT_DEVICE_ID_DEFAULT, op1),
+                    testService.getUidMode(uid, PERSISTENT_DEVICE_ID_DEFAULT, op2)
             );
         }
         for (UserPackage pkg : testService.getPackagesWithNonDefaultModes()) {
@@ -275,7 +276,9 @@
                 } else {
                     expectedMode = previousMode;
                 }
-                int mode = testService.getUidMode(uid, OP_SCHEDULE_EXACT_ALARM);
+                int mode =
+                        testService.getUidMode(
+                                uid, PERSISTENT_DEVICE_ID_DEFAULT, OP_SCHEDULE_EXACT_ALARM);
                 assertEquals(expectedMode, mode);
             }
         }
@@ -284,7 +287,9 @@
         int[] unrelatedUidsInFile = {10225, 10178};
 
         for (int uid : unrelatedUidsInFile) {
-            int mode = testService.getUidMode(uid, OP_SCHEDULE_EXACT_ALARM);
+            int mode =
+                    testService.getUidMode(
+                            uid, PERSISTENT_DEVICE_ID_DEFAULT, OP_SCHEDULE_EXACT_ALARM);
             assertEquals(AppOpsManager.opToDefaultMode(OP_SCHEDULE_EXACT_ALARM), mode);
         }
     }
@@ -331,7 +336,9 @@
                 final int uid = UserHandle.getUid(userId, appId);
                 final int expectedMode = AppOpsManager.opToDefaultMode(OP_USE_FULL_SCREEN_INTENT);
                 synchronized (testService) {
-                    int mode = testService.getUidMode(uid, OP_USE_FULL_SCREEN_INTENT);
+                    int mode =
+                            testService.getUidMode(
+                                    uid, PERSISTENT_DEVICE_ID_DEFAULT, OP_USE_FULL_SCREEN_INTENT);
                     assertEquals(expectedMode, mode);
                 }
             }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java
index 2003d04..ca7de7c 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java
@@ -90,6 +90,10 @@
 
     private AggregatedPowerStats prepareAggregatePowerStats() {
         AggregatedPowerStats stats = new AggregatedPowerStats(mAggregatedPowerStatsConfig);
+
+        PowerStats ps = new PowerStats(mPowerComponentDescriptor);
+        stats.addPowerStats(ps, 0);
+
         stats.addClockUpdate(1000, 456);
         stats.setDuration(789);
 
@@ -100,7 +104,6 @@
         stats.setUidState(APP_2, AggregatedPowerStatsConfig.STATE_PROCESS_STATE,
                 BatteryConsumer.PROCESS_STATE_FOREGROUND, 2000);
 
-        PowerStats ps = new PowerStats(mPowerComponentDescriptor);
         ps.stats[0] = 100;
         ps.stats[1] = 987;
 
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java
index 663af5d..9c2834d 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java
@@ -215,7 +215,7 @@
 
     public class TestBatteryStatsImpl extends BatteryStatsImpl {
         public TestBatteryStatsImpl(Context context) {
-            super(Clock.SYSTEM_CLOCK, null);
+            super(Clock.SYSTEM_CLOCK, null, null, null);
             mPowerProfile = new PowerProfile(context, true /* forTest */);
 
             SparseArray<int[]> cpusByPolicy = new SparseArray<>();
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsCpuTimesTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsCpuTimesTest.java
index 55ffa1a..f9f32b2 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsCpuTimesTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsCpuTimesTest.java
@@ -37,6 +37,8 @@
 import static org.mockito.Mockito.when;
 
 import android.os.BatteryStats;
+import android.os.Handler;
+import android.os.Looper;
 import android.os.UserHandle;
 import android.util.SparseArray;
 import android.util.SparseLongArray;
@@ -97,6 +99,7 @@
     BatteryStatsImpl.UserInfoProvider mUserInfoProvider;
 
     private MockClock mClocks;
+    private PowerStatsUidResolver mPowerStatsUidResolver;
     private MockBatteryStatsImpl mBatteryStatsImpl;
     private KernelCpuSpeedReader[] mKernelCpuSpeedReaders;
 
@@ -105,7 +108,9 @@
         MockitoAnnotations.initMocks(this);
 
         mClocks = new MockClock();
-        mBatteryStatsImpl = new MockBatteryStatsImpl(mClocks)
+        Handler handler = new Handler(Looper.getMainLooper());
+        mPowerStatsUidResolver = new PowerStatsUidResolver();
+        mBatteryStatsImpl = new MockBatteryStatsImpl(mClocks, null, handler, mPowerStatsUidResolver)
                 .setTestCpuScalingPolicies()
                 .setKernelCpuUidUserSysTimeReader(mCpuUidUserSysTimeReader)
                 .setKernelCpuUidFreqTimeReader(mCpuUidFreqTimeReader)
@@ -374,7 +379,7 @@
 
         // PRECONDITIONS
         final int ownerUid = UserHandle.getUid(testUserId, FIRST_APPLICATION_UID + 42);
-        mBatteryStatsImpl.addIsolatedUidLocked(isolatedUid, ownerUid);
+        mPowerStatsUidResolver.noteIsolatedUidAdded(isolatedUid, ownerUid);
         final long[][] deltasUs = {
                 {9379, 3332409833484L}, {493247, 723234}, {3247819, 123348}
         };
@@ -965,7 +970,7 @@
 
         // PRECONDITIONS
         final int ownerUid = UserHandle.getUid(testUserId, FIRST_APPLICATION_UID + 42);
-        mBatteryStatsImpl.addIsolatedUidLocked(isolatedUid, ownerUid);
+        mPowerStatsUidResolver.noteIsolatedUidAdded(isolatedUid, ownerUid);
         final long[][] deltasMs = {
                 {3, 12, 55, 100, 32},
                 {32483274, 232349349, 123, 2398, 0},
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
index 5ebc6ca..8d51592 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
@@ -39,14 +39,22 @@
 import android.app.ActivityManager;
 import android.bluetooth.BluetoothActivityEnergyInfo;
 import android.bluetooth.UidTraffic;
+import android.content.Context;
+import android.os.BatteryConsumer;
+import android.os.BatteryManager;
 import android.os.BatteryStats;
+import android.os.BatteryUsageStats;
 import android.os.BluetoothBatteryStats;
+import android.os.ConditionVariable;
+import android.os.Handler;
+import android.os.HandlerThread;
 import android.os.Parcel;
 import android.os.WakeLockStats;
 import android.os.WorkSource;
 import android.util.SparseArray;
 import android.view.Display;
 
+import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.LargeTest;
 import androidx.test.runner.AndroidJUnit4;
 
@@ -65,6 +73,8 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.io.File;
+import java.time.Instant;
 import java.util.List;
 
 @LargeTest
@@ -93,6 +103,11 @@
 
     private final MockClock mMockClock = new MockClock();
     private MockBatteryStatsImpl mBatteryStatsImpl;
+    private Handler mHandler;
+    private PowerStatsStore mPowerStatsStore;
+    private BatteryUsageStatsProvider mBatteryUsageStatsProvider;
+    @Mock
+    private PowerStatsExporter mPowerStatsExporter;
 
     @Before
     public void setUp() {
@@ -103,12 +118,23 @@
         when(mKernelSingleUidTimeReader.singleUidCpuTimesAvailable()).thenReturn(true);
         when(mKernelWakelockReader.readKernelWakelockStats(
                 any(KernelWakelockStats.class))).thenReturn(mKernelWakelockStats);
-        mBatteryStatsImpl = new MockBatteryStatsImpl(mMockClock)
+        HandlerThread bgThread = new HandlerThread("bg thread");
+        bgThread.start();
+        mHandler = new Handler(bgThread.getLooper());
+        mBatteryStatsImpl = new MockBatteryStatsImpl(mMockClock, null, mHandler)
                 .setPowerProfile(mPowerProfile)
                 .setCpuScalingPolicies(mCpuScalingPolicies)
                 .setKernelCpuUidFreqTimeReader(mKernelUidCpuFreqTimeReader)
                 .setKernelSingleUidTimeReader(mKernelSingleUidTimeReader)
                 .setKernelWakelockReader(mKernelWakelockReader);
+
+        final Context context = InstrumentationRegistry.getContext();
+        File systemDir = context.getCacheDir();
+        mPowerStatsStore = new PowerStatsStore(systemDir, mHandler,
+                new AggregatedPowerStatsConfig());
+        mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context, mPowerStatsExporter,
+                mPowerProfile, mBatteryStatsImpl.getCpuScalingPolicies(), mPowerStatsStore,
+                mMockClock);
     }
 
     @Test
@@ -754,4 +780,76 @@
         parcel.recycle();
         return info;
     }
+
+    @Test
+    public void storeBatteryUsageStatsOnReset() {
+        mBatteryStatsImpl.forceRecordAllHistory();
+
+        mMockClock.currentTime = Instant.parse("2023-01-02T03:04:05.00Z").toEpochMilli();
+        mMockClock.realtime = 7654321;
+
+        synchronized (mBatteryStatsImpl) {
+            mBatteryStatsImpl.setOnBatteryLocked(mMockClock.realtime, mMockClock.uptime, true,
+                    BatteryManager.BATTERY_STATUS_DISCHARGING, 50, 0);
+            // Will not save to PowerStatsStore because "saveBatteryUsageStatsOnReset" has not
+            // been called yet.
+            mBatteryStatsImpl.resetAllStatsAndHistoryLocked(
+                    BatteryStatsImpl.RESET_REASON_ADB_COMMAND);
+        }
+
+        assertThat(mPowerStatsStore.getTableOfContents()).isEmpty();
+
+        mBatteryStatsImpl.saveBatteryUsageStatsOnReset(mBatteryUsageStatsProvider,
+                mPowerStatsStore);
+
+        synchronized (mBatteryStatsImpl) {
+            mBatteryStatsImpl.noteFlashlightOnLocked(42, mMockClock.realtime, mMockClock.uptime);
+        }
+
+        mMockClock.realtime += 60000;
+        mMockClock.currentTime += 60000;
+
+        synchronized (mBatteryStatsImpl) {
+            mBatteryStatsImpl.noteFlashlightOffLocked(42, mMockClock.realtime, mMockClock.uptime);
+        }
+
+        mMockClock.realtime += 60000;
+        mMockClock.currentTime += 60000;
+
+        // Battery stats reset should have the side-effect of saving accumulated battery usage stats
+        synchronized (mBatteryStatsImpl) {
+            mBatteryStatsImpl.resetAllStatsAndHistoryLocked(
+                    BatteryStatsImpl.RESET_REASON_ADB_COMMAND);
+        }
+
+        // Await completion
+        ConditionVariable done = new ConditionVariable();
+        mHandler.post(done::open);
+        done.block();
+
+        List<PowerStatsSpan.Metadata> contents = mPowerStatsStore.getTableOfContents();
+        assertThat(contents).hasSize(1);
+
+        PowerStatsSpan.Metadata metadata = contents.get(0);
+
+        PowerStatsSpan span = mPowerStatsStore.loadPowerStatsSpan(metadata.getId(),
+                BatteryUsageStatsSection.TYPE);
+        assertThat(span).isNotNull();
+
+        List<PowerStatsSpan.TimeFrame> timeFrames = span.getMetadata().getTimeFrames();
+        assertThat(timeFrames).hasSize(1);
+        assertThat(timeFrames.get(0).startMonotonicTime).isEqualTo(7654321);
+        assertThat(timeFrames.get(0).duration).isEqualTo(120000);
+
+        List<PowerStatsSpan.Section> sections = span.getSections();
+        assertThat(sections).hasSize(1);
+
+        PowerStatsSpan.Section section = sections.get(0);
+        assertThat(section.getType()).isEqualTo(BatteryUsageStatsSection.TYPE);
+        BatteryUsageStats bus = ((BatteryUsageStatsSection) section).getBatteryUsageStats();
+        assertThat(bus.getAggregateBatteryConsumer(
+                        BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
+                .getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
+                .isEqualTo(60000);
+    }
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java
index 7ef1a3f..24c67f8 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java
@@ -35,6 +35,8 @@
 import android.os.BatteryStats;
 import android.os.BatteryStats.HistoryItem;
 import android.os.BatteryStats.Uid.Sensor;
+import android.os.Handler;
+import android.os.Looper;
 import android.os.Process;
 import android.os.UserHandle;
 import android.os.WorkSource;
@@ -155,7 +157,9 @@
     @SmallTest
     public void testNoteStartWakeLocked_isolatedUid() throws Exception {
         final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
-        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+        PowerStatsUidResolver uidResolver = new PowerStatsUidResolver();
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks, null,
+                new Handler(Looper.getMainLooper()), uidResolver);
 
         int pid = 10;
         String name = "name";
@@ -165,7 +169,7 @@
         isolatedWorkChain.addNode(ISOLATED_UID, name);
 
         // Map ISOLATED_UID to UID.
-        bi.addIsolatedUidLocked(ISOLATED_UID, UID);
+        uidResolver.noteIsolatedUidAdded(ISOLATED_UID, UID);
 
         bi.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_TOP);
@@ -195,7 +199,9 @@
     @SmallTest
     public void testNoteStartWakeLocked_isolatedUidRace() throws Exception {
         final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
-        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
+        PowerStatsUidResolver uidResolver = new PowerStatsUidResolver();
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks, null,
+                new Handler(Looper.getMainLooper()), uidResolver);
 
         int pid = 10;
         String name = "name";
@@ -205,7 +211,7 @@
         isolatedWorkChain.addNode(ISOLATED_UID, name);
 
         // Map ISOLATED_UID to UID.
-        bi.addIsolatedUidLocked(ISOLATED_UID, UID);
+        uidResolver.noteIsolatedUidAdded(ISOLATED_UID, UID);
 
         bi.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_TOP);
@@ -216,7 +222,7 @@
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
 
         clocks.realtime = clocks.uptime = 150;
-        bi.maybeRemoveIsolatedUidLocked(ISOLATED_UID, clocks.realtime, clocks.uptime);
+        uidResolver.releaseIsolatedUid(ISOLATED_UID);
 
         clocks.realtime = clocks.uptime = 220;
         bi.noteStopWakeLocked(ISOLATED_UID, pid, isolatedWorkChain, name, historyName,
@@ -237,8 +243,9 @@
     @SmallTest
     public void testNoteLongPartialWakelockStart_isolatedUid() throws Exception {
         final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
-        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
-
+        PowerStatsUidResolver uidResolver = new PowerStatsUidResolver();
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks, null,
+                new Handler(Looper.getMainLooper()), uidResolver);
 
         bi.setRecordAllHistoryLocked(true);
         bi.forceRecordAllHistory();
@@ -251,7 +258,7 @@
         isolatedWorkChain.addNode(ISOLATED_UID, name);
 
         // Map ISOLATED_UID to UID.
-        bi.addIsolatedUidLocked(ISOLATED_UID, UID);
+        uidResolver.noteIsolatedUidAdded(ISOLATED_UID, UID);
 
         bi.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_TOP);
@@ -290,8 +297,9 @@
     @SmallTest
     public void testNoteLongPartialWakelockStart_isolatedUidRace() throws Exception {
         final MockClock clocks = new MockClock(); // holds realtime and uptime in ms
-        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks);
-
+        PowerStatsUidResolver uidResolver = new PowerStatsUidResolver();
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(clocks, null,
+                new Handler(Looper.getMainLooper()), uidResolver);
 
         bi.setRecordAllHistoryLocked(true);
         bi.forceRecordAllHistory();
@@ -304,7 +312,7 @@
         isolatedWorkChain.addNode(ISOLATED_UID, name);
 
         // Map ISOLATED_UID to UID.
-        bi.addIsolatedUidLocked(ISOLATED_UID, UID);
+        uidResolver.noteIsolatedUidAdded(ISOLATED_UID, UID);
 
         bi.updateTimeBasesLocked(true, Display.STATE_OFF, 0, 0);
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_TOP);
@@ -314,7 +322,7 @@
         bi.noteUidProcessStateLocked(UID, ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND);
 
         clocks.realtime = clocks.uptime = 150;
-        bi.maybeRemoveIsolatedUidLocked(ISOLATED_UID, clocks.realtime, clocks.uptime);
+        uidResolver.releaseIsolatedUid(ISOLATED_UID);
 
         clocks.realtime = clocks.uptime = 220;
         bi.noteLongPartialWakelockFinish(name, historyName, ISOLATED_UID);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
index b1da1fc..7148b16 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
@@ -28,9 +28,7 @@
 import android.os.BatteryStats;
 import android.os.BatteryUsageStats;
 import android.os.BatteryUsageStatsQuery;
-import android.os.Handler;
-import android.os.Looper;
-import android.os.Message;
+import android.os.ConditionVariable;
 import android.os.Parcel;
 import android.os.Process;
 import android.os.UidBatteryConsumer;
@@ -40,7 +38,6 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.os.BatteryStatsHistoryIterator;
-import com.android.internal.os.MonotonicClock;
 import com.android.internal.os.PowerProfile;
 
 import org.junit.Rule;
@@ -72,10 +69,11 @@
         BatteryStatsImpl batteryStats = prepareBatteryStats();
 
         Context context = InstrumentationRegistry.getContext();
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context, batteryStats);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context, null,
+                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), null, mMockClock);
 
         final BatteryUsageStats batteryUsageStats =
-                provider.getBatteryUsageStats(BatteryUsageStatsQuery.DEFAULT);
+                provider.getBatteryUsageStats(batteryStats, BatteryUsageStatsQuery.DEFAULT);
 
         final List<UidBatteryConsumer> uidBatteryConsumers =
                 batteryUsageStats.getUidBatteryConsumers();
@@ -99,10 +97,11 @@
         BatteryStatsImpl batteryStats = prepareBatteryStats();
 
         Context context = InstrumentationRegistry.getContext();
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context, batteryStats);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context, null,
+                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), null, mMockClock);
 
         final BatteryUsageStats batteryUsageStats =
-                provider.getBatteryUsageStats(
+                provider.getBatteryUsageStats(batteryStats,
                         new BatteryUsageStatsQuery.Builder()
                                 .includePowerComponents(
                                         new int[]{BatteryConsumer.POWER_COMPONENT_AUDIO})
@@ -204,10 +203,11 @@
         }
 
         Context context = InstrumentationRegistry.getContext();
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context, batteryStats);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context, null,
+                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), null, mMockClock);
 
         final BatteryUsageStats batteryUsageStats =
-                provider.getBatteryUsageStats(
+                provider.getBatteryUsageStats(batteryStats,
                         new BatteryUsageStatsQuery.Builder().includeBatteryHistory().build());
 
         Parcel in = Parcel.obtain();
@@ -292,11 +292,11 @@
         }
 
         Context context = InstrumentationRegistry.getContext();
-        BatteryUsageStatsProvider
-                provider = new BatteryUsageStatsProvider(context, batteryStats);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context, null,
+                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), null, mMockClock);
 
         final BatteryUsageStats batteryUsageStats =
-                provider.getBatteryUsageStats(
+                provider.getBatteryUsageStats(batteryStats,
                         new BatteryUsageStatsQuery.Builder().includeBatteryHistory().build());
 
         Parcel parcel = Parcel.obtain();
@@ -352,27 +352,22 @@
 
     @Test
     public void shouldUpdateStats() {
-        Context context = InstrumentationRegistry.getContext();
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context,
-                mStatsRule.getBatteryStats());
-
         final List<BatteryUsageStatsQuery> queries = List.of(
                 new BatteryUsageStatsQuery.Builder().setMaxStatsAgeMs(1000).build(),
                 new BatteryUsageStatsQuery.Builder().setMaxStatsAgeMs(2000).build()
         );
 
-        mStatsRule.setTime(10500, 0);
-        assertThat(provider.shouldUpdateStats(queries, 10000)).isFalse();
+        assertThat(BatteryUsageStatsProvider.shouldUpdateStats(queries,
+                10500, 10000)).isFalse();
 
-        mStatsRule.setTime(11500, 0);
-        assertThat(provider.shouldUpdateStats(queries, 10000)).isTrue();
+        assertThat(BatteryUsageStatsProvider.shouldUpdateStats(queries,
+                11500, 10000)).isTrue();
     }
 
     @Test
     public void testAggregateBatteryStats() {
         Context context = InstrumentationRegistry.getContext();
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
-        MonotonicClock monotonicClock = new MonotonicClock(0, mStatsRule.getMockClock());
 
         setTime(5 * MINUTE_IN_MS);
         synchronized (batteryStats) {
@@ -381,14 +376,17 @@
 
         PowerStatsStore powerStatsStore = new PowerStatsStore(
                 new File(context.getCacheDir(), "BatteryUsageStatsProviderTest"),
-                new TestHandler(), null);
+                mStatsRule.getHandler(), null);
+        powerStatsStore.reset();
 
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context,
-                batteryStats, powerStatsStore);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context, null,
+                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), powerStatsStore,
+                mMockClock);
 
-        batteryStats.setBatteryResetListener(reason ->
-                powerStatsStore.storeBatteryUsageStats(monotonicClock.monotonicTime(),
-                        provider.getBatteryUsageStats(BatteryUsageStatsQuery.DEFAULT)));
+        batteryStats.saveBatteryUsageStatsOnReset(provider, powerStatsStore);
+        synchronized (batteryStats) {
+            batteryStats.resetAllStatsAndHistoryLocked(BatteryStatsImpl.RESET_REASON_ADB_COMMAND);
+        }
 
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOnLocked(APP_UID,
@@ -441,11 +439,16 @@
         }
         setTime(95 * MINUTE_IN_MS);
 
+        // Await completion
+        ConditionVariable done = new ConditionVariable();
+        mStatsRule.getHandler().post(done::open);
+        done.block();
+
         // Include the first and the second snapshot, but not the third or current
         BatteryUsageStatsQuery query = new BatteryUsageStatsQuery.Builder()
                 .aggregateSnapshots(20 * MINUTE_IN_MS, 60 * MINUTE_IN_MS)
                 .build();
-        final BatteryUsageStats stats = provider.getBatteryUsageStats(query);
+        final BatteryUsageStats stats = provider.getBatteryUsageStats(batteryStats, query);
 
         assertThat(stats.getStatsStartTimestamp()).isEqualTo(5 * MINUTE_IN_MS);
         assertThat(stats.getStatsEndTimestamp()).isEqualTo(55 * MINUTE_IN_MS);
@@ -499,30 +502,19 @@
         when(powerStatsStore.loadPowerStatsSpan(1, BatteryUsageStatsSection.TYPE))
                 .thenReturn(span1);
 
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context,
-                batteryStats, powerStatsStore);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context, null,
+                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), powerStatsStore,
+                mMockClock);
 
         BatteryUsageStatsQuery query = new BatteryUsageStatsQuery.Builder()
                 .aggregateSnapshots(0, 3000)
                 .build();
-        final BatteryUsageStats stats = provider.getBatteryUsageStats(query);
+        final BatteryUsageStats stats = provider.getBatteryUsageStats(batteryStats, query);
         assertThat(stats.getCustomPowerComponentNames())
                 .isEqualTo(batteryStats.getCustomEnergyConsumerNames());
         assertThat(stats.getStatsDuration()).isEqualTo(1234);
     }
 
-    private static class TestHandler extends Handler {
-        TestHandler() {
-            super(Looper.getMainLooper());
-        }
-
-        @Override
-        public boolean sendMessageAtTime(Message msg, long uptimeMillis) {
-            msg.getCallback().run();
-            return true;
-        }
-    }
-
     private static final Random sRandom = new Random();
 
     /**
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
index 0b10954..e61dd0b 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsRule.java
@@ -28,6 +28,8 @@
 import android.os.BatteryStats;
 import android.os.BatteryUsageStats;
 import android.os.BatteryUsageStatsQuery;
+import android.os.Handler;
+import android.os.HandlerThread;
 import android.os.UidBatteryConsumer;
 import android.os.UserBatteryConsumer;
 import android.util.SparseArray;
@@ -57,6 +59,7 @@
     private final PowerProfile mPowerProfile;
     private final MockClock mMockClock = new MockClock();
     private final MockBatteryStatsImpl mBatteryStats;
+    private final Handler mHandler;
 
     private BatteryUsageStats mBatteryUsageStats;
     private boolean mScreenOn;
@@ -73,10 +76,13 @@
     }
 
     public BatteryUsageStatsRule(long currentTime, File historyDir) {
+        HandlerThread bgThread = new HandlerThread("bg thread");
+        bgThread.start();
+        mHandler = new Handler(bgThread.getLooper());
         mContext = InstrumentationRegistry.getContext();
         mPowerProfile = spy(new PowerProfile(mContext, true /* forTest */));
         mMockClock.currentTime = currentTime;
-        mBatteryStats = new MockBatteryStatsImpl(mMockClock, historyDir);
+        mBatteryStats = new MockBatteryStatsImpl(mMockClock, historyDir, mHandler);
         mBatteryStats.setPowerProfile(mPowerProfile);
 
         mCpusByPolicy.put(0, new int[]{0, 1, 2, 3});
@@ -92,6 +98,10 @@
         return mMockClock;
     }
 
+    public Handler getHandler() {
+        return mHandler;
+    }
+
     public BatteryUsageStatsRule setTestPowerProfile(@XmlRes int xmlId) {
         mPowerProfile.forceInitForTesting(mContext, xmlId);
         return this;
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuAggregatedPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuAggregatedPowerStatsProcessorTest.java
index 79084cc..8ca4ff6 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuAggregatedPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuAggregatedPowerStatsProcessorTest.java
@@ -192,7 +192,7 @@
 
     private static class MockPowerComponentAggregatedPowerStats extends
             PowerComponentAggregatedPowerStats {
-        private final CpuPowerStatsCollector.StatsArrayLayout mStatsLayout;
+        private final CpuPowerStatsCollector.CpuStatsArrayLayout mStatsLayout;
         private final PowerStats.Descriptor mDescriptor;
         private HashMap<String, long[]> mDeviceStats = new HashMap<>();
         private HashMap<String, long[]> mUidStats = new HashMap<>();
@@ -203,10 +203,10 @@
         MockPowerComponentAggregatedPowerStats(AggregatedPowerStatsConfig.PowerComponent config,
                 boolean useEnergyConsumers) {
             super(config);
-            mStatsLayout = new CpuPowerStatsCollector.StatsArrayLayout();
+            mStatsLayout = new CpuPowerStatsCollector.CpuStatsArrayLayout();
             mStatsLayout.addDeviceSectionCpuTimeByScalingStep(3);
             mStatsLayout.addDeviceSectionCpuTimeByCluster(2);
-            mStatsLayout.addDeviceSectionUptime();
+            mStatsLayout.addDeviceSectionUsageDuration();
             if (useEnergyConsumers) {
                 mStatsLayout.addDeviceSectionEnergyConsumers(2);
             }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java
index bc211df..64d5414 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsCollectorTest.java
@@ -19,6 +19,7 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
@@ -56,6 +57,9 @@
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class CpuPowerStatsCollectorTest {
+    private static final int ISOLATED_UID = 99123;
+    private static final int UID_1 = 42;
+    private static final int UID_2 = 99;
     private Context mContext;
     private final MockClock mMockClock = new MockClock();
     private final HandlerThread mHandlerThread = new HandlerThread("test");
@@ -63,6 +67,8 @@
     private PowerStats mCollectedStats;
     private PowerProfile mPowerProfile;
     @Mock
+    private PowerStatsUidResolver mUidResolver;
+    @Mock
     private CpuPowerStatsCollector.KernelCpuStatsReader mMockKernelCpuStatsReader;
     @Mock
     private PowerStatsInternal mPowerStatsInternal;
@@ -76,6 +82,14 @@
         mHandlerThread.start();
         mHandler = mHandlerThread.getThreadHandler();
         when(mMockKernelCpuStatsReader.nativeIsSupportedFeature()).thenReturn(true);
+        when(mUidResolver.mapUid(anyInt())).thenAnswer(invocation -> {
+            int uid = invocation.getArgument(0);
+            if (uid == ISOLATED_UID) {
+                return UID_2;
+            } else {
+                return uid;
+            }
+        });
     }
 
     @Test
@@ -156,14 +170,14 @@
         assertThat(descriptor.name).isEqualTo("cpu");
         assertThat(descriptor.statsArrayLength).isEqualTo(13);
         assertThat(descriptor.uidStatsArrayLength).isEqualTo(5);
-        CpuPowerStatsCollector.StatsArrayLayout layout =
-                new CpuPowerStatsCollector.StatsArrayLayout();
+        CpuPowerStatsCollector.CpuStatsArrayLayout layout =
+                new CpuPowerStatsCollector.CpuStatsArrayLayout();
         layout.fromExtras(descriptor.extras);
 
         long[] deviceStats = new long[descriptor.statsArrayLength];
         layout.setTimeByScalingStep(deviceStats, 2, 42);
         layout.setConsumedEnergy(deviceStats, 1, 43);
-        layout.setUptime(deviceStats, 44);
+        layout.setUsageDuration(deviceStats, 44);
         layout.setDevicePowerEstimate(deviceStats, 45);
 
         long[] uidStats = new long[descriptor.uidStatsArrayLength];
@@ -173,10 +187,10 @@
         assertThat(layout.getCpuScalingStepCount()).isEqualTo(7);
         assertThat(layout.getTimeByScalingStep(deviceStats, 2)).isEqualTo(42);
 
-        assertThat(layout.getCpuClusterEnergyConsumerCount()).isEqualTo(2);
+        assertThat(layout.getEnergyConsumerCount()).isEqualTo(2);
         assertThat(layout.getConsumedEnergy(deviceStats, 1)).isEqualTo(43);
 
-        assertThat(layout.getUptime(deviceStats)).isEqualTo(44);
+        assertThat(layout.getUsageDuration(deviceStats)).isEqualTo(44);
 
         assertThat(layout.getDevicePowerEstimate(deviceStats)).isEqualTo(45);
 
@@ -195,14 +209,15 @@
         mockEnergyConsumers();
 
         CpuPowerStatsCollector collector = createCollector(8, 0);
-        CpuPowerStatsCollector.StatsArrayLayout layout =
-                new CpuPowerStatsCollector.StatsArrayLayout();
+        CpuPowerStatsCollector.CpuStatsArrayLayout layout =
+                new CpuPowerStatsCollector.CpuStatsArrayLayout();
         layout.fromExtras(collector.getPowerStatsDescriptor().extras);
 
         mockKernelCpuStats(new long[]{1111, 2222, 3333},
                 new SparseArray<>() {{
-                    put(42, new long[]{100, 200});
-                    put(99, new long[]{300, 600});
+                    put(UID_1, new long[]{100, 200});
+                    put(UID_2, new long[]{100, 150});
+                    put(ISOLATED_UID, new long[]{200, 450});
                 }}, 0, 1234);
 
         mMockClock.uptime = 1000;
@@ -219,19 +234,19 @@
         assertThat(layout.getConsumedEnergy(mCollectedStats.stats, 0)).isEqualTo(0);
         assertThat(layout.getConsumedEnergy(mCollectedStats.stats, 1)).isEqualTo(0);
 
-        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(42), 0))
+        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(UID_1), 0))
                 .isEqualTo(100);
-        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(42), 1))
+        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(UID_1), 1))
                 .isEqualTo(200);
-        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(99), 0))
+        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(UID_2), 0))
                 .isEqualTo(300);
-        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(99), 1))
+        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(UID_2), 1))
                 .isEqualTo(600);
 
         mockKernelCpuStats(new long[]{5555, 4444, 3333},
                 new SparseArray<>() {{
-                    put(42, new long[]{123, 234});
-                    put(99, new long[]{345, 678});
+                    put(UID_1, new long[]{123, 234});
+                    put(ISOLATED_UID, new long[]{245, 528});
                 }}, 1234, 3421);
 
         mMockClock.uptime = 2000;
@@ -249,13 +264,13 @@
         // 700 * 1000 / 3500
         assertThat(layout.getConsumedEnergy(mCollectedStats.stats, 1)).isEqualTo(200);
 
-        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(42), 0))
+        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(UID_1), 0))
                 .isEqualTo(23);
-        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(42), 1))
+        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(UID_1), 1))
                 .isEqualTo(34);
-        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(99), 0))
+        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(UID_2), 0))
                 .isEqualTo(45);
-        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(99), 1))
+        assertThat(layout.getUidTimeByPowerBracket(mCollectedStats.uidStats.get(UID_2), 1))
                 .isEqualTo(78);
     }
 
@@ -282,9 +297,9 @@
     private CpuPowerStatsCollector createCollector(int defaultCpuPowerBrackets,
             int defaultCpuPowerBracketsPerEnergyConsumer) {
         CpuPowerStatsCollector collector = new CpuPowerStatsCollector(mCpuScalingPolicies,
-                mPowerProfile, mHandler, mMockKernelCpuStatsReader, () -> mPowerStatsInternal,
-                () -> 3500, 60_000, mMockClock, defaultCpuPowerBrackets,
-                defaultCpuPowerBracketsPerEnergyConsumer);
+                mPowerProfile, mHandler, mMockKernelCpuStatsReader, mUidResolver,
+                () -> mPowerStatsInternal, () -> 3500, 60_000, mMockClock,
+                defaultCpuPowerBrackets, defaultCpuPowerBracketsPerEnergyConsumer);
         collector.addConsumer(stats -> mCollectedStats = stats);
         collector.setEnabled(true);
         return collector;
@@ -375,8 +390,8 @@
     }
 
     private static int[] getScalingStepToPowerBracketMap(CpuPowerStatsCollector collector) {
-        CpuPowerStatsCollector.StatsArrayLayout layout =
-                new CpuPowerStatsCollector.StatsArrayLayout();
+        CpuPowerStatsCollector.CpuStatsArrayLayout layout =
+                new CpuPowerStatsCollector.CpuStatsArrayLayout();
         layout.fromExtras(collector.getPowerStatsDescriptor().extras);
         return layout.getScalingStepToPowerBracketMap();
     }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java b/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
index 4150972a..fb71ac8 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
@@ -61,16 +61,23 @@
     }
 
     MockBatteryStatsImpl(Clock clock, File historyDirectory) {
-        super(clock, historyDirectory);
+        this(clock, historyDirectory, new Handler(Looper.getMainLooper()));
+    }
+
+    MockBatteryStatsImpl(Clock clock, File historyDirectory, Handler handler) {
+        this(clock, historyDirectory, handler, new PowerStatsUidResolver());
+    }
+
+    MockBatteryStatsImpl(Clock clock, File historyDirectory, Handler handler,
+            PowerStatsUidResolver powerStatsUidResolver) {
+        super(clock, historyDirectory, handler, powerStatsUidResolver);
         initTimersAndCounters();
         setMaxHistoryBuffer(128 * 1024);
 
         setExternalStatsSyncLocked(mExternalStatsSync);
         informThatAllExternalStatsAreFlushed();
 
-        // A no-op handler.
-        mHandler = new Handler(Looper.getMainLooper()) {
-        };
+        mHandler = handler;
 
         mCpuUidFreqTimeReader = mock(KernelCpuUidFreqTimeReader.class);
         mKernelWakelockReader = null;
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsAggregatorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsAggregatorTest.java
index b52fc8a..6704987 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsAggregatorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsAggregatorTest.java
@@ -76,9 +76,18 @@
 
     @Test
     public void stateUpdates() {
+        PowerStats.Descriptor descriptor =
+                new PowerStats.Descriptor(TEST_POWER_COMPONENT, "majorDrain", 1, 1,
+                        new PersistableBundle());
+        PowerStats powerStats = new PowerStats(descriptor);
+
         mClock.currentTime = 1222156800000L;    // An important date in world history
 
         mHistory.forceRecordAllHistory();
+        powerStats.stats = new long[]{0};
+        powerStats.uidStats.put(TEST_UID, new long[]{0});
+        mHistory.recordPowerStats(mClock.realtime, mClock.uptime, powerStats);
+
         mHistory.recordBatteryState(mClock.realtime, mClock.uptime, 10, /* plugged */ true);
         mHistory.recordStateStartEvent(mClock.realtime, mClock.uptime,
                 BatteryStats.HistoryItem.STATE_SCREEN_ON_FLAG);
@@ -87,10 +96,6 @@
 
         advance(1000);
 
-        PowerStats.Descriptor descriptor =
-                new PowerStats.Descriptor(TEST_POWER_COMPONENT, "majorDrain", 1, 1,
-                        new PersistableBundle());
-        PowerStats powerStats = new PowerStats(descriptor);
         powerStats.stats = new long[]{10000};
         powerStats.uidStats.put(TEST_UID, new long[]{1234});
         mHistory.recordPowerStats(mClock.realtime, mClock.uptime, powerStats);
@@ -181,17 +186,21 @@
 
     @Test
     public void incompatiblePowerStats() {
+        PowerStats.Descriptor descriptor =
+                new PowerStats.Descriptor(TEST_POWER_COMPONENT, "majorDrain", 1, 1,
+                        new PersistableBundle());
+        PowerStats powerStats = new PowerStats(descriptor);
+
         mHistory.forceRecordAllHistory();
+        powerStats.stats = new long[]{0};
+        powerStats.uidStats.put(TEST_UID, new long[]{0});
+        mHistory.recordPowerStats(mClock.realtime, mClock.uptime, powerStats);
         mHistory.recordBatteryState(mClock.realtime, mClock.uptime, 10, /* plugged */ true);
         mHistory.recordProcessStateChange(mClock.realtime, mClock.uptime, TEST_UID,
                 BatteryConsumer.PROCESS_STATE_FOREGROUND);
 
         advance(1000);
 
-        PowerStats.Descriptor descriptor =
-                new PowerStats.Descriptor(TEST_POWER_COMPONENT, "majorDrain", 1, 1,
-                        new PersistableBundle());
-        PowerStats powerStats = new PowerStats(descriptor);
         powerStats.stats = new long[]{10000};
         powerStats.uidStats.put(TEST_UID, new long[]{1234});
         mHistory.recordPowerStats(mClock.realtime, mClock.uptime, powerStats);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java
new file mode 100644
index 0000000..3c48262
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java
@@ -0,0 +1,316 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import static androidx.test.InstrumentationRegistry.getContext;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
+
+import static org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.mock;
+
+import android.os.AggregateBatteryConsumer;
+import android.os.BatteryConsumer;
+import android.os.BatteryStats;
+import android.os.BatteryUsageStats;
+import android.os.Handler;
+import android.os.Looper;
+import android.os.Message;
+import android.os.Parcel;
+import android.os.PersistableBundle;
+import android.os.UidBatteryConsumer;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.internal.os.BatteryStatsHistory;
+import com.android.internal.os.MonotonicClock;
+import com.android.internal.os.PowerProfile;
+import com.android.internal.os.PowerStats;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.util.List;
+
+@RunWith(AndroidJUnit4.class)
+public class PowerStatsExporterTest {
+
+    private static final int APP_UID1 = 42;
+    private static final int APP_UID2 = 84;
+    private static final double TOLERANCE = 0.01;
+
+    @Rule
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .setAveragePower(PowerProfile.POWER_CPU_ACTIVE, 720)
+            .setCpuScalingPolicy(0, new int[]{0}, new int[]{100})
+            .setAveragePowerForCpuScalingPolicy(0, 360)
+            .setAveragePowerForCpuScalingStep(0, 0, 300)
+            .setCpuPowerBracketCount(1)
+            .setCpuPowerBracket(0, 0, 0);
+
+    private MockClock mClock = new MockClock();
+    private MonotonicClock mMonotonicClock = new MonotonicClock(0, mClock);
+    private PowerStatsStore mPowerStatsStore;
+    private PowerStatsAggregator mPowerStatsAggregator;
+    private BatteryStatsHistory mHistory;
+    private CpuPowerStatsCollector.CpuStatsArrayLayout mCpuStatsArrayLayout;
+    private PowerStats.Descriptor mPowerStatsDescriptor;
+
+    @Before
+    public void setup() {
+        File storeDirectory = new File(getContext().getCacheDir(), getClass().getSimpleName());
+        clearDirectory(storeDirectory);
+
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_CPU)
+                .trackDeviceStates(AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessor(
+                        new CpuAggregatedPowerStatsProcessor(mStatsRule.getPowerProfile(),
+                                mStatsRule.getCpuScalingPolicies()));
+
+        mPowerStatsStore = new PowerStatsStore(storeDirectory, new TestHandler(), config);
+        mHistory = new BatteryStatsHistory(Parcel.obtain(), storeDirectory, 0, 10000,
+                mock(BatteryStatsHistory.HistoryStepDetailsCalculator.class), mClock,
+                mMonotonicClock, null);
+        mPowerStatsAggregator = new PowerStatsAggregator(config, mHistory);
+
+        mCpuStatsArrayLayout = new CpuPowerStatsCollector.CpuStatsArrayLayout();
+        mCpuStatsArrayLayout.addDeviceSectionCpuTimeByScalingStep(1);
+        mCpuStatsArrayLayout.addDeviceSectionCpuTimeByCluster(1);
+        mCpuStatsArrayLayout.addDeviceSectionPowerEstimate();
+        mCpuStatsArrayLayout.addUidSectionCpuTimeByPowerBracket(new int[]{0});
+        mCpuStatsArrayLayout.addUidSectionPowerEstimate();
+        PersistableBundle extras = new PersistableBundle();
+        mCpuStatsArrayLayout.toExtras(extras);
+
+        mPowerStatsDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_CPU,
+                mCpuStatsArrayLayout.getDeviceStatsArrayLength(),
+                mCpuStatsArrayLayout.getUidStatsArrayLength(), extras);
+    }
+
+    @Test
+    public void breakdownByProcState_fullRange() throws Exception {
+        BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(
+                new String[0], /* includePowerModels */ false,
+                /* includeProcessStateData */ true, /* powerThreshold */ 0);
+        exportAggregatedPowerStats(builder, 1000, 10000);
+
+        BatteryUsageStats actual = builder.build();
+        String message = "Actual BatteryUsageStats: " + actual;
+
+        assertDevicePowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 25.53);
+        assertAllAppsPowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 25.53);
+
+        assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_ANY, 13.5);
+        assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND, 7.47);
+        assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND, 6.03);
+
+        assertUidPowerEstimate(message, actual, APP_UID2, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_ANY, 12.03);
+        assertUidPowerEstimate(message, actual, APP_UID2, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE, 12.03);
+
+        actual.close();
+    }
+
+    @Test
+    public void breakdownByProcState_subRange() throws Exception {
+        BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(
+                new String[0], /* includePowerModels */ false,
+                /* includeProcessStateData */ true, /* powerThreshold */ 0);
+        exportAggregatedPowerStats(builder, 3700, 6700);
+
+        BatteryUsageStats actual = builder.build();
+        String message = "Actual BatteryUsageStats: " + actual;
+
+        assertDevicePowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 15.4);
+        assertAllAppsPowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 15.4);
+
+        assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_ANY, 4.06);
+        assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND, 1.35);
+        assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND, 2.70);
+
+        assertUidPowerEstimate(message, actual, APP_UID2, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_ANY, 11.33);
+        assertUidPowerEstimate(message, actual, APP_UID2, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE, 11.33);
+
+        actual.close();
+    }
+
+    @Test
+    public void combinedProcessStates() throws Exception {
+        BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(
+                new String[0], /* includePowerModels */ false,
+                /* includeProcessStateData */ false, /* powerThreshold */ 0);
+        exportAggregatedPowerStats(builder, 1000, 10000);
+
+        BatteryUsageStats actual = builder.build();
+        String message = "Actual BatteryUsageStats: " + actual;
+
+        assertDevicePowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 25.53);
+        assertAllAppsPowerEstimate(message, actual, BatteryConsumer.POWER_COMPONENT_CPU, 25.53);
+
+        assertUidPowerEstimate(message, actual, APP_UID1, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_ANY, 13.5);
+        assertUidPowerEstimate(message, actual, APP_UID2, BatteryConsumer.POWER_COMPONENT_CPU,
+                BatteryConsumer.PROCESS_STATE_ANY, 12.03);
+        UidBatteryConsumer uidScope = actual.getUidBatteryConsumers().stream()
+                .filter(us -> us.getUid() == APP_UID1).findFirst().orElse(null);
+        // There shouldn't be any per-procstate data
+        assertThrows(
+                IllegalArgumentException.class,
+                () -> uidScope.getConsumedPower(new BatteryConsumer.Dimensions(
+                        BatteryConsumer.POWER_COMPONENT_CPU,
+                        BatteryConsumer.PROCESS_STATE_FOREGROUND)));
+
+
+        actual.close();
+    }
+
+    private void recordBatteryHistory() {
+        PowerStats powerStats = new PowerStats(mPowerStatsDescriptor);
+        long[] uidStats1 = new long[mCpuStatsArrayLayout.getUidStatsArrayLength()];
+        powerStats.uidStats.put(APP_UID1, uidStats1);
+        long[] uidStats2 = new long[mCpuStatsArrayLayout.getUidStatsArrayLength()];
+        powerStats.uidStats.put(APP_UID2, uidStats2);
+
+        mHistory.forceRecordAllHistory();
+
+        mHistory.startRecordingHistory(1000, 1000, false);
+        mHistory.recordPowerStats(1000, 1000, powerStats);
+        mHistory.recordBatteryState(1000, 1000, 70, /* plugged */ false);
+        mHistory.recordStateStartEvent(1000, 1000, BatteryStats.HistoryItem.STATE_SCREEN_ON_FLAG);
+        mHistory.recordProcessStateChange(1000, 1000, APP_UID1,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND);
+        mHistory.recordProcessStateChange(1000, 1000, APP_UID2,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE);
+
+        mCpuStatsArrayLayout.setTimeByScalingStep(powerStats.stats, 0, 11111);
+        mCpuStatsArrayLayout.setUidTimeByPowerBracket(uidStats1, 0, 10000);
+        mCpuStatsArrayLayout.setUidTimeByPowerBracket(uidStats2, 0, 1111);
+        mHistory.recordPowerStats(1000, 1000, powerStats);
+
+        mCpuStatsArrayLayout.setTimeByScalingStep(powerStats.stats, 0, 12345);
+        mCpuStatsArrayLayout.setUidTimeByPowerBracket(uidStats1, 0, 9876);
+        mCpuStatsArrayLayout.setUidTimeByPowerBracket(uidStats2, 0, 2469);
+        mHistory.recordPowerStats(3000, 3000, powerStats);
+
+        mPowerStatsAggregator.aggregatePowerStats(0, 3500, stats -> {
+            mPowerStatsStore.storeAggregatedPowerStats(stats);
+        });
+
+        mHistory.recordProcessStateChange(4000, 4000, APP_UID1,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND);
+
+        mHistory.recordStateStopEvent(4000, 4000, BatteryStats.HistoryItem.STATE_SCREEN_ON_FLAG);
+
+        mCpuStatsArrayLayout.setTimeByScalingStep(powerStats.stats, 0, 54321);
+        mCpuStatsArrayLayout.setUidTimeByPowerBracket(uidStats1, 0, 14321);
+        mCpuStatsArrayLayout.setUidTimeByPowerBracket(uidStats2, 0, 40000);
+        mHistory.recordPowerStats(6000, 6000, powerStats);
+
+        mPowerStatsAggregator.aggregatePowerStats(3500, 6500, stats -> {
+            mPowerStatsStore.storeAggregatedPowerStats(stats);
+        });
+
+        mHistory.recordStateStartEvent(7000, 7000, BatteryStats.HistoryItem.STATE_SCREEN_ON_FLAG);
+        mHistory.recordProcessStateChange(7000, 7000, APP_UID1,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND);
+        mHistory.recordProcessStateChange(7000, 7000, APP_UID2,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED);
+
+        mCpuStatsArrayLayout.setTimeByScalingStep(powerStats.stats, 0, 23456);
+        mCpuStatsArrayLayout.setUidTimeByPowerBracket(uidStats1, 0, 23456);
+        mCpuStatsArrayLayout.setUidTimeByPowerBracket(uidStats2, 0, 0);
+        mHistory.recordPowerStats(8000, 8000, powerStats);
+
+        assertThat(mPowerStatsStore.getTableOfContents()).hasSize(2);
+    }
+
+    private void exportAggregatedPowerStats(BatteryUsageStats.Builder builder,
+            int monotonicStartTime, int monotonicEndTime) {
+        recordBatteryHistory();
+        PowerStatsExporter exporter = new PowerStatsExporter(mPowerStatsStore,
+                mPowerStatsAggregator, /* batterySessionTimeSpanSlackMillis */ 0);
+        exporter.exportAggregatedPowerStats(builder, monotonicStartTime, monotonicEndTime);
+    }
+
+    private void assertDevicePowerEstimate(String message, BatteryUsageStats bus, int componentId,
+            double expected) {
+        AggregateBatteryConsumer consumer = bus.getAggregateBatteryConsumer(
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE);
+        assertWithMessage(message).that(consumer.getConsumedPower(componentId))
+                .isWithin(TOLERANCE).of(expected);
+    }
+
+    private void assertAllAppsPowerEstimate(String message, BatteryUsageStats bus, int componentId,
+            double expected) {
+        AggregateBatteryConsumer consumer = bus.getAggregateBatteryConsumer(
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS);
+        assertWithMessage(message).that(consumer.getConsumedPower(componentId))
+                .isWithin(TOLERANCE).of(expected);
+    }
+
+    private void assertUidPowerEstimate(String message, BatteryUsageStats bus, int uid,
+            int componentId, int processState, double expected) {
+        List<UidBatteryConsumer> uidScopes = bus.getUidBatteryConsumers();
+        final UidBatteryConsumer uidScope = uidScopes.stream()
+                .filter(us -> us.getUid() == uid).findFirst().orElse(null);
+        assertWithMessage(message).that(uidScope).isNotNull();
+        assertWithMessage(message).that(uidScope.getConsumedPower(
+                new BatteryConsumer.Dimensions(componentId, processState)))
+                .isWithin(TOLERANCE).of(expected);
+    }
+
+    private void clearDirectory(File dir) {
+        if (dir.exists()) {
+            for (File child : dir.listFiles()) {
+                if (child.isDirectory()) {
+                    clearDirectory(child);
+                }
+                child.delete();
+            }
+        }
+    }
+
+    private static class TestHandler extends Handler {
+        TestHandler() {
+            super(Looper.getMainLooper());
+        }
+
+        @Override
+        public boolean sendMessageAtTime(Message msg, long uptimeMillis) {
+            msg.getCallback().run();
+            return true;
+        }
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsSchedulerTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsSchedulerTest.java
index 0e58787..7257a94 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsSchedulerTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsSchedulerTest.java
@@ -27,9 +27,6 @@
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
-import android.os.BatteryConsumer;
-import android.os.BatteryManager;
-import android.os.BatteryUsageStats;
 import android.os.ConditionVariable;
 import android.os.Handler;
 import android.os.HandlerThread;
@@ -59,13 +56,13 @@
     private MockClock mClock = new MockClock();
     private MonotonicClock mMonotonicClock = new MonotonicClock(0, mClock);
     private MockBatteryStatsImpl mBatteryStats;
-    private BatteryUsageStatsProvider mBatteryUsageStatsProvider;
     private PowerStatsScheduler mPowerStatsScheduler;
     private PowerProfile mPowerProfile;
     private PowerStatsAggregator mPowerStatsAggregator;
     private AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
 
     @Before
+    @SuppressWarnings("GuardedBy")
     public void setup() {
         final Context context = InstrumentationRegistry.getContext();
 
@@ -83,11 +80,10 @@
         mPowerProfile = mock(PowerProfile.class);
         when(mPowerProfile.getAveragePower(PowerProfile.POWER_FLASHLIGHT)).thenReturn(1000000.0);
         mBatteryStats = new MockBatteryStatsImpl(mClock).setPowerProfile(mPowerProfile);
-        mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context, mBatteryStats);
         mPowerStatsAggregator = mock(PowerStatsAggregator.class);
         mPowerStatsScheduler = new PowerStatsScheduler(context, mPowerStatsAggregator,
                 TimeUnit.MINUTES.toMillis(30), TimeUnit.HOURS.toMillis(1), mPowerStatsStore, mClock,
-                mMonotonicClock, mHandler, mBatteryStats, mBatteryUsageStatsProvider);
+                mMonotonicClock, mHandler, mBatteryStats);
     }
 
     @Test
@@ -176,70 +172,6 @@
     }
 
     @Test
-    public void storeBatteryUsageStatsOnReset() {
-        mBatteryStats.forceRecordAllHistory();
-        synchronized (mBatteryStats) {
-            mBatteryStats.setOnBatteryLocked(mClock.realtime, mClock.uptime, true,
-                    BatteryManager.BATTERY_STATUS_DISCHARGING, 50, 0);
-        }
-
-        mPowerStatsScheduler.start(/* schedulePeriodicPowerStatsCollection */false);
-
-        assertThat(mPowerStatsStore.getTableOfContents()).isEmpty();
-
-        mPowerStatsScheduler.start(true);
-
-        synchronized (mBatteryStats) {
-            mBatteryStats.noteFlashlightOnLocked(42, mClock.realtime, mClock.uptime);
-        }
-
-        mClock.realtime += 60000;
-        mClock.currentTime += 60000;
-
-        synchronized (mBatteryStats) {
-            mBatteryStats.noteFlashlightOffLocked(42, mClock.realtime, mClock.uptime);
-        }
-
-        mClock.realtime += 60000;
-        mClock.currentTime += 60000;
-
-        // Battery stats reset should have the side-effect of saving accumulated battery usage stats
-        synchronized (mBatteryStats) {
-            mBatteryStats.resetAllStatsAndHistoryLocked(BatteryStatsImpl.RESET_REASON_ADB_COMMAND);
-        }
-
-        // Await completion
-        ConditionVariable done = new ConditionVariable();
-        mHandler.post(done::open);
-        done.block();
-
-        List<PowerStatsSpan.Metadata> contents = mPowerStatsStore.getTableOfContents();
-        assertThat(contents).hasSize(1);
-
-        PowerStatsSpan.Metadata metadata = contents.get(0);
-
-        PowerStatsSpan span = mPowerStatsStore.loadPowerStatsSpan(metadata.getId(),
-                BatteryUsageStatsSection.TYPE);
-        assertThat(span).isNotNull();
-
-        List<PowerStatsSpan.TimeFrame> timeFrames = span.getMetadata().getTimeFrames();
-        assertThat(timeFrames).hasSize(1);
-        assertThat(timeFrames.get(0).startMonotonicTime).isEqualTo(7654321);
-        assertThat(timeFrames.get(0).duration).isEqualTo(120000);
-
-        List<PowerStatsSpan.Section> sections = span.getSections();
-        assertThat(sections).hasSize(1);
-
-        PowerStatsSpan.Section section = sections.get(0);
-        assertThat(section.getType()).isEqualTo(BatteryUsageStatsSection.TYPE);
-        BatteryUsageStats bus = ((BatteryUsageStatsSection) section).getBatteryUsageStats();
-        assertThat(bus.getAggregateBatteryConsumer(
-                        BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
-                .getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
-                .isEqualTo(60000);
-    }
-
-    @Test
     public void alignToWallClock() {
         // Expect the aligned value to be adjusted by 1 min 30 sec - rounded to the next 15 min
         assertThat(PowerStatsScheduler.alignToWallClock(123, TimeUnit.MINUTES.toMillis(15),
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsUidResolverTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsUidResolverTest.java
new file mode 100644
index 0000000..60b2541
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsUidResolverTest.java
@@ -0,0 +1,115 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@RunWith(AndroidJUnit4.class)
+public class PowerStatsUidResolverTest {
+
+    private final PowerStatsUidResolver mResolver = new PowerStatsUidResolver();
+    @Mock
+    PowerStatsUidResolver.Listener mListener;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+    }
+
+    @Test
+    public void addAndRemoveIsolatedUid() {
+        mResolver.addListener(mListener);
+        mResolver.noteIsolatedUidAdded(42, 314);
+        verify(mListener).onIsolatedUidAdded(42, 314);
+        assertThat(mResolver.mapUid(42)).isEqualTo(314);
+
+        mResolver.noteIsolatedUidRemoved(42, 314);
+        verify(mListener).onBeforeIsolatedUidRemoved(42, 314);
+        verify(mListener).onAfterIsolatedUidRemoved(42, 314);
+        assertThat(mResolver.mapUid(42)).isEqualTo(42);
+
+        verifyNoMoreInteractions(mListener);
+    }
+
+    @Test
+    public void retainAndRemoveIsolatedUid() {
+        mResolver.addListener(mListener);
+        mResolver.noteIsolatedUidAdded(42, 314);
+        verify(mListener).onIsolatedUidAdded(42, 314);
+        assertThat(mResolver.mapUid(42)).isEqualTo(314);
+
+        mResolver.retainIsolatedUid(42);
+
+        mResolver.noteIsolatedUidRemoved(42, 314);
+        verify(mListener).onBeforeIsolatedUidRemoved(42, 314);
+        assertThat(mResolver.mapUid(42)).isEqualTo(314);
+        verifyNoMoreInteractions(mListener);
+
+        mResolver.releaseIsolatedUid(42);
+        verify(mListener).onAfterIsolatedUidRemoved(42, 314);
+        assertThat(mResolver.mapUid(42)).isEqualTo(42);
+
+        verifyNoMoreInteractions(mListener);
+    }
+
+    @Test
+    public void removeUidsInRange() {
+        mResolver.noteIsolatedUidAdded(1, 314);
+        mResolver.noteIsolatedUidAdded(2, 314);
+        mResolver.noteIsolatedUidAdded(3, 314);
+        mResolver.noteIsolatedUidAdded(4, 314);
+        mResolver.noteIsolatedUidAdded(6, 314);
+        mResolver.noteIsolatedUidAdded(8, 314);
+        mResolver.noteIsolatedUidAdded(10, 314);
+
+        mResolver.addListener(mListener);
+
+        mResolver.releaseUidsInRange(4, 4);     // Single
+        verify(mListener).onAfterIsolatedUidRemoved(4, 314);
+        verifyNoMoreInteractions(mListener);
+
+        // Now: [1, 2, 3, 6, 8, 10]
+
+        mResolver.releaseUidsInRange(2, 3);     // Inclusive
+        verify(mListener).onAfterIsolatedUidRemoved(2, 314);
+        verify(mListener).onAfterIsolatedUidRemoved(3, 314);
+        verifyNoMoreInteractions(mListener);
+
+        // Now: [1, 6, 8, 10]
+
+        mResolver.releaseUidsInRange(5, 9);     // Exclusive
+        verify(mListener).onAfterIsolatedUidRemoved(6, 314);
+        verify(mListener).onAfterIsolatedUidRemoved(8, 314);
+        verifyNoMoreInteractions(mListener);
+
+        // Now: [1, 10]
+
+        mResolver.releaseUidsInRange(5, 9);     // Empty
+        verifyNoMoreInteractions(mListener);
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/PinnerServiceTest.java b/services/tests/servicestests/src/com/android/server/PinnerServiceTest.java
index 1002fba..88ca029 100644
--- a/services/tests/servicestests/src/com/android/server/PinnerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/PinnerServiceTest.java
@@ -24,6 +24,7 @@
 import static org.mockito.Mockito.spy;
 
 import android.app.ActivityManagerInternal;
+import android.app.pinner.PinnedFileStat;
 import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.content.pm.ResolveInfo;
@@ -47,22 +48,19 @@
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
-import java.io.BufferedReader;
 import java.io.CharArrayWriter;
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
-import java.io.StringReader;
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
-import java.util.Optional;
+import java.util.List;
 import java.util.concurrent.TimeUnit;
 
 @SmallTest
@@ -138,6 +136,13 @@
             protected void publishBinderService(PinnerService service, Binder binderService) {
                 // Suppress this for testing, it's not needed and causes conflitcs.
             }
+
+            @Override
+            protected PinnerService.PinnedFile pinFileInternal(String fileToPin,
+                    int maxBytesToPin, boolean attemptPinIntrospection) {
+                return new PinnerService.PinnedFile(-1,
+                        maxBytesToPin, fileToPin, maxBytesToPin);
+            }
         };
     }
 
@@ -202,20 +207,27 @@
         return cw.toString();
     }
 
-    private int getPinnedSize(PinnerService pinnerService) throws Exception {
-        return getPinnedSizeImpl(pinnerService, "Total size: ");
+    private long getPinnedSize(PinnerService pinnerService) {
+        long totalBytesPinned = 0;
+        for (PinnedFileStat stat : pinnerService.getPinnerStats()) {
+            totalBytesPinned += stat.getBytesPinned();
+        }
+        return totalBytesPinned;
     }
 
-    private int getPinnedAnonSize(PinnerService pinnerService) throws Exception {
-        return getPinnedSizeImpl(pinnerService, "Pinned anon region: ");
+    private int getPinnedAnonSize(PinnerService pinnerService) {
+        List<PinnedFileStat> anonStats = pinnerService.getPinnerStats().stream()
+                .filter(pf -> pf.getGroupName().equals(PinnerService.ANON_REGION_STAT_NAME))
+                .toList();
+        int totalAnon = 0;
+        for (PinnedFileStat anonStat : anonStats) {
+            totalAnon += anonStat.getBytesPinned();
+        }
+        return totalAnon;
     }
 
-    private int getPinnedSizeImpl(PinnerService pinnerService, String sizeToken) throws Exception {
-        String dumpOutput = getPinnerServiceDump(pinnerService);
-        BufferedReader bufReader = new BufferedReader(new StringReader(dumpOutput));
-        Optional<Integer> size = bufReader.lines().filter(s -> s.contains(sizeToken))
-                .map(s -> Integer.valueOf(s.substring(sizeToken.length()))).findAny();
-        return size.orElse(-1);
+    private long getTotalPinnedFiles(PinnerService pinnerService) {
+        return pinnerService.getPinnerStats().stream().count();
     }
 
     private void setDeviceConfigPinnedAnonSize(long size) {
@@ -227,7 +239,6 @@
     }
 
     @Test
-    @Ignore("b/309853498, pinning home app can fail with ENOMEM")
     public void testPinHomeApp() throws Exception {
         // Enable HOME app pinning
         mContext.getOrCreateTestableResources()
@@ -245,15 +256,13 @@
         ArrayMap<Integer, Object> pinnedApps = getPinnedApps(pinnerService);
         assertThat(pinnedApps.get(KEY_HOME)).isNotNull();
 
-        // Check if dump() reports total pinned bytes
-        int totalPinnedSizeBytes = getPinnedSize(pinnerService);
-        assertThat(totalPinnedSizeBytes).isGreaterThan(0);
+        assertThat(getPinnedSize(pinnerService)).isGreaterThan(0);
+        assertThat(getTotalPinnedFiles(pinnerService)).isGreaterThan(0);
 
         unpinAll(pinnerService);
     }
 
     @Test
-    @Ignore("b/309853498, pinning home app can fail with ENOMEM")
     public void testPinHomeAppOnBootCompleted() throws Exception {
         // Enable HOME app pinning
         mContext.getOrCreateTestableResources()
@@ -271,9 +280,7 @@
         ArrayMap<Integer, Object> pinnedApps = getPinnedApps(pinnerService);
         assertThat(pinnedApps.get(KEY_HOME)).isNotNull();
 
-        // Check if dump() reports total pinned bytes
-        int totalPinnedSizeBytes = getPinnedSize(pinnerService);
-        assertThat(totalPinnedSizeBytes).isGreaterThan(0);
+        assertThat(getPinnedSize(pinnerService)).isGreaterThan(0);
 
         unpinAll(pinnerService);
     }
@@ -294,12 +301,24 @@
         ArrayMap<Integer, Object> pinnedApps = getPinnedApps(pinnerService);
         assertThat(pinnedApps).isEmpty();
 
-        // Check if dump() reports total pinned bytes
-        int totalPinnedSizeBytes = getPinnedSize(pinnerService);
+        long totalPinnedSizeBytes = getPinnedSize(pinnerService);
         assertThat(totalPinnedSizeBytes).isEqualTo(0);
 
         int pinnedAnonSizeBytes = getPinnedAnonSize(pinnerService);
-        assertThat(pinnedAnonSizeBytes).isEqualTo(-1);
+        assertThat(pinnedAnonSizeBytes).isEqualTo(0);
+
+        unpinAll(pinnerService);
+    }
+
+    @Test
+    public void testPinFile() throws Exception {
+        PinnerService pinnerService = new PinnerService(mContext, mInjector);
+        pinnerService.onStart();
+
+        pinnerService.pinFile("test_file", 4096, null, "my_group");
+
+        assertThat(getPinnedSize(pinnerService)).isGreaterThan(0);
+        assertThat(getTotalPinnedFiles(pinnerService)).isGreaterThan(0);
 
         unpinAll(pinnerService);
     }
@@ -341,11 +360,8 @@
         waitForPinnerService(pinnerService);
         // An empty anon region should clear the associated status entry.
         pinnedAnonSizeBytes = getPinnedAnonSize(pinnerService);
-        assertThat(pinnedAnonSizeBytes).isEqualTo(-1);
+        assertThat(pinnedAnonSizeBytes).isEqualTo(0);
 
         unpinAll(pinnerService);
     }
-
-    // TODO: Add test to check that the pages we expect to be pinned are actually pinned
-
 }
diff --git a/services/tests/servicestests/src/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandlerTest.java b/services/tests/servicestests/src/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandlerTest.java
index fd2cf6d..3b39160 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandlerTest.java
+++ b/services/tests/servicestests/src/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandlerTest.java
@@ -540,18 +540,23 @@
         twoFingerTap();
 
         assertIn(STATE_ACTIVATED);
+        verify(mMockMagnificationLogger, never()).logMagnificationTripleTap(anyBoolean());
+        verify(mMockMagnificationLogger).logMagnificationTwoFingerTripleTap(true);
     }
 
     @Test
     @RequiresFlagsEnabled(Flags.FLAG_ENABLE_MAGNIFICATION_MULTIPLE_FINGER_MULTIPLE_TAP_GESTURE)
     public void testTwoFingerTripleTap_StateIsActivated_shouldInIdle() {
         goFromStateIdleTo(STATE_ACTIVATED);
+        reset(mMockMagnificationLogger);
 
         twoFingerTap();
         twoFingerTap();
         twoFingerTap();
 
         assertIn(STATE_IDLE);
+        verify(mMockMagnificationLogger, never()).logMagnificationTripleTap(anyBoolean());
+        verify(mMockMagnificationLogger).logMagnificationTwoFingerTripleTap(false);
     }
 
     @Test
@@ -564,6 +569,8 @@
         twoFingerTapAndHold();
 
         assertIn(STATE_NON_ACTIVATED_ZOOMED_TMP);
+        verify(mMockMagnificationLogger, never()).logMagnificationTripleTap(anyBoolean());
+        verify(mMockMagnificationLogger).logMagnificationTwoFingerTripleTap(true);
     }
 
     @Test
@@ -576,6 +583,8 @@
         twoFingerSwipeAndHold();
 
         assertIn(STATE_NON_ACTIVATED_ZOOMED_TMP);
+        verify(mMockMagnificationLogger, never()).logMagnificationTripleTap(anyBoolean());
+        verify(mMockMagnificationLogger).logMagnificationTwoFingerTripleTap(true);
     }
 
     @Test
diff --git a/tests/CtsSurfaceControlTestsStaging/src/main/java/android/view/surfacecontroltests/GraphicsActivity.java b/tests/CtsSurfaceControlTestsStaging/src/main/java/android/view/surfacecontroltests/GraphicsActivity.java
index 4548a7d..1e5f33f 100644
--- a/tests/CtsSurfaceControlTestsStaging/src/main/java/android/view/surfacecontroltests/GraphicsActivity.java
+++ b/tests/CtsSurfaceControlTestsStaging/src/main/java/android/view/surfacecontroltests/GraphicsActivity.java
@@ -710,9 +710,15 @@
             float childFrameRate = Collections.max(frameRates);
             float parentFrameRate = childFrameRate / 2;
             int initialNumEvents = mModeChangedEvents.size();
-            parent.setFrameRate(parentFrameRate);
             parent.setFrameRateSelectionStrategy(parentStrategy);
-            child.setFrameRate(childFrameRate);
+
+            // For Self case, we want to test that child gets default behavior
+            if (parentStrategy == SurfaceControl.FRAME_RATE_SELECTION_STRATEGY_SELF) {
+                parent.setFrameRateCategory(Surface.FRAME_RATE_CATEGORY_NO_PREFERENCE);
+            } else {
+                parent.setFrameRate(parentFrameRate);
+                child.setFrameRate(childFrameRate);
+            }
 
             // Verify
             float expectedFrameRate =
diff --git a/tests/CtsSurfaceControlTestsStaging/src/main/java/android/view/surfacecontroltests/SurfaceControlTest.java b/tests/CtsSurfaceControlTestsStaging/src/main/java/android/view/surfacecontroltests/SurfaceControlTest.java
index bed9cff..29f6879 100644
--- a/tests/CtsSurfaceControlTestsStaging/src/main/java/android/view/surfacecontroltests/SurfaceControlTest.java
+++ b/tests/CtsSurfaceControlTestsStaging/src/main/java/android/view/surfacecontroltests/SurfaceControlTest.java
@@ -16,11 +16,8 @@
 
 package android.view.surfacecontroltests;
 
-import static org.junit.Assume.assumeTrue;
-
 import android.Manifest;
 import android.hardware.display.DisplayManager;
-import android.os.SystemProperties;
 import android.support.test.uiautomator.UiDevice;
 import android.view.Surface;
 import android.view.SurfaceControl;
@@ -54,10 +51,6 @@
 
         UiDevice uiDevice = UiDevice.getInstance(InstrumentationRegistry.getInstrumentation());
 
-        // TODO(b/290634611): clean this up once SF new front end is enabled by default
-        assumeTrue(SystemProperties.getBoolean(
-                "persist.debug.sf.enable_layer_lifecycle_manager", false));
-
         uiDevice.wakeUp();
         uiDevice.executeShellCommand("wm dismiss-keyguard");
 
@@ -118,10 +111,11 @@
     }
 
     @Test
-    public void testSurfaceControlFrameRateSelectionStrategySelf() throws InterruptedException {
+    public void testSurfaceControlFrameRateSelectionStrategyPropagate()
+            throws InterruptedException {
         GraphicsActivity activity = mActivityRule.getActivity();
         activity.testSurfaceControlFrameRateSelectionStrategy(
-                SurfaceControl.FRAME_RATE_SELECTION_STRATEGY_SELF);
+                SurfaceControl.FRAME_RATE_SELECTION_STRATEGY_PROPAGATE);
     }
 
     @Test
@@ -131,4 +125,12 @@
         activity.testSurfaceControlFrameRateSelectionStrategy(
                 SurfaceControl.FRAME_RATE_SELECTION_STRATEGY_OVERRIDE_CHILDREN);
     }
+
+    @Test
+    public void testSurfaceControlFrameRateSelectionStrategySelf()
+            throws InterruptedException {
+        GraphicsActivity activity = mActivityRule.getActivity();
+        activity.testSurfaceControlFrameRateSelectionStrategy(
+                SurfaceControl.FRAME_RATE_SELECTION_STRATEGY_SELF);
+    }
 }
diff --git a/tools/hoststubgen/hoststubgen/invoketest/hoststubgen-invoke-test.sh b/tools/hoststubgen/hoststubgen/invoketest/hoststubgen-invoke-test.sh
index 85038be..91e6814 100755
--- a/tools/hoststubgen/hoststubgen/invoketest/hoststubgen-invoke-test.sh
+++ b/tools/hoststubgen/hoststubgen/invoketest/hoststubgen-invoke-test.sh
@@ -51,6 +51,8 @@
 
 HOSTSTUBGEN_OUT=$TEMP/output.txt
 
+EXTRA_ARGS=""
+
 # Because of `set -e`, we can't return non-zero from functions, so we store
 # HostStubGen result in it.
 HOSTSTUBGEN_RC=0
@@ -115,6 +117,7 @@
       --keep-static-initializer-annotation \
           android.hosttest.annotation.HostSideTestStaticInitializerKeep \
       $filter_arg \
+      $EXTRA_ARGS \
       |& tee $HOSTSTUBGEN_OUT
   HOSTSTUBGEN_RC=${PIPESTATUS[0]}
   echo "HostStubGen exited with $HOSTSTUBGEN_RC"
@@ -209,7 +212,6 @@
 com.unsupported.*
 "
 
-
 run_hoststubgen_for_failure "One specific class disallowed" \
     "TinyFrameworkClassAnnotations is not allowed to have Ravenwood annotations" \
     "
@@ -229,6 +231,14 @@
 
 STUB="" IMPL="" run_hoststubgen_for_success "No stub, no impl generation" ""
 
+EXTRA_ARGS="--in-jar abc" run_hoststubgen_for_failure "Duplicate arg" \
+    "Duplicate or conflicting argument found: --in-jar" \
+    ""
+
+EXTRA_ARGS="--quiet" run_hoststubgen_for_failure "Conflicting arg" \
+    "Duplicate or conflicting argument found: --quiet" \
+    ""
+
 
 echo "All tests passed"
 exit 0
\ No newline at end of file
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
index 3cdddc2..dbcf3a5 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
@@ -22,7 +22,6 @@
 import com.android.hoststubgen.filters.DefaultHookInjectingFilter
 import com.android.hoststubgen.filters.FilterPolicy
 import com.android.hoststubgen.filters.ImplicitOutputFilter
-import com.android.hoststubgen.filters.KeepAllClassesFilter
 import com.android.hoststubgen.filters.OutputFilter
 import com.android.hoststubgen.filters.StubIntersectingFilter
 import com.android.hoststubgen.filters.createFilterFromTextPolicyFile
@@ -52,15 +51,15 @@
         val errors = HostStubGenErrors()
 
         // Load all classes.
-        val allClasses = loadClassStructures(options.inJar)
+        val allClasses = loadClassStructures(options.inJar.get)
 
         // Dump the classes, if specified.
-        options.inputJarDumpFile?.let {
+        options.inputJarDumpFile.ifSet {
             PrintWriter(it).use { pw -> allClasses.dump(pw) }
             log.i("Dump file created at $it")
         }
 
-        options.inputJarAsKeepAllFile?.let {
+        options.inputJarAsKeepAllFile.ifSet {
             PrintWriter(it).use {
                 pw -> allClasses.forEach {
                     classNode -> printAsTextPolicy(pw, classNode)
@@ -74,11 +73,11 @@
 
         // Transform the jar.
         convert(
-                options.inJar,
-                options.outStubJar,
-                options.outImplJar,
+                options.inJar.get,
+                options.outStubJar.get,
+                options.outImplJar.get,
                 filter,
-                options.enableClassChecker,
+                options.enableClassChecker.get,
                 allClasses,
                 errors,
         )
@@ -153,7 +152,7 @@
         // text-file based filter, which is handled by parseTextFilterPolicyFile.
 
         // The first filter is for the default policy from the command line options.
-        var filter: OutputFilter = ConstantFilter(options.defaultPolicy, "default-by-options")
+        var filter: OutputFilter = ConstantFilter(options.defaultPolicy.get, "default-by-options")
 
         // Next, we need a filter that resolves "class-wide" policies.
         // This is used when a member (methods, fields, nested classes) don't get any polices
@@ -163,16 +162,16 @@
 
         // Inject default hooks from options.
         filter = DefaultHookInjectingFilter(
-            options.defaultClassLoadHook,
-            options.defaultMethodCallHook,
+            options.defaultClassLoadHook.get,
+            options.defaultMethodCallHook.get,
             filter
         )
 
-        val annotationAllowedClassesFilter = options.annotationAllowedClassesFile.let { filename ->
-            if (filename == null) {
+        val annotationAllowedClassesFilter = options.annotationAllowedClassesFile.get.let { file ->
+            if (file == null) {
                 ClassFilter.newNullFilter(true) // Allow all classes
             } else {
-                ClassFilter.loadFromFile(filename, false)
+                ClassFilter.loadFromFile(file, false)
             }
         }
 
@@ -196,7 +195,7 @@
 
         // Next, "text based" filter, which allows to override polices without touching
         // the target code.
-        options.policyOverrideFile?.let {
+        options.policyOverrideFile.ifSet {
             filter = createFilterFromTextPolicyFile(it, allClasses, filter)
         }
 
@@ -212,11 +211,6 @@
         // Apply the implicit filter.
         filter = ImplicitOutputFilter(errors, allClasses, filter)
 
-        // Optionally keep all classes.
-        if (options.keepAllClasses) {
-            filter = KeepAllClassesFilter(filter)
-        }
-
         return filter
     }
 
@@ -422,9 +416,9 @@
             outVisitor = CheckClassAdapter(outVisitor)
         }
         val visitorOptions = BaseAdapter.Options(
-                enablePreTrace = options.enablePreTrace,
-                enablePostTrace = options.enablePostTrace,
-                enableNonStubMethodCallDetection = options.enableNonStubMethodCallDetection,
+                enablePreTrace = options.enablePreTrace.get,
+                enablePostTrace = options.enablePostTrace.get,
+                enableNonStubMethodCallDetection = options.enableNonStubMethodCallDetection.get,
                 errors = errors,
         )
         outVisitor = BaseAdapter.getVisitor(classInternalName, classes, outVisitor, filter,
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt
index 83f873d..0ae52af 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt
@@ -21,21 +21,60 @@
 import java.io.FileReader
 
 /**
+ * A single value that can only set once.
+ */
+class SetOnce<T>(
+        private var value: T,
+) {
+    class SetMoreThanOnceException : Exception()
+
+    private var set = false
+
+    fun set(v: T) {
+        if (set) {
+            throw SetMoreThanOnceException()
+        }
+        if (v == null) {
+            throw NullPointerException("This shouldn't happen")
+        }
+        set = true
+        value = v
+    }
+
+    val get: T
+        get() = this.value
+
+    val isSet: Boolean
+        get() = this.set
+
+    fun <R> ifSet(block: (T & Any) -> R): R? {
+        if (isSet) {
+            return block(value!!)
+        }
+        return null
+    }
+
+    override fun toString(): String {
+        return "$value"
+    }
+}
+
+/**
  * Options that can be set from command line arguments.
  */
 class HostStubGenOptions(
         /** Input jar file*/
-        var inJar: String = "",
+        var inJar: SetOnce<String> = SetOnce(""),
 
         /** Output stub jar file */
-        var outStubJar: String? = null,
+        var outStubJar: SetOnce<String?> = SetOnce(null),
 
         /** Output implementation jar file */
-        var outImplJar: String? = null,
+        var outImplJar: SetOnce<String?> = SetOnce(null),
 
-        var inputJarDumpFile: String? = null,
+        var inputJarDumpFile: SetOnce<String?> = SetOnce(null),
 
-        var inputJarAsKeepAllFile: String? = null,
+        var inputJarAsKeepAllFile: SetOnce<String?> = SetOnce(null),
 
         var stubAnnotations: MutableSet<String> = mutableSetOf(),
         var keepAnnotations: MutableSet<String> = mutableSetOf(),
@@ -51,27 +90,26 @@
 
         var packageRedirects: MutableList<Pair<String, String>> = mutableListOf(),
 
-        var annotationAllowedClassesFile: String? = null,
+        var annotationAllowedClassesFile: SetOnce<String?> = SetOnce(null),
 
-        var defaultClassLoadHook: String? = null,
-        var defaultMethodCallHook: String? = null,
+        var defaultClassLoadHook: SetOnce<String?> = SetOnce(null),
+        var defaultMethodCallHook: SetOnce<String?> = SetOnce(null),
 
         var intersectStubJars: MutableSet<String> = mutableSetOf(),
 
-        var policyOverrideFile: String? = null,
+        var policyOverrideFile: SetOnce<String?> = SetOnce(null),
 
-        var defaultPolicy: FilterPolicy = FilterPolicy.Remove,
-        var keepAllClasses: Boolean = false,
+        var defaultPolicy: SetOnce<FilterPolicy> = SetOnce(FilterPolicy.Remove),
 
-        var logLevel: LogLevel = LogLevel.Info,
+        var logLevel: SetOnce<LogLevel> = SetOnce(LogLevel.Info),
 
-        var cleanUpOnError: Boolean = false,
+        var cleanUpOnError: SetOnce<Boolean> = SetOnce(true),
 
-        var enableClassChecker: Boolean = false,
-        var enablePreTrace: Boolean = false,
-        var enablePostTrace: Boolean = false,
+        var enableClassChecker: SetOnce<Boolean> = SetOnce(false),
+        var enablePreTrace: SetOnce<Boolean> = SetOnce(false),
+        var enablePostTrace: SetOnce<Boolean> = SetOnce(false),
 
-        var enableNonStubMethodCallDetection: Boolean = false,
+        var enableNonStubMethodCallDetection: SetOnce<Boolean> = SetOnce(false),
 ) {
     companion object {
 
@@ -111,110 +149,120 @@
                     break
                 }
 
-                when (arg) {
-                    // TODO: Write help
-                    "-h", "--h" -> TODO("Help is not implemented yet")
+                // Define some shorthands...
+                fun nextArg(): String = ai.nextArgRequired(arg)
+                fun SetOnce<String>.setNextStringArg(): String = nextArg().also { this.set(it) }
+                fun SetOnce<String?>.setNextStringArg(): String = nextArg().also { this.set(it) }
+                fun MutableSet<String>.addUniqueAnnotationArg(): String =
+                        nextArg().also { this += ensureUniqueAnnotation(it) }
 
-                    "-v", "--verbose" -> ret.logLevel = LogLevel.Verbose
-                    "-d", "--debug" -> ret.logLevel = LogLevel.Debug
-                    "-q", "--quiet" -> ret.logLevel = LogLevel.None
+                try {
+                    when (arg) {
+                        // TODO: Write help
+                        "-h", "--help" -> TODO("Help is not implemented yet")
 
-                    "--in-jar" -> ret.inJar = ai.nextArgRequired(arg).ensureFileExists()
-                    "--out-stub-jar" -> ret.outStubJar = ai.nextArgRequired(arg)
-                    "--out-impl-jar" -> ret.outImplJar = ai.nextArgRequired(arg)
+                        "-v", "--verbose" -> ret.logLevel.set(LogLevel.Verbose)
+                        "-d", "--debug" -> ret.logLevel.set(LogLevel.Debug)
+                        "-q", "--quiet" -> ret.logLevel.set(LogLevel.None)
 
-                    "--policy-override-file" ->
-                        ret.policyOverrideFile = ai.nextArgRequired(arg).ensureFileExists()
+                        "--in-jar" -> ret.inJar.setNextStringArg().ensureFileExists()
+                        "--out-stub-jar" -> ret.outStubJar.setNextStringArg()
+                        "--out-impl-jar" -> ret.outImplJar.setNextStringArg()
 
-                    "--clean-up-on-error" -> ret.cleanUpOnError = true
-                    "--no-clean-up-on-error" -> ret.cleanUpOnError = false
+                        "--policy-override-file" ->
+                            ret.policyOverrideFile.setNextStringArg().ensureFileExists()
 
-                    "--default-remove" -> ret.defaultPolicy = FilterPolicy.Remove
-                    "--default-throw" -> ret.defaultPolicy = FilterPolicy.Throw
-                    "--default-keep" -> ret.defaultPolicy = FilterPolicy.Keep
-                    "--default-stub" -> ret.defaultPolicy = FilterPolicy.Stub
+                        "--clean-up-on-error" -> ret.cleanUpOnError.set(true)
+                        "--no-clean-up-on-error" -> ret.cleanUpOnError.set(false)
 
-                    "--keep-all-classes" -> ret.keepAllClasses = true
-                    "--no-keep-all-classes" -> ret.keepAllClasses = false
+                        "--default-remove" -> ret.defaultPolicy.set(FilterPolicy.Remove)
+                        "--default-throw" -> ret.defaultPolicy.set(FilterPolicy.Throw)
+                        "--default-keep" -> ret.defaultPolicy.set(FilterPolicy.Keep)
+                        "--default-stub" -> ret.defaultPolicy.set(FilterPolicy.Stub)
 
-                    "--stub-annotation" ->
-                        ret.stubAnnotations += ensureUniqueAnnotation(ai.nextArgRequired(arg))
+                        "--stub-annotation" ->
+                            ret.stubAnnotations.addUniqueAnnotationArg()
 
-                    "--keep-annotation" ->
-                        ret.keepAnnotations += ensureUniqueAnnotation(ai.nextArgRequired(arg))
+                        "--keep-annotation" ->
+                            ret.keepAnnotations.addUniqueAnnotationArg()
 
-                    "--stub-class-annotation" ->
-                        ret.stubClassAnnotations += ensureUniqueAnnotation(ai.nextArgRequired(arg))
+                        "--stub-class-annotation" ->
+                            ret.stubClassAnnotations.addUniqueAnnotationArg()
 
-                    "--keep-class-annotation" ->
-                        ret.keepClassAnnotations += ensureUniqueAnnotation(ai.nextArgRequired(arg))
+                        "--keep-class-annotation" ->
+                            ret.keepClassAnnotations.addUniqueAnnotationArg()
 
-                    "--throw-annotation" ->
-                        ret.throwAnnotations += ensureUniqueAnnotation(ai.nextArgRequired(arg))
+                        "--throw-annotation" ->
+                            ret.throwAnnotations.addUniqueAnnotationArg()
 
-                    "--remove-annotation" ->
-                        ret.removeAnnotations += ensureUniqueAnnotation(ai.nextArgRequired(arg))
+                        "--remove-annotation" ->
+                            ret.removeAnnotations.addUniqueAnnotationArg()
 
-                    "--substitute-annotation" ->
-                        ret.substituteAnnotations += ensureUniqueAnnotation(ai.nextArgRequired(arg))
+                        "--substitute-annotation" ->
+                            ret.substituteAnnotations.addUniqueAnnotationArg()
 
-                    "--native-substitute-annotation" ->
-                        ret.nativeSubstituteAnnotations +=
-                                ensureUniqueAnnotation(ai.nextArgRequired(arg))
+                        "--native-substitute-annotation" ->
+                            ret.nativeSubstituteAnnotations.addUniqueAnnotationArg()
 
-                    "--class-load-hook-annotation" ->
-                        ret.classLoadHookAnnotations +=
-                                ensureUniqueAnnotation(ai.nextArgRequired(arg))
+                        "--class-load-hook-annotation" ->
+                            ret.classLoadHookAnnotations.addUniqueAnnotationArg()
 
-                    "--keep-static-initializer-annotation" ->
-                        ret.keepStaticInitializerAnnotations +=
-                                ensureUniqueAnnotation(ai.nextArgRequired(arg))
+                        "--keep-static-initializer-annotation" ->
+                            ret.keepStaticInitializerAnnotations.addUniqueAnnotationArg()
 
-                    "--package-redirect" ->
-                        ret.packageRedirects += parsePackageRedirect(ai.nextArgRequired(arg))
+                        "--package-redirect" ->
+                            ret.packageRedirects += parsePackageRedirect(ai.nextArgRequired(arg))
 
-                    "--annotation-allowed-classes-file" ->
-                        ret.annotationAllowedClassesFile = ai.nextArgRequired(arg)
+                        "--annotation-allowed-classes-file" ->
+                            ret.annotationAllowedClassesFile.setNextStringArg()
 
-                    "--default-class-load-hook" ->
-                        ret.defaultClassLoadHook = ai.nextArgRequired(arg)
+                        "--default-class-load-hook" ->
+                            ret.defaultClassLoadHook.setNextStringArg()
 
-                    "--default-method-call-hook" ->
-                        ret.defaultMethodCallHook = ai.nextArgRequired(arg)
+                        "--default-method-call-hook" ->
+                            ret.defaultMethodCallHook.setNextStringArg()
 
-                    "--intersect-stub-jar" ->
-                        ret.intersectStubJars += ai.nextArgRequired(arg).ensureFileExists()
+                        "--intersect-stub-jar" ->
+                            ret.intersectStubJars += nextArg().ensureFileExists()
 
-                    "--gen-keep-all-file" ->
-                        ret.inputJarAsKeepAllFile = ai.nextArgRequired(arg)
+                        "--gen-keep-all-file" ->
+                            ret.inputJarAsKeepAllFile.setNextStringArg()
 
-                    // Following options are for debugging.
-                    "--enable-class-checker" -> ret.enableClassChecker = true
-                    "--no-class-checker" -> ret.enableClassChecker = false
+                        // Following options are for debugging.
+                        "--enable-class-checker" -> ret.enableClassChecker.set(true)
+                        "--no-class-checker" -> ret.enableClassChecker.set(false)
 
-                    "--enable-pre-trace" -> ret.enablePreTrace = true
-                    "--no-pre-trace" -> ret.enablePreTrace = false
+                        "--enable-pre-trace" -> ret.enablePreTrace.set(true)
+                        "--no-pre-trace" -> ret.enablePreTrace.set(false)
 
-                    "--enable-post-trace" -> ret.enablePostTrace = true
-                    "--no-post-trace" -> ret.enablePostTrace = false
+                        "--enable-post-trace" -> ret.enablePostTrace.set(true)
+                        "--no-post-trace" -> ret.enablePostTrace.set(false)
 
-                    "--enable-non-stub-method-check" -> ret.enableNonStubMethodCallDetection = true
-                    "--no-non-stub-method-check" -> ret.enableNonStubMethodCallDetection = false
+                        "--enable-non-stub-method-check" ->
+                            ret.enableNonStubMethodCallDetection.set(true)
 
-                    "--gen-input-dump-file" -> ret.inputJarDumpFile = ai.nextArgRequired(arg)
+                        "--no-non-stub-method-check" ->
+                            ret.enableNonStubMethodCallDetection.set(false)
 
-                    else -> throw ArgumentsException("Unknown option: $arg")
+                        "--gen-input-dump-file" -> ret.inputJarDumpFile.setNextStringArg()
+
+                        else -> throw ArgumentsException("Unknown option: $arg")
+                    }
+                } catch (e: SetOnce.SetMoreThanOnceException) {
+                    throw ArgumentsException("Duplicate or conflicting argument found: $arg")
                 }
             }
-            if (ret.inJar.isEmpty()) {
+            log.w(ret.toString())
+
+            if (!ret.inJar.isSet) {
                 throw ArgumentsException("Required option missing: --in-jar")
             }
-            if (ret.outStubJar == null && ret.outImplJar == null) {
+            if (!ret.outStubJar.isSet && !ret.outImplJar.isSet) {
                 log.w("Neither --out-stub-jar nor --out-impl-jar is set." +
                         " $COMMAND_NAME will not generate jar files.")
             }
 
-            if (ret.enableNonStubMethodCallDetection) {
+            if (ret.enableNonStubMethodCallDetection.get) {
                 log.w("--enable-non-stub-method-check is not fully implemented yet." +
                     " See the todo in doesMethodNeedNonStubCallCheck().")
             }
@@ -329,7 +377,6 @@
               intersectStubJars=$intersectStubJars,
               policyOverrideFile=$policyOverrideFile,
               defaultPolicy=$defaultPolicy,
-              keepAllClasses=$keepAllClasses,
               logLevel=$logLevel,
               cleanUpOnError=$cleanUpOnError,
               enableClassChecker=$enableClassChecker,
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/Main.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/Main.kt
index 0321d9d..38ba0cc 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/Main.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/Main.kt
@@ -28,9 +28,9 @@
     try {
         // Parse the command line arguments.
         val options = HostStubGenOptions.parseArgs(args)
-        clanupOnError = options.cleanUpOnError
+        clanupOnError = options.cleanUpOnError.get
 
-        log.level = options.logLevel
+        log.level = options.logLevel.get
 
         log.v("HostStubGen started")
         log.v("Options: $options")
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/KeepAllClassesFilter.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/KeepAllClassesFilter.kt
deleted file mode 100644
index 45dd38d1..0000000
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/KeepAllClassesFilter.kt
+++ /dev/null
@@ -1,30 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package com.android.hoststubgen.filters
-
-/**
- * An [OutputFilter] that keeps all classes by default. (but none of its members)
- *
- * We're not currently using it, but using it *might* make certain things easier. For example, with
- * this, all classes would at least be loadable.
- */
-class KeepAllClassesFilter(fallback: OutputFilter) : DelegatingFilter(fallback) {
-    override fun getPolicyForClass(className: String): FilterPolicyWithReason {
-        // If the default visibility wouldn't keep it, change it to "keep".
-        val f = super.getPolicyForClass(className)
-        return f.promoteToKeep("keep-all-classes")
-    }
-}
\ No newline at end of file