Store aggregated power stats in PowerStatsStore

Bug: 297274264
Test: atest PowerStatsTests

Change-Id: I8b6af08d8ade14503fe7e5d3258ee22f7051ce71
diff --git a/core/java/com/android/internal/os/LongArrayMultiStateCounter.java b/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
index 5ea6ba8..1f44b33 100644
--- a/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
+++ b/core/java/com/android/internal/os/LongArrayMultiStateCounter.java
@@ -191,6 +191,27 @@
     }
 
     /**
+     * Sets the new values for the given state.
+     */
+    public void setValues(int state, long[] values) {
+        if (state < 0 || state >= mStateCount) {
+            throw new IllegalArgumentException(
+                    "State: " + state + ", outside the range: [0-" + (mStateCount - 1) + "]");
+        }
+        if (values.length != mLength) {
+            throw new IllegalArgumentException(
+                    "Invalid array length: " + values.length + ", expected: " + mLength);
+        }
+        LongArrayContainer container = sTmpArrayContainer.getAndSet(null);
+        if (container == null || container.mLength != values.length) {
+            container = new LongArrayContainer(values.length);
+        }
+        container.setValues(values);
+        native_setValues(mNativeObject, state, container.mNativeObject);
+        sTmpArrayContainer.set(container);
+    }
+
+    /**
      * Sets the new values.  The delta between the previously set values and these values
      * is distributed among the state according to the time the object spent in those states
      * since the previous call to updateValues.
@@ -317,6 +338,10 @@
     private static native void native_setState(long nativeObject, int state, long timestampMs);
 
     @CriticalNative
+    private static native void native_setValues(long nativeObject, int state,
+            long longArrayContainerNativeObject);
+
+    @CriticalNative
     private static native void native_updateValues(long nativeObject,
             long longArrayContainerNativeObject, long timestampMs);
 
diff --git a/core/java/com/android/internal/os/MultiStateStats.java b/core/java/com/android/internal/os/MultiStateStats.java
index f971849..dc5055a 100644
--- a/core/java/com/android/internal/os/MultiStateStats.java
+++ b/core/java/com/android/internal/os/MultiStateStats.java
@@ -16,9 +16,17 @@
 
 package com.android.internal.os;
 
+import android.util.Slog;
+
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.Preconditions;
+import com.android.modules.utils.TypedXmlPullParser;
+import com.android.modules.utils.TypedXmlSerializer;
 
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
 import java.io.PrintWriter;
 import java.util.Arrays;
 
@@ -29,15 +37,21 @@
  * values;
  */
 public class MultiStateStats {
+    private static final String TAG = "MultiStateStats";
+
+    private static final String XML_TAG_STATS = "stats";
+
     /**
      * A set of states, e.g. on-battery, screen-on, procstate.  The state values are integers
      * from 0 to States.mLabels.length
      */
     public static class States {
+        final String mName;
         final boolean mTracked;
         final String[] mLabels;
 
-        public States(boolean tracked, String... labels) {
+        public States(String name, boolean tracked, String... labels) {
+            mName = name;
             this.mTracked = tracked;
             this.mLabels = labels;
         }
@@ -155,11 +169,28 @@
                    >>> mStateBitFieldShifts[stateIndex];
         }
 
-        private int setStateInComposite(int baseCompositeState, int stateIndex, int value) {
+        int setStateInComposite(int baseCompositeState, int stateIndex, int value) {
             return (baseCompositeState & ~mStateBitFieldMasks[stateIndex])
                     | (value << mStateBitFieldShifts[stateIndex]);
         }
 
+        int setStateInComposite(int compositeState, String stateName, String stateLabel) {
+            for (int stateIndex = 0; stateIndex < mStates.length; stateIndex++) {
+                States stateConfig = mStates[stateIndex];
+                if (stateConfig.mName.equals(stateName)) {
+                    for (int state = 0; state < stateConfig.mLabels.length; state++) {
+                        if (stateConfig.mLabels[state].equals(stateLabel)) {
+                            return setStateInComposite(compositeState, stateIndex, state);
+                        }
+                    }
+                    Slog.e(TAG, "Unexpected label '" + stateLabel + "' for state: " + stateName);
+                    return -1;
+                }
+            }
+            Slog.e(TAG, "Unsupported state: " + stateName);
+            return -1;
+        }
+
         /**
          * Allocates a new stats container using this Factory's configuration.
          */
@@ -195,6 +226,10 @@
             }
             return serialState;
         }
+
+        int getSerialState(int compositeState) {
+            return mCompositeToSerialState[compositeState];
+        }
     }
 
     private final Factory mFactory;
@@ -254,6 +289,106 @@
     }
 
     /**
+     * Stores contents in an XML doc.
+     */
+    public void writeXml(TypedXmlSerializer serializer) throws IOException {
+        long[] tmpArray = new long[mCounter.getArrayLength()];
+        writeXmlAllStates(serializer, new int[mFactory.mStates.length], 0, tmpArray);
+    }
+
+    private void writeXmlAllStates(TypedXmlSerializer serializer, int[] states, int stateIndex,
+            long[] values) throws IOException {
+        if (stateIndex < states.length) {
+            if (!mFactory.mStates[stateIndex].mTracked) {
+                writeXmlAllStates(serializer, states, stateIndex + 1, values);
+                return;
+            }
+
+            for (int i = 0; i < mFactory.mStates[stateIndex].mLabels.length; i++) {
+                states[stateIndex] = i;
+                writeXmlAllStates(serializer, states, stateIndex + 1, values);
+            }
+            return;
+        }
+
+        mCounter.getCounts(values, mFactory.getSerialState(states));
+        boolean nonZero = false;
+        for (long value : values) {
+            if (value != 0) {
+                nonZero = true;
+                break;
+            }
+        }
+        if (!nonZero) {
+            return;
+        }
+
+        serializer.startTag(null, XML_TAG_STATS);
+
+        for (int i = 0; i < states.length; i++) {
+            if (mFactory.mStates[i].mTracked && states[i] != 0) {
+                serializer.attribute(null, mFactory.mStates[i].mName,
+                        mFactory.mStates[i].mLabels[states[i]]);
+            }
+        }
+        for (int i = 0; i < values.length; i++) {
+            if (values[i] != 0) {
+                serializer.attributeLong(null, "_" + i, values[i]);
+            }
+        }
+        serializer.endTag(null, XML_TAG_STATS);
+    }
+
+    /**
+     * Populates the object with contents in an XML doc. The parser is expected to be
+     * positioned on the opening tag of the corresponding element.
+     */
+    public boolean readFromXml(TypedXmlPullParser parser) throws XmlPullParserException,
+            IOException {
+        String outerTag = parser.getName();
+        long[] tmpArray = new long[mCounter.getArrayLength()];
+        int eventType = parser.getEventType();
+        while (eventType != XmlPullParser.END_DOCUMENT
+               && !(eventType == XmlPullParser.END_TAG && parser.getName().equals(outerTag))) {
+            if (eventType == XmlPullParser.START_TAG) {
+                if (parser.getName().equals(XML_TAG_STATS)) {
+                    Arrays.fill(tmpArray, 0);
+                    int compositeState = 0;
+                    int attributeCount = parser.getAttributeCount();
+                    for (int i = 0; i < attributeCount; i++) {
+                        String attributeName = parser.getAttributeName(i);
+                        if (attributeName.startsWith("_")) {
+                            int index;
+                            try {
+                                index = Integer.parseInt(attributeName.substring(1));
+                            } catch (NumberFormatException e) {
+                                throw new XmlPullParserException(
+                                        "Unexpected index syntax: " + attributeName, parser, e);
+                            }
+                            if (index < 0 || index >= tmpArray.length) {
+                                Slog.e(TAG, "State index out of bounds: " + index
+                                            + " length: " + tmpArray.length);
+                                return false;
+                            }
+                            tmpArray[index] = parser.getAttributeLong(i);
+                        } else {
+                            String attributeValue = parser.getAttributeValue(i);
+                            compositeState = mFactory.setStateInComposite(compositeState,
+                                    attributeName, attributeValue);
+                            if (compositeState == -1) {
+                                return false;
+                            }
+                        }
+                    }
+                    mCounter.setValues(mFactory.getSerialState(compositeState), tmpArray);
+                }
+            }
+            eventType = parser.next();
+        }
+        return true;
+    }
+
+    /**
      * Prints the accumulated stats, one line of every combination of states that has data.
      */
     public void dump(PrintWriter pw) {
diff --git a/core/java/com/android/internal/os/PowerStats.java b/core/java/com/android/internal/os/PowerStats.java
index 8f66d1f..1130a45 100644
--- a/core/java/com/android/internal/os/PowerStats.java
+++ b/core/java/com/android/internal/os/PowerStats.java
@@ -24,9 +24,16 @@
 import android.os.PersistableBundle;
 import android.os.UserHandle;
 import android.util.IndentingPrintWriter;
-import android.util.Log;
+import android.util.Slog;
 import android.util.SparseArray;
 
+import com.android.modules.utils.TypedXmlPullParser;
+import com.android.modules.utils.TypedXmlSerializer;
+
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.Objects;
 
@@ -61,6 +68,13 @@
      * to adjust the algorithm in accordance with the stats available on the device.
      */
     public static class Descriptor {
+        public static final String XML_TAG_DESCRIPTOR = "descriptor";
+        private static final String XML_ATTR_ID = "id";
+        private static final String XML_ATTR_NAME = "name";
+        private static final String XML_ATTR_STATS_ARRAY_LENGTH = "stats-array-length";
+        private static final String XML_ATTR_UID_STATS_ARRAY_LENGTH = "uid-stats-array-length";
+        private static final String XML_TAG_EXTRAS = "extras";
+
         /**
          * {@link BatteryConsumer.PowerComponent} (e.g. CPU, WIFI etc) that this snapshot relates
          * to; or a custom power component ID (if the value
@@ -126,7 +140,7 @@
             int firstWord = parcel.readInt();
             int version = (firstWord & PARCEL_FORMAT_VERSION_MASK) >>> PARCEL_FORMAT_VERSION_SHIFT;
             if (version != PARCEL_FORMAT_VERSION) {
-                Log.w(TAG, "Cannot read PowerStats from Parcel - the parcel format version has "
+                Slog.w(TAG, "Cannot read PowerStats from Parcel - the parcel format version has "
                            + "changed from " + version + " to " + PARCEL_FORMAT_VERSION);
                 return null;
             }
@@ -155,6 +169,71 @@
                     that.extras);  // Since the Parcel is now unparceled, do a deep comparison
         }
 
+        /**
+         * Stores contents in an XML doc.
+         */
+        public void writeXml(TypedXmlSerializer serializer) throws IOException {
+            serializer.startTag(null, XML_TAG_DESCRIPTOR);
+            serializer.attributeInt(null, XML_ATTR_ID, powerComponentId);
+            serializer.attribute(null, XML_ATTR_NAME, name);
+            serializer.attributeInt(null, XML_ATTR_STATS_ARRAY_LENGTH, statsArrayLength);
+            serializer.attributeInt(null, XML_ATTR_UID_STATS_ARRAY_LENGTH, uidStatsArrayLength);
+            try {
+                serializer.startTag(null, XML_TAG_EXTRAS);
+                extras.saveToXml(serializer);
+                serializer.endTag(null, XML_TAG_EXTRAS);
+            } catch (XmlPullParserException e) {
+                throw new IOException(e);
+            }
+            serializer.endTag(null, XML_TAG_DESCRIPTOR);
+        }
+
+        /**
+         * Creates a Descriptor by parsing an XML doc.  The parser is expected to be positioned
+         * on or before the opening "descriptor" tag.
+         */
+        public static Descriptor createFromXml(TypedXmlPullParser parser)
+                throws XmlPullParserException, IOException {
+            int powerComponentId = -1;
+            String name = null;
+            int statsArrayLength = 0;
+            int uidStatsArrayLength = 0;
+            PersistableBundle extras = null;
+            int eventType = parser.getEventType();
+            while (eventType != XmlPullParser.END_DOCUMENT
+                   && !(eventType == XmlPullParser.END_TAG
+                        && parser.getName().equals(XML_TAG_DESCRIPTOR))) {
+                if (eventType == XmlPullParser.START_TAG) {
+                    switch (parser.getName()) {
+                        case XML_TAG_DESCRIPTOR:
+                            powerComponentId = parser.getAttributeInt(null, XML_ATTR_ID);
+                            name = parser.getAttributeValue(null, XML_ATTR_NAME);
+                            statsArrayLength = parser.getAttributeInt(null,
+                                    XML_ATTR_STATS_ARRAY_LENGTH);
+                            uidStatsArrayLength = parser.getAttributeInt(null,
+                                    XML_ATTR_UID_STATS_ARRAY_LENGTH);
+                            break;
+                        case XML_TAG_EXTRAS:
+                            extras = PersistableBundle.restoreFromXml(parser);
+                            break;
+                    }
+                }
+                eventType = parser.next();
+            }
+            if (powerComponentId == -1) {
+                return null;
+            } else if (powerComponentId >= BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID) {
+                return new Descriptor(powerComponentId, name, statsArrayLength, uidStatsArrayLength,
+                        extras);
+            } else if (powerComponentId < BatteryConsumer.POWER_COMPONENT_COUNT) {
+                return new Descriptor(powerComponentId, statsArrayLength, uidStatsArrayLength,
+                        extras);
+            } else {
+                Slog.e(TAG, "Unrecognized power component: " + powerComponentId);
+                return null;
+            }
+        }
+
         @Override
         public int hashCode() {
             return Objects.hash(powerComponentId);
@@ -259,7 +338,7 @@
 
             Descriptor descriptor = registry.get(powerComponentId);
             if (descriptor == null) {
-                Log.e(TAG, "Unsupported PowerStats for power component ID: " + powerComponentId);
+                Slog.e(TAG, "Unsupported PowerStats for power component ID: " + powerComponentId);
                 return null;
             }
             PowerStats stats = new PowerStats(descriptor);
@@ -274,7 +353,7 @@
                 stats.uidStats.put(uid, uidStats);
             }
             if (parcel.dataPosition() != endPos) {
-                Log.e(TAG, "Corrupted PowerStats parcel. Expected length: " + length
+                Slog.e(TAG, "Corrupted PowerStats parcel. Expected length: " + length
                            + ", actual length: " + (parcel.dataPosition() - startPos));
                 return null;
             }
diff --git a/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp b/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
index 6920211..1f29735 100644
--- a/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
+++ b/core/jni/com_android_internal_os_LongArrayMultiStateCounter.cpp
@@ -55,6 +55,15 @@
     counter->setState(state, timestamp);
 }
 
+static void native_setValues(jlong nativePtr, jint state, jlong longArrayContainerNativePtr) {
+    battery::LongArrayMultiStateCounter *counter =
+            reinterpret_cast<battery::LongArrayMultiStateCounter *>(nativePtr);
+    std::vector<uint64_t> *vector =
+            reinterpret_cast<std::vector<uint64_t> *>(longArrayContainerNativePtr);
+
+    counter->setValue(state, *vector);
+}
+
 static void native_updateValues(jlong nativePtr, jlong longArrayContainerNativePtr,
                                 jlong timestamp) {
     battery::LongArrayMultiStateCounter *counter =
@@ -210,6 +219,8 @@
         // @CriticalNative
         {"native_setState", "(JIJ)V", (void *)native_setState},
         // @CriticalNative
+        {"native_setValues", "(JIJ)V", (void *)native_setValues},
+        // @CriticalNative
         {"native_updateValues", "(JJJ)V", (void *)native_updateValues},
         // @CriticalNative
         {"native_incrementValues", "(JJJ)V", (void *)native_incrementValues},
diff --git a/core/res/res/values/config_battery_stats.xml b/core/res/res/values/config_battery_stats.xml
index 8fb48bc..e42962c 100644
--- a/core/res/res/values/config_battery_stats.xml
+++ b/core/res/res/values/config_battery_stats.xml
@@ -31,4 +31,13 @@
     is a relatively expensive operation, this throttle period may need to be adjusted for low-power
     devices-->
     <integer name="config_defaultPowerStatsThrottlePeriodCpu">60000</integer>
+
+    <!-- PowerStats aggregation period in milliseconds. This is the interval at which the power
+    stats aggregation procedure is performed and the results stored in PowerStatsStore. -->
+    <integer name="config_powerStatsAggregationPeriod">14400000</integer>
+
+    <!-- PowerStats aggregation span duration in milliseconds. This is the length of battery
+    history time for every aggregated power stats span that is stored stored in PowerStatsStore.
+    It should not be larger than config_powerStatsAggregationPeriod (but it can be the same) -->
+    <integer name="config_aggregatedPowerStatsSpanDuration">3600000</integer>
 </resources>
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index c1ecb05..76744ea 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -5124,6 +5124,8 @@
   <java-symbol type="bool" name="config_batteryStatsResetOnUnplugHighBatteryLevel" />
   <java-symbol type="bool" name="config_batteryStatsResetOnUnplugAfterSignificantCharge" />
   <java-symbol type="integer" name="config_defaultPowerStatsThrottlePeriodCpu" />
+  <java-symbol type="integer" name="config_powerStatsAggregationPeriod" />
+  <java-symbol type="integer" name="config_aggregatedPowerStatsSpanDuration" />
 
   <java-symbol name="materialColorOnSecondaryFixedVariant" type="attr"/>
   <java-symbol name="materialColorOnTertiaryFixedVariant" type="attr"/>
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 56061ff..0ab81a5 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -117,9 +117,11 @@
 import com.android.server.net.BaseNetworkObserver;
 import com.android.server.pm.UserManagerInternal;
 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.BatteryStatsImpl;
 import com.android.server.power.stats.BatteryUsageStatsProvider;
+import com.android.server.power.stats.PowerStatsAggregator;
 import com.android.server.power.stats.PowerStatsScheduler;
 import com.android.server.power.stats.PowerStatsStore;
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
@@ -170,6 +172,7 @@
     final BatteryStatsImpl mStats;
     final CpuWakeupStats mCpuWakeupStats;
     private final PowerStatsStore mPowerStatsStore;
+    private final PowerStatsAggregator mPowerStatsAggregator;
     private final PowerStatsScheduler mPowerStatsScheduler;
     private final BatteryStatsImpl.UserInfoProvider mUserManagerUserInfoProvider;
     private final Context mContext;
@@ -410,15 +413,37 @@
         mStats.setRadioScanningTimeoutLocked(mContext.getResources().getInteger(
                 com.android.internal.R.integer.config_radioScanningTimeout) * 1000L);
         mStats.startTrackingSystemServerCpuTime();
-        mPowerStatsStore = new PowerStatsStore(systemDir, mHandler);
+
+        AggregatedPowerStatsConfig aggregatedPowerStatsConfig = getAggregatedPowerStatsConfig();
+        mPowerStatsStore = new PowerStatsStore(systemDir, mHandler, aggregatedPowerStatsConfig);
         mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context, mStats,
                 mPowerStatsStore);
-        mPowerStatsScheduler = new PowerStatsScheduler(mPowerStatsStore, mMonotonicClock, mHandler,
-                mStats, mBatteryUsageStatsProvider);
+        mPowerStatsAggregator = new PowerStatsAggregator(aggregatedPowerStatsConfig,
+                mStats.getHistory());
+        final long aggregatedPowerStatsSpanDuration = context.getResources().getInteger(
+                com.android.internal.R.integer.config_aggregatedPowerStatsSpanDuration);
+        final long powerStatsAggregationPeriod = context.getResources().getInteger(
+                com.android.internal.R.integer.config_powerStatsAggregationPeriod);
+        mPowerStatsScheduler = new PowerStatsScheduler(context, mPowerStatsAggregator,
+                aggregatedPowerStatsSpanDuration, powerStatsAggregationPeriod, mPowerStatsStore,
+                Clock.SYSTEM_CLOCK, mMonotonicClock, mHandler, mStats, mBatteryUsageStatsProvider);
         mCpuWakeupStats = new CpuWakeupStats(context, R.xml.irq_device_map, mHandler);
         mConfigFile = new AtomicFile(new File(systemDir, "battery_usage_stats_config"));
     }
 
+    private AggregatedPowerStatsConfig getAggregatedPowerStatsConfig() {
+        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);
+        return config;
+    }
+
     /**
      * Creates an instance of BatteryStatsService and restores data from stored state.
      */
@@ -475,7 +500,7 @@
      */
     public void onSystemReady() {
         mStats.onSystemReady();
-        mPowerStatsScheduler.start();
+        mPowerStatsScheduler.start(Flags.streamlinedBatteryStats());
     }
 
     private final class LocalService extends BatteryStatsInternal {
@@ -2646,7 +2671,7 @@
     }
 
     private void dumpAggregatedStats(PrintWriter pw) {
-        mStats.dumpAggregatedStats(pw, /* startTime */ 0, /* endTime */0);
+        mPowerStatsScheduler.aggregateAndDumpPowerStats(pw);
     }
 
     private void dumpPowerStatsStore(PrintWriter pw) {
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 288912a..bc90f5c 100644
--- a/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java
+++ b/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java
@@ -18,6 +18,8 @@
 
 import android.annotation.CurrentTimeMillisLong;
 import android.annotation.DurationMillisLong;
+import android.annotation.NonNull;
+import android.os.BatteryConsumer;
 import android.os.UserHandle;
 import android.text.format.DateFormat;
 import android.util.IndentingPrintWriter;
@@ -25,8 +27,13 @@
 import android.util.TimeUtils;
 
 import com.android.internal.os.PowerStats;
+import com.android.modules.utils.TypedXmlPullParser;
+import com.android.modules.utils.TypedXmlSerializer;
 
-import java.io.PrintWriter;
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
@@ -40,6 +47,8 @@
 class AggregatedPowerStats {
     private static final String TAG = "AggregatedPowerStats";
     private static final int MAX_CLOCK_UPDATES = 100;
+    private static final String XML_TAG_AGGREGATED_POWER_STATS = "agg-power-stats";
+
     private final PowerComponentAggregatedPowerStats[] mPowerComponentStats;
 
     static class ClockUpdate {
@@ -52,8 +61,23 @@
     @DurationMillisLong
     private long mDurationMs;
 
-    AggregatedPowerStats(PowerComponentAggregatedPowerStats... powerComponentAggregatedPowerStats) {
-        this.mPowerComponentStats = powerComponentAggregatedPowerStats;
+    AggregatedPowerStats(AggregatedPowerStatsConfig aggregatedPowerStatsConfig) {
+        List<AggregatedPowerStatsConfig.PowerComponent> configs =
+                aggregatedPowerStatsConfig.getPowerComponentsAggregatedStatsConfigs();
+        mPowerComponentStats = new PowerComponentAggregatedPowerStats[configs.size()];
+        for (int i = 0; i < configs.size(); i++) {
+            mPowerComponentStats[i] = createPowerComponentAggregatedPowerStats(configs.get(i));
+        }
+    }
+
+    private PowerComponentAggregatedPowerStats createPowerComponentAggregatedPowerStats(
+            AggregatedPowerStatsConfig.PowerComponent config) {
+        switch (config.getPowerComponentId()) {
+            case BatteryConsumer.POWER_COMPONENT_CPU:
+                return new CpuAggregatedPowerStats(config);
+            default:
+                return new PowerComponentAggregatedPowerStats(config);
+        }
     }
 
     /**
@@ -110,13 +134,14 @@
         return null;
     }
 
-    void setDeviceState(@PowerStatsAggregator.TrackedState int stateId, int state, long time) {
+    void setDeviceState(@AggregatedPowerStatsConfig.TrackedState int stateId, int state,
+            long time) {
         for (PowerComponentAggregatedPowerStats stats : mPowerComponentStats) {
             stats.setState(stateId, state, time);
         }
     }
 
-    void setUidState(int uid, @PowerStatsAggregator.TrackedState int stateId, int state,
+    void setUidState(int uid, @AggregatedPowerStatsConfig.TrackedState int stateId, int state,
             long time) {
         for (PowerComponentAggregatedPowerStats stats : mPowerComponentStats) {
             stats.setUidState(uid, stateId, state, time);
@@ -150,8 +175,56 @@
         }
     }
 
-    void dump(PrintWriter pw) {
-        IndentingPrintWriter ipw = new IndentingPrintWriter(pw);
+    public void writeXml(TypedXmlSerializer serializer) throws IOException {
+        serializer.startTag(null, XML_TAG_AGGREGATED_POWER_STATS);
+        for (PowerComponentAggregatedPowerStats stats : mPowerComponentStats) {
+            stats.writeXml(serializer);
+        }
+        serializer.endTag(null, XML_TAG_AGGREGATED_POWER_STATS);
+        serializer.flush();
+    }
+
+    @NonNull
+    public static AggregatedPowerStats createFromXml(
+            TypedXmlPullParser parser, AggregatedPowerStatsConfig aggregatedPowerStatsConfig)
+            throws XmlPullParserException, IOException {
+        AggregatedPowerStats stats = new AggregatedPowerStats(aggregatedPowerStatsConfig);
+        boolean inElement = false;
+        boolean skipToEnd = false;
+        int eventType = parser.getEventType();
+        while (eventType != XmlPullParser.END_DOCUMENT
+                   && !(eventType == XmlPullParser.END_TAG
+                        && parser.getName().equals(XML_TAG_AGGREGATED_POWER_STATS))) {
+            if (!skipToEnd && eventType == XmlPullParser.START_TAG) {
+                switch (parser.getName()) {
+                    case XML_TAG_AGGREGATED_POWER_STATS:
+                        inElement = true;
+                        break;
+                    case PowerComponentAggregatedPowerStats.XML_TAG_POWER_COMPONENT:
+                        if (!inElement) {
+                            break;
+                        }
+
+                        int powerComponentId = parser.getAttributeInt(null,
+                                PowerComponentAggregatedPowerStats.XML_ATTR_ID);
+                        for (PowerComponentAggregatedPowerStats powerComponent :
+                                stats.mPowerComponentStats) {
+                            if (powerComponent.powerComponentId == powerComponentId) {
+                                if (!powerComponent.readFromXml(parser)) {
+                                    skipToEnd = true;
+                                }
+                                break;
+                            }
+                        }
+                        break;
+                }
+            }
+            eventType = parser.next();
+        }
+        return stats;
+    }
+
+    void dump(IndentingPrintWriter ipw) {
         StringBuilder sb = new StringBuilder();
         long baseTime = 0;
         for (int i = 0; i < mClockUpdates.size(); i++) {
diff --git a/services/core/java/com/android/server/power/stats/AggregatedPowerStatsConfig.java b/services/core/java/com/android/server/power/stats/AggregatedPowerStatsConfig.java
new file mode 100644
index 0000000..477c228
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/AggregatedPowerStatsConfig.java
@@ -0,0 +1,156 @@
+/*
+ * 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.annotation.IntDef;
+import android.os.BatteryConsumer;
+
+import com.android.internal.os.MultiStateStats;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Configuration that controls how power stats are aggregated.  It determines which state changes
+ * are to be considered as essential dimensions ("tracked states") for each power component (CPU,
+ * WiFi, etc).  Also, it determines which states are tracked globally and which ones on a per-UID
+ * basis.
+ */
+public class AggregatedPowerStatsConfig {
+    public static final int STATE_POWER = 0;
+    public static final int STATE_SCREEN = 1;
+    public static final int STATE_PROCESS_STATE = 2;
+
+    @IntDef({
+            STATE_POWER,
+            STATE_SCREEN,
+            STATE_PROCESS_STATE,
+    })
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface TrackedState {
+    }
+
+    static final String STATE_NAME_POWER = "pwr";
+    static final int POWER_STATE_BATTERY = 0;
+    static final int POWER_STATE_OTHER = 1;   // Plugged in, or on wireless charger, etc.
+    static final String[] STATE_LABELS_POWER = {"pwr-battery", "pwr-other"};
+
+    static final String STATE_NAME_SCREEN = "scr";
+    static final int SCREEN_STATE_ON = 0;
+    static final int SCREEN_STATE_OTHER = 1;  // Off, doze etc
+    static final String[] STATE_LABELS_SCREEN = {"scr-on", "scr-other"};
+
+    static final String STATE_NAME_PROCESS_STATE = "ps";
+    static final String[] STATE_LABELS_PROCESS_STATE;
+
+    static {
+        String[] procStateLabels = new String[BatteryConsumer.PROCESS_STATE_COUNT];
+        for (int i = 0; i < BatteryConsumer.PROCESS_STATE_COUNT; i++) {
+            procStateLabels[i] = BatteryConsumer.processStateToString(i);
+        }
+        STATE_LABELS_PROCESS_STATE = procStateLabels;
+    }
+
+    /**
+     * Configuration for a give power component (CPU, WiFi, etc)
+     */
+    public static class PowerComponent {
+        private final int mPowerComponentId;
+        private @TrackedState int[] mTrackedDeviceStates;
+        private @TrackedState int[] mTrackedUidStates;
+
+        PowerComponent(int powerComponentId) {
+            this.mPowerComponentId = powerComponentId;
+        }
+
+        /**
+         * Configures which states should be tracked as separate dimensions for the entire device.
+         */
+        public PowerComponent trackDeviceStates(@TrackedState int... states) {
+            mTrackedDeviceStates = states;
+            return this;
+        }
+
+        /**
+         * Configures which states should be tracked as separate dimensions on a per-UID basis.
+         */
+        public PowerComponent trackUidStates(@TrackedState int... states) {
+            mTrackedUidStates = states;
+            return this;
+        }
+
+        public int getPowerComponentId() {
+            return mPowerComponentId;
+        }
+
+        public MultiStateStats.States[] getDeviceStateConfig() {
+            return new MultiStateStats.States[]{
+                    new MultiStateStats.States(STATE_NAME_POWER,
+                            isTracked(mTrackedDeviceStates, STATE_POWER),
+                            STATE_LABELS_POWER),
+                    new MultiStateStats.States(STATE_NAME_SCREEN,
+                            isTracked(mTrackedDeviceStates, STATE_SCREEN),
+                            STATE_LABELS_SCREEN),
+            };
+        }
+
+        public MultiStateStats.States[] getUidStateConfig() {
+            return new MultiStateStats.States[]{
+                    new MultiStateStats.States(STATE_NAME_POWER,
+                            isTracked(mTrackedUidStates, STATE_POWER),
+                            AggregatedPowerStatsConfig.STATE_LABELS_POWER),
+                    new MultiStateStats.States(STATE_NAME_SCREEN,
+                            isTracked(mTrackedUidStates, STATE_SCREEN),
+                            AggregatedPowerStatsConfig.STATE_LABELS_SCREEN),
+                    new MultiStateStats.States(STATE_NAME_PROCESS_STATE,
+                            isTracked(mTrackedUidStates, STATE_PROCESS_STATE),
+                            AggregatedPowerStatsConfig.STATE_LABELS_PROCESS_STATE),
+            };
+        }
+
+        private boolean isTracked(int[] trackedStates, int state) {
+            if (trackedStates == null) {
+                return false;
+            }
+
+            for (int trackedState : trackedStates) {
+                if (trackedState == state) {
+                    return true;
+                }
+            }
+            return false;
+        }
+    }
+
+    private final List<PowerComponent> mPowerComponents = new ArrayList<>();
+
+    /**
+     * Creates a configuration for the specified power component, which may be one of the
+     * standard power component IDs, e.g. {@link BatteryConsumer#POWER_COMPONENT_CPU}, or
+     * a custom power component.
+     */
+    public PowerComponent trackPowerComponent(int powerComponentId) {
+        PowerComponent builder = new PowerComponent(powerComponentId);
+        mPowerComponents.add(builder);
+        return builder;
+    }
+
+    public List<PowerComponent> getPowerComponentsAggregatedStatsConfigs() {
+        return mPowerComponents;
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/AggregatedPowerStatsSection.java b/services/core/java/com/android/server/power/stats/AggregatedPowerStatsSection.java
new file mode 100644
index 0000000..7ba4330
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/AggregatedPowerStatsSection.java
@@ -0,0 +1,48 @@
+/*
+ * 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.IndentingPrintWriter;
+
+import com.android.modules.utils.TypedXmlSerializer;
+
+import java.io.IOException;
+
+class AggregatedPowerStatsSection extends PowerStatsSpan.Section {
+    public static final String TYPE = "aggregated-power-stats";
+
+    private final AggregatedPowerStats mAggregatedPowerStats;
+
+    AggregatedPowerStatsSection(AggregatedPowerStats aggregatedPowerStats) {
+        super(TYPE);
+        mAggregatedPowerStats = aggregatedPowerStats;
+    }
+
+    public AggregatedPowerStats getAggregatedPowerStats() {
+        return mAggregatedPowerStats;
+    }
+
+    @Override
+    void write(TypedXmlSerializer serializer) throws IOException {
+        mAggregatedPowerStats.writeXml(serializer);
+    }
+
+    @Override
+    public void dump(IndentingPrintWriter ipw) {
+        mAggregatedPowerStats.dump(ipw);
+    }
+}
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 1c79edbc..a9c2bc2 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -284,7 +284,6 @@
     private final LongSparseArray<SamplingTimer> mKernelMemoryStats = new LongSparseArray<>();
     private int[] mCpuPowerBracketMap;
     private final CpuPowerStatsCollector mCpuPowerStatsCollector;
-    private final PowerStatsAggregator mPowerStatsAggregator;
 
     public LongSparseArray<SamplingTimer> getKernelMemoryStats() {
         return mKernelMemoryStats;
@@ -357,6 +356,11 @@
     protected Queue<UidToRemove> mPendingRemovedUids = new LinkedList<>();
 
     @NonNull
+    public BatteryStatsHistory getHistory() {
+        return mHistory;
+    }
+
+    @NonNull
     BatteryStatsHistory copyHistory() {
         return mHistory.copy();
     }
@@ -1744,7 +1748,6 @@
         mEnergyConsumerRetriever = null;
         mUserInfoProvider = null;
         mCpuPowerStatsCollector = null;
-        mPowerStatsAggregator = null;
     }
 
     private void init(Clock clock) {
@@ -10937,18 +10940,6 @@
                 mHandler, mBatteryStatsConfig.getPowerStatsThrottlePeriodCpu());
         mCpuPowerStatsCollector.addConsumer(this::recordPowerStats);
 
-        PowerStatsAggregator.Builder builder = new PowerStatsAggregator.Builder(mHistory);
-        builder.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_CPU)
-                .trackDeviceStates(
-                        PowerStatsAggregator.STATE_POWER,
-                        PowerStatsAggregator.STATE_SCREEN)
-                .trackUidStates(
-                        PowerStatsAggregator.STATE_POWER,
-                        PowerStatsAggregator.STATE_SCREEN,
-                        PowerStatsAggregator.STATE_PROCESS_STATE);
-
-        mPowerStatsAggregator = builder.build();
-
         mStartCount++;
         initTimersAndCounters();
         mOnBattery = mOnBatteryInternal = false;
@@ -15690,20 +15681,23 @@
     }
 
     /**
+     * Schedules an immediate (but asynchronous) collection of PowerStats samples.
+     * Callers will need to wait for the collection to complete on the handler thread.
+     */
+    public void schedulePowerStatsSampleCollection() {
+        if (mCpuPowerStatsCollector == null) {
+            return;
+        }
+        mCpuPowerStatsCollector.forceSchedule();
+    }
+
+    /**
      * Grabs one sample of PowerStats and prints it.
      */
     public void dumpStatsSample(PrintWriter pw) {
         mCpuPowerStatsCollector.collectAndDump(pw);
     }
 
-    /**
-     * Aggregates power stats between the specified times and prints them.
-     */
-    public void dumpAggregatedStats(PrintWriter pw, long startTimeMs, long endTimeMs) {
-        mPowerStatsAggregator.aggregateBatteryStats(startTimeMs, endTimeMs,
-                stats-> stats.dump(pw));
-    }
-
     private final Runnable mWriteAsyncRunnable = () -> {
         synchronized (BatteryStatsImpl.this) {
             writeSyncLocked();
diff --git a/services/core/java/com/android/server/power/stats/CpuAggregatedPowerStats.java b/services/core/java/com/android/server/power/stats/CpuAggregatedPowerStats.java
index 5b3fe06..fbf6928 100644
--- a/services/core/java/com/android/server/power/stats/CpuAggregatedPowerStats.java
+++ b/services/core/java/com/android/server/power/stats/CpuAggregatedPowerStats.java
@@ -16,14 +16,8 @@
 
 package com.android.server.power.stats;
 
-import android.os.BatteryConsumer;
-
-import com.android.internal.os.MultiStateStats;
-
 class CpuAggregatedPowerStats extends PowerComponentAggregatedPowerStats {
-
-    CpuAggregatedPowerStats(MultiStateStats.States[] deviceStates,
-            MultiStateStats.States[] uidStates) {
-        super(BatteryConsumer.POWER_COMPONENT_CPU, deviceStates, uidStates);
+    CpuAggregatedPowerStats(AggregatedPowerStatsConfig.PowerComponent config) {
+        super(config);
     }
 }
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 686268f..05c0a13 100644
--- a/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
+++ b/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
@@ -21,7 +21,13 @@
 
 import com.android.internal.os.MultiStateStats;
 import com.android.internal.os.PowerStats;
+import com.android.modules.utils.TypedXmlPullParser;
+import com.android.modules.utils.TypedXmlSerializer;
 
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
 import java.util.Collection;
 
 /**
@@ -31,6 +37,12 @@
  * as part of the {@link PowerStats.Descriptor}.
  */
 class PowerComponentAggregatedPowerStats {
+    static final String XML_TAG_POWER_COMPONENT = "power_component";
+    static final String XML_ATTR_ID = "id";
+    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";
+
     public final int powerComponentId;
     private final MultiStateStats.States[] mDeviceStateConfig;
     private final MultiStateStats.States[] mUidStateConfig;
@@ -49,22 +61,24 @@
         public MultiStateStats stats;
     }
 
-    PowerComponentAggregatedPowerStats(int powerComponentId,
-            MultiStateStats.States[] deviceStates,
-            MultiStateStats.States[] uidStates) {
-        this.powerComponentId = powerComponentId;
-        mDeviceStateConfig = deviceStates;
-        mUidStateConfig = uidStates;
+    PowerComponentAggregatedPowerStats(AggregatedPowerStatsConfig.PowerComponent config) {
+        this.powerComponentId = config.getPowerComponentId();
+        mDeviceStateConfig = config.getDeviceStateConfig();
+        mUidStateConfig = config.getUidStateConfig();
         mDeviceStates = new int[mDeviceStateConfig.length];
         mDeviceStateTimestamps = new long[mDeviceStateConfig.length];
     }
 
-    void setState(@PowerStatsAggregator.TrackedState int stateId, int state, long time) {
+    public PowerStats.Descriptor getPowerStatsDescriptor() {
+        return mPowerStatsDescriptor;
+    }
+
+    void setState(@AggregatedPowerStatsConfig.TrackedState int stateId, int state, long time) {
         mDeviceStates[stateId] = state;
         mDeviceStateTimestamps[stateId] = time;
 
         if (mDeviceStateConfig[stateId].isTracked()) {
-            if (mDeviceStats != null || createDeviceStats()) {
+            if (mDeviceStats != null) {
                 mDeviceStats.setState(stateId, state, time);
             }
         }
@@ -72,14 +86,14 @@
         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)) {
+                if (uidStats.stats != null) {
                     uidStats.stats.setState(stateId, state, time);
                 }
             }
         }
     }
 
-    void setUidState(int uid, @PowerStatsAggregator.TrackedState int stateId, int state,
+    void setUidState(int uid, @AggregatedPowerStatsConfig.TrackedState int stateId, int state,
             long time) {
         if (!mUidStateConfig[stateId].isTracked()) {
             return;
@@ -89,7 +103,7 @@
         uidStats.states[stateId] = state;
         uidStats.stateTimestampMs[stateId] = time;
 
-        if (uidStats.stats != null || createUidStats(uidStats)) {
+        if (uidStats.stats != null) {
             uidStats.stats.setState(stateId, state, time);
         }
     }
@@ -102,13 +116,6 @@
         mPowerStatsDescriptor = powerStats.descriptor;
 
         if (mDeviceStats == null) {
-            if (mStatsFactory == null) {
-                mStatsFactory = new MultiStateStats.Factory(
-                        mPowerStatsDescriptor.statsArrayLength, mDeviceStateConfig);
-                mUidStatsFactory = new MultiStateStats.Factory(
-                        mPowerStatsDescriptor.uidStatsArrayLength, mUidStateConfig);
-            }
-
             createDeviceStats();
         }
 
@@ -183,7 +190,11 @@
 
     private boolean createDeviceStats() {
         if (mStatsFactory == null) {
-            return false;
+            if (mPowerStatsDescriptor == null) {
+                return false;
+            }
+            mStatsFactory = new MultiStateStats.Factory(
+                    mPowerStatsDescriptor.statsArrayLength, mDeviceStateConfig);
         }
 
         mDeviceStats = mStatsFactory.create();
@@ -196,7 +207,11 @@
 
     private boolean createUidStats(UidStats uidStats) {
         if (mUidStatsFactory == null) {
-            return false;
+            if (mPowerStatsDescriptor == null) {
+                return false;
+            }
+            mUidStatsFactory = new MultiStateStats.Factory(
+                    mPowerStatsDescriptor.uidStatsArrayLength, mUidStateConfig);
         }
 
         uidStats.stats = mUidStatsFactory.create();
@@ -211,6 +226,74 @@
         return true;
     }
 
+    public void writeXml(TypedXmlSerializer serializer) throws IOException {
+        // No stats aggregated - can skip writing XML altogether
+        if (mPowerStatsDescriptor == null) {
+            return;
+        }
+
+        serializer.startTag(null, XML_TAG_POWER_COMPONENT);
+        serializer.attributeInt(null, XML_ATTR_ID, powerComponentId);
+        mPowerStatsDescriptor.writeXml(serializer);
+
+        if (mDeviceStats != null) {
+            serializer.startTag(null, XML_TAG_DEVICE_STATS);
+            mDeviceStats.writeXml(serializer);
+            serializer.endTag(null, XML_TAG_DEVICE_STATS);
+        }
+
+        for (int i = mUidStats.size() - 1; i >= 0; i--) {
+            int uid = mUidStats.keyAt(i);
+            UidStats uidStats = mUidStats.valueAt(i);
+            if (uidStats.stats != null) {
+                serializer.startTag(null, XML_TAG_UID_STATS);
+                serializer.attributeInt(null, XML_ATTR_UID, uid);
+                uidStats.stats.writeXml(serializer);
+                serializer.endTag(null, XML_TAG_UID_STATS);
+            }
+        }
+
+        serializer.endTag(null, XML_TAG_POWER_COMPONENT);
+        serializer.flush();
+    }
+
+    public boolean readFromXml(TypedXmlPullParser parser) throws XmlPullParserException,
+            IOException {
+        int eventType = parser.getEventType();
+        while (eventType != XmlPullParser.END_DOCUMENT) {
+            if (eventType == XmlPullParser.START_TAG) {
+                switch (parser.getName()) {
+                    case PowerStats.Descriptor.XML_TAG_DESCRIPTOR:
+                        mPowerStatsDescriptor = PowerStats.Descriptor.createFromXml(parser);
+                        if (mPowerStatsDescriptor == null) {
+                            return false;
+                        }
+                        break;
+                    case XML_TAG_DEVICE_STATS:
+                        if (mDeviceStats == null) {
+                            createDeviceStats();
+                        }
+                        if (!mDeviceStats.readFromXml(parser)) {
+                            return false;
+                        }
+                        break;
+                    case XML_TAG_UID_STATS:
+                        int uid = parser.getAttributeInt(null, XML_ATTR_UID);
+                        UidStats uidStats = getUidStats(uid);
+                        if (uidStats.stats == null) {
+                            createUidStats(uidStats);
+                        }
+                        if (!uidStats.stats.readFromXml(parser)) {
+                            return false;
+                        }
+                        break;
+                }
+            }
+            eventType = parser.next();
+        }
+        return true;
+    }
+
     void dumpDevice(IndentingPrintWriter ipw) {
         if (mDeviceStats != null) {
             ipw.println(mPowerStatsDescriptor.name);
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 cfb081d..f374fb7 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java
@@ -15,59 +15,26 @@
  */
 package com.android.server.power.stats;
 
-import android.annotation.IntDef;
-import android.os.BatteryConsumer;
 import android.os.BatteryStats;
 
 import com.android.internal.os.BatteryStatsHistory;
 import com.android.internal.os.BatteryStatsHistoryIterator;
-import com.android.internal.os.MultiStateStats;
 
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-import java.util.ArrayList;
-import java.util.List;
 import java.util.function.Consumer;
 
-class PowerStatsAggregator {
-    public static final int STATE_POWER = 0;
-    public static final int STATE_SCREEN = 1;
-    public static final int STATE_PROCESS_STATE = 2;
-
-    @IntDef({
-            STATE_POWER,
-            STATE_SCREEN,
-            STATE_PROCESS_STATE,
-    })
-    @Retention(RetentionPolicy.SOURCE)
-    public @interface TrackedState {
-    }
-
-    static final int POWER_STATE_BATTERY = 0;
-    static final int POWER_STATE_OTHER = 1;   // Plugged in, or on wireless charger, etc.
-    static final String[] STATE_LABELS_POWER = {"pwr-battery", "pwr-other"};
-
-    static final int SCREEN_STATE_ON = 0;
-    static final int SCREEN_STATE_OTHER = 1;  // Off, doze etc
-    static final String[] STATE_LABELS_SCREEN = {"scr-on", "scr-other"};
-
-    static final String[] STATE_LABELS_PROCESS_STATE;
-
-    static {
-        String[] procStateLabels = new String[BatteryConsumer.PROCESS_STATE_COUNT];
-        for (int i = 0; i < BatteryConsumer.PROCESS_STATE_COUNT; i++) {
-            procStateLabels[i] = BatteryConsumer.processStateToString(i);
-        }
-        STATE_LABELS_PROCESS_STATE = procStateLabels;
-    }
-
-    private final BatteryStatsHistory mHistory;
+/**
+ * Power stats aggregator. It reads through portions of battery stats history, finds
+ * relevant items (state changes, power stats etc) and produces one or more
+ * {@link AggregatedPowerStats} that adds up power stats from the samples found in battery history.
+ */
+public class PowerStatsAggregator {
     private final AggregatedPowerStats mStats;
+    private final BatteryStatsHistory mHistory;
 
-    private PowerStatsAggregator(BatteryStatsHistory history,
-            AggregatedPowerStats aggregatedPowerStats) {
+    public PowerStatsAggregator(AggregatedPowerStatsConfig aggregatedPowerStatsConfig,
+            BatteryStatsHistory history) {
+        mStats = new AggregatedPowerStats(aggregatedPowerStatsConfig);
         mHistory = history;
-        mStats = aggregatedPowerStats;
     }
 
     /**
@@ -82,12 +49,10 @@
      * Note: the AggregatedPowerStats object is reused, so the consumer should fully consume
      * the stats in the <code>accept</code> method and never cache it.
      */
-    void aggregateBatteryStats(long startTimeMs, long endTimeMs,
+    public void aggregatePowerStats(long startTimeMs, long endTimeMs,
             Consumer<AggregatedPowerStats> consumer) {
-        mStats.reset();
-
-        int currentBatteryState = POWER_STATE_BATTERY;
-        int currentScreenState = SCREEN_STATE_OTHER;
+        int currentBatteryState = AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
+        int currentScreenState = AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
         long baseTime = -1;
         long lastTime = 0;
         try (BatteryStatsHistoryIterator iterator =
@@ -107,29 +72,36 @@
 
                 int batteryState =
                         (item.states & BatteryStats.HistoryItem.STATE_BATTERY_PLUGGED_FLAG) != 0
-                                ? POWER_STATE_OTHER : POWER_STATE_BATTERY;
+                                ? AggregatedPowerStatsConfig.POWER_STATE_OTHER
+                                : AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
                 if (batteryState != currentBatteryState) {
-                    mStats.setDeviceState(STATE_POWER, batteryState, item.time);
+                    mStats.setDeviceState(AggregatedPowerStatsConfig.STATE_POWER, batteryState,
+                            item.time);
                     currentBatteryState = batteryState;
                 }
 
                 int screenState =
                         (item.states & BatteryStats.HistoryItem.STATE_SCREEN_ON_FLAG) != 0
-                                ? SCREEN_STATE_ON : SCREEN_STATE_OTHER;
+                                ? AggregatedPowerStatsConfig.SCREEN_STATE_ON
+                                : AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
                 if (screenState != currentScreenState) {
-                    mStats.setDeviceState(STATE_SCREEN, screenState, item.time);
+                    mStats.setDeviceState(AggregatedPowerStatsConfig.STATE_SCREEN, screenState,
+                            item.time);
                     currentScreenState = screenState;
                 }
 
                 if (item.processStateChange != null) {
-                    mStats.setUidState(item.processStateChange.uid, STATE_PROCESS_STATE,
+                    mStats.setUidState(item.processStateChange.uid,
+                            AggregatedPowerStatsConfig.STATE_PROCESS_STATE,
                             item.processStateChange.processState, item.time);
                 }
 
                 if (item.powerStats != null) {
                     if (!mStats.isCompatible(item.powerStats)) {
-                        mStats.setDuration(lastTime - baseTime);
-                        consumer.accept(mStats);
+                        if (lastTime > baseTime) {
+                            mStats.setDuration(lastTime - baseTime);
+                            consumer.accept(mStats);
+                        }
                         mStats.reset();
                         mStats.addClockUpdate(item.time, item.currentTime);
                         baseTime = lastTime = item.time;
@@ -138,92 +110,11 @@
                 }
             }
         }
-        mStats.setDuration(lastTime - baseTime);
-        consumer.accept(mStats);
-    }
-
-    static class Builder {
-        static class PowerComponentAggregateStatsBuilder {
-            private final int mPowerComponentId;
-            private @TrackedState int[] mTrackedDeviceStates;
-            private @TrackedState int[] mTrackedUidStates;
-
-            PowerComponentAggregateStatsBuilder(int powerComponentId) {
-                this.mPowerComponentId = powerComponentId;
-            }
-
-            public PowerComponentAggregateStatsBuilder trackDeviceStates(
-                    @TrackedState int... states) {
-                mTrackedDeviceStates = states;
-                return this;
-            }
-
-            public PowerComponentAggregateStatsBuilder trackUidStates(@TrackedState int... states) {
-                mTrackedUidStates = states;
-                return this;
-            }
-
-            private PowerComponentAggregatedPowerStats build() {
-                MultiStateStats.States[] deviceStates = new MultiStateStats.States[]{
-                        new MultiStateStats.States(isTracked(mTrackedDeviceStates, STATE_POWER),
-                                PowerStatsAggregator.STATE_LABELS_POWER),
-                        new MultiStateStats.States(isTracked(mTrackedDeviceStates, STATE_SCREEN),
-                                PowerStatsAggregator.STATE_LABELS_SCREEN),
-                };
-
-                MultiStateStats.States[] uidStates = new MultiStateStats.States[]{
-                        new MultiStateStats.States(isTracked(mTrackedUidStates, STATE_POWER),
-                                PowerStatsAggregator.STATE_LABELS_POWER),
-                        new MultiStateStats.States(isTracked(mTrackedUidStates, STATE_SCREEN),
-                                PowerStatsAggregator.STATE_LABELS_SCREEN),
-                        new MultiStateStats.States(
-                                isTracked(mTrackedUidStates, STATE_PROCESS_STATE),
-                                PowerStatsAggregator.STATE_LABELS_PROCESS_STATE),
-                };
-
-                switch (mPowerComponentId) {
-                    case BatteryConsumer.POWER_COMPONENT_CPU:
-                        return new CpuAggregatedPowerStats(deviceStates, uidStates);
-                    default:
-                        return new PowerComponentAggregatedPowerStats(mPowerComponentId,
-                                deviceStates, uidStates);
-                }
-            }
-
-            private boolean isTracked(int[] trackedStates, int state) {
-                if (trackedStates == null) {
-                    return false;
-                }
-
-                for (int trackedState : trackedStates) {
-                    if (trackedState == state) {
-                        return true;
-                    }
-                }
-                return false;
-            }
+        if (lastTime > baseTime) {
+            mStats.setDuration(lastTime - baseTime);
+            consumer.accept(mStats);
         }
 
-        private final BatteryStatsHistory mHistory;
-        private final List<PowerComponentAggregateStatsBuilder> mPowerComponents =
-                new ArrayList<>();
-
-        Builder(BatteryStatsHistory history) {
-            mHistory = history;
-        }
-
-        PowerComponentAggregateStatsBuilder trackPowerComponent(int powerComponentId) {
-            PowerComponentAggregateStatsBuilder builder = new PowerComponentAggregateStatsBuilder(
-                    powerComponentId);
-            mPowerComponents.add(builder);
-            return builder;
-        }
-
-        PowerStatsAggregator build() {
-            return new PowerStatsAggregator(mHistory, new AggregatedPowerStats(
-                    mPowerComponents.stream()
-                            .map(PowerComponentAggregateStatsBuilder::build)
-                            .toArray(PowerComponentAggregatedPowerStats[]::new)));
-        }
+        mStats.reset();     // to free up memory
     }
 }
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 93b0bac..58619c7 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsScheduler.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsScheduler.java
@@ -16,27 +16,57 @@
 
 package com.android.server.power.stats;
 
+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;
 
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.os.Clock;
 import com.android.internal.os.MonotonicClock;
 
+import java.io.PrintWriter;
+import java.util.Calendar;
+import java.util.concurrent.TimeUnit;
+
 /**
  * Controls the frequency at which {@link PowerStatsSpan}'s are generated and stored in
  * {@link PowerStatsStore}.
  */
 public class PowerStatsScheduler {
+    private static final long MINUTE_IN_MILLIS = TimeUnit.MINUTES.toMillis(1);
+    private static final long HOUR_IN_MILLIS = TimeUnit.HOURS.toMillis(1);
+
+    private final Context mContext;
+    private boolean mEnablePeriodicPowerStatsCollection;
+    @DurationMillisLong
+    private final long mAggregatedPowerStatsSpanDuration;
+    @DurationMillisLong
+    private final long mPowerStatsAggregationPeriod;
     private final PowerStatsStore mPowerStatsStore;
+    private final Clock mClock;
     private final MonotonicClock mMonotonicClock;
     private final Handler mHandler;
     private final BatteryStatsImpl mBatteryStats;
     private final BatteryUsageStatsProvider mBatteryUsageStatsProvider;
+    private final PowerStatsAggregator mPowerStatsAggregator;
+    private long mLastSavedSpanEndMonotonicTime;
 
-    public PowerStatsScheduler(PowerStatsStore powerStatsStore, MonotonicClock monotonicClock,
-            Handler handler, BatteryStatsImpl batteryStats,
-            BatteryUsageStatsProvider batteryUsageStatsProvider) {
+    public PowerStatsScheduler(Context context, PowerStatsAggregator powerStatsAggregator,
+            @DurationMillisLong long aggregatedPowerStatsSpanDuration,
+            @DurationMillisLong long powerStatsAggregationPeriod, PowerStatsStore powerStatsStore,
+            Clock clock, MonotonicClock monotonicClock, Handler handler,
+            BatteryStatsImpl batteryStats, BatteryUsageStatsProvider batteryUsageStatsProvider) {
+        mContext = context;
+        mPowerStatsAggregator = powerStatsAggregator;
+        mAggregatedPowerStatsSpanDuration = aggregatedPowerStatsSpanDuration;
+        mPowerStatsAggregationPeriod = powerStatsAggregationPeriod;
         mPowerStatsStore = powerStatsStore;
+        mClock = clock;
         mMonotonicClock = monotonicClock;
         mHandler = handler;
         mBatteryStats = batteryStats;
@@ -46,8 +76,159 @@
     /**
      * Kicks off the scheduling of power stats aggregation spans.
      */
-    public void start() {
+    public void start(boolean enablePeriodicPowerStatsCollection) {
         mBatteryStats.setBatteryResetListener(this::storeBatteryUsageStatsOnReset);
+        mEnablePeriodicPowerStatsCollection = enablePeriodicPowerStatsCollection;
+        if (mEnablePeriodicPowerStatsCollection) {
+            scheduleNextPowerStatsAggregation();
+        }
+    }
+
+    private void scheduleNextPowerStatsAggregation() {
+        AlarmManager alarmManager = mContext.getSystemService(AlarmManager.class);
+        alarmManager.set(AlarmManager.ELAPSED_REALTIME,
+                mClock.elapsedRealtime() + mPowerStatsAggregationPeriod, "PowerStats",
+                () -> {
+                    schedulePowerStatsAggregation();
+                    mHandler.post(this::scheduleNextPowerStatsAggregation);
+                }, mHandler);
+    }
+
+    /**
+     * Initiate an asynchronous process of aggregation of power stats.
+     */
+    @VisibleForTesting
+    public void schedulePowerStatsAggregation() {
+        // Catch up the power stats collectors
+        mBatteryStats.schedulePowerStatsSampleCollection();
+        mHandler.post(this::aggregateAndStorePowerStats);
+    }
+
+    private void aggregateAndStorePowerStats() {
+        long currentTimeMillis = mClock.currentTimeMillis();
+        long currentMonotonicTime = mMonotonicClock.monotonicTime();
+        long startTime = getLastSavedSpanEndMonotonicTime();
+        long endTimeMs = alignToWallClock(startTime + mAggregatedPowerStatsSpanDuration,
+                mAggregatedPowerStatsSpanDuration, currentMonotonicTime, currentTimeMillis);
+        while (endTimeMs <= currentMonotonicTime) {
+            mPowerStatsAggregator.aggregatePowerStats(startTime, endTimeMs,
+                    stats -> {
+                        storeAggregatedPowerStats(stats);
+                        mLastSavedSpanEndMonotonicTime = stats.getStartTime() + stats.getDuration();
+                    });
+
+            startTime = endTimeMs;
+            endTimeMs += mAggregatedPowerStatsSpanDuration;
+        }
+    }
+
+    /**
+     * Performs a power stats aggregation pass and then dumps all stored aggregated power stats
+     * spans followed by the remainder that has not been stored yet.
+     */
+    public void aggregateAndDumpPowerStats(PrintWriter pw) {
+        if (mHandler.getLooper().isCurrentThread()) {
+            throw new IllegalStateException("Should not be executed on the bg handler thread.");
+        }
+
+        schedulePowerStatsAggregation();
+
+        // Wait for the aggregation process to finish storing aggregated stats spans in the store.
+        awaitCompletion();
+
+        IndentingPrintWriter ipw = new IndentingPrintWriter(pw);
+        mHandler.post(() -> {
+            mPowerStatsStore.dump(ipw);
+            // Aggregate the remainder of power stats and dump the results without storing them yet.
+            long powerStoreEndMonotonicTime = getLastSavedSpanEndMonotonicTime();
+            mPowerStatsAggregator.aggregatePowerStats(powerStoreEndMonotonicTime, 0,
+                    stats -> {
+                        // Create a PowerStatsSpan for consistency of the textual output
+                        PowerStatsSpan span = PowerStatsStore.createPowerStatsSpan(stats);
+                        if (span != null) {
+                            span.dump(ipw);
+                        }
+                    });
+        });
+
+        awaitCompletion();
+    }
+
+    /**
+     * Align the supplied time to the wall clock, for aesthetic purposes. For example, if
+     * the schedule is configured with a 15-min interval, the captured aggregated stats will
+     * be for spans XX:00-XX:15, XX:15-XX:30, XX:30-XX:45 and XX:45-XX:60. Only the current
+     * time is used for the alignment, so if the wall clock changed during an aggregation span,
+     * or if the device was off (which stops the monotonic clock), the alignment may be
+     * temporarily broken.
+     */
+    @VisibleForTesting
+    public static long alignToWallClock(long targetMonotonicTime, long interval,
+            long currentMonotonicTime, long currentTimeMillis) {
+
+        // Estimate the wall clock time for the requested targetMonotonicTime
+        long targetWallClockTime = currentTimeMillis + (targetMonotonicTime - currentMonotonicTime);
+
+        if (interval >= MINUTE_IN_MILLIS && TimeUnit.HOURS.toMillis(1) % interval == 0) {
+            // If the interval is a divisor of an hour, e.g. 10 minutes, 15 minutes, etc
+
+            // First, round up to the next whole minute
+            Calendar cal = Calendar.getInstance();
+            cal.setTimeInMillis(targetWallClockTime + MINUTE_IN_MILLIS - 1);
+            cal.set(Calendar.SECOND, 0);
+            cal.set(Calendar.MILLISECOND, 0);
+
+            // Now set the minute to a multiple of the requested interval
+            int intervalInMinutes = (int) (interval / MINUTE_IN_MILLIS);
+            cal.set(Calendar.MINUTE,
+                    ((cal.get(Calendar.MINUTE) + intervalInMinutes - 1) / intervalInMinutes)
+                            * intervalInMinutes);
+
+            long adjustment = cal.getTimeInMillis() - targetWallClockTime;
+            return targetMonotonicTime + adjustment;
+        } else if (interval >= HOUR_IN_MILLIS && TimeUnit.DAYS.toMillis(1) % interval == 0) {
+            // If the interval is a divisor of a day, e.g. 2h, 3h, etc
+
+            // First, round up to the next whole hour
+            Calendar cal = Calendar.getInstance();
+            cal.setTimeInMillis(targetWallClockTime + HOUR_IN_MILLIS - 1);
+            cal.set(Calendar.MINUTE, 0);
+            cal.set(Calendar.SECOND, 0);
+            cal.set(Calendar.MILLISECOND, 0);
+
+            // Now set the hour of day to a multiple of the requested interval
+            int intervalInHours = (int) (interval / HOUR_IN_MILLIS);
+            cal.set(Calendar.HOUR_OF_DAY,
+                    ((cal.get(Calendar.HOUR_OF_DAY) + intervalInHours - 1) / intervalInHours)
+                    * intervalInHours);
+
+            long adjustment = cal.getTimeInMillis() - targetWallClockTime;
+            return targetMonotonicTime + adjustment;
+        }
+
+        return targetMonotonicTime;
+    }
+
+    private long getLastSavedSpanEndMonotonicTime() {
+        if (mLastSavedSpanEndMonotonicTime != 0) {
+            return mLastSavedSpanEndMonotonicTime;
+        }
+
+        for (PowerStatsSpan.Metadata metadata : mPowerStatsStore.getTableOfContents()) {
+            if (metadata.getSections().contains(AggregatedPowerStatsSection.TYPE)) {
+                for (PowerStatsSpan.TimeFrame timeFrame : metadata.getTimeFrames()) {
+                    long endMonotonicTime = timeFrame.startMonotonicTime + timeFrame.duration;
+                    if (endMonotonicTime > mLastSavedSpanEndMonotonicTime) {
+                        mLastSavedSpanEndMonotonicTime = endMonotonicTime;
+                    }
+                }
+            }
+        }
+        return mLastSavedSpanEndMonotonicTime;
+    }
+
+    private void storeAggregatedPowerStats(AggregatedPowerStats stats) {
+        mPowerStatsStore.storeAggregatedPowerStats(stats);
     }
 
     private void storeBatteryUsageStatsOnReset(int resetReason) {
@@ -63,11 +244,18 @@
                                 .includeProcessStateData()
                                 .build());
 
-        // TODO(b/188068523): BatteryUsageStats should contain monotonic time for start and end
-        // When that is done, we will be able to use the BatteryUsageStats' monotonic start time
+        // 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);
+        done.block();
+    }
 }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsStore.java b/services/core/java/com/android/server/power/stats/PowerStatsStore.java
index ec84fd7..7123bcb 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsStore.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsStore.java
@@ -41,6 +41,7 @@
 import java.nio.charset.StandardCharsets;
 import java.nio.file.StandardOpenOption;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -66,16 +67,12 @@
     private final long mMaxStorageBytes;
     private final Handler mHandler;
     private final PowerStatsSpan.SectionReader mSectionReader;
-    private final List<PowerStatsSpan.Metadata> mTableOfContents = new ArrayList<>();
-    private boolean mTableOfContentsLoaded;
+    private volatile List<PowerStatsSpan.Metadata> mTableOfContents;
 
-    public PowerStatsStore(@NonNull File systemDir, Handler handler) {
-        this(systemDir, MAX_POWER_STATS_SPAN_STORAGE_BYTES, handler);
-    }
-
-    @VisibleForTesting
-    public PowerStatsStore(@NonNull File systemDir, long maxStorageBytes, Handler handler) {
-        this(systemDir, maxStorageBytes, handler, new DefaultSectionReader());
+    public PowerStatsStore(@NonNull File systemDir, Handler handler,
+            AggregatedPowerStatsConfig aggregatedPowerStatsConfig) {
+        this(systemDir, MAX_POWER_STATS_SPAN_STORAGE_BYTES, handler,
+                new DefaultSectionReader(aggregatedPowerStatsConfig));
     }
 
     @VisibleForTesting
@@ -95,13 +92,15 @@
      */
     @NonNull
     public List<PowerStatsSpan.Metadata> getTableOfContents() {
-        if (mTableOfContentsLoaded) {
-            return mTableOfContents;
+        List<PowerStatsSpan.Metadata> toc = mTableOfContents;
+        if (toc != null) {
+            return toc;
         }
 
         TypedXmlPullParser parser = Xml.newBinaryPullParser();
         lockStoreDirectory();
         try {
+            toc = new ArrayList<>();
             for (File file : mStoreDir.listFiles()) {
                 String fileName = file.getName();
                 if (!fileName.endsWith(POWER_STATS_SPAN_FILE_EXTENSION)) {
@@ -111,7 +110,7 @@
                     parser.setInput(inputStream, StandardCharsets.UTF_8.name());
                     PowerStatsSpan.Metadata metadata = PowerStatsSpan.Metadata.read(parser);
                     if (metadata != null) {
-                        mTableOfContents.add(metadata);
+                        toc.add(metadata);
                     } else {
                         Slog.e(TAG, "Removing incompatible PowerStatsSpan file: " + fileName);
                         file.delete();
@@ -120,13 +119,13 @@
                     Slog.wtf(TAG, "Cannot read PowerStatsSpan file: " + fileName);
                 }
             }
-            mTableOfContentsLoaded = true;
+            toc.sort(PowerStatsSpan.Metadata.COMPARATOR);
+            mTableOfContents = Collections.unmodifiableList(toc);
         } finally {
             unlockStoreDirectory();
         }
 
-        mTableOfContents.sort(PowerStatsSpan.Metadata.COMPARATOR);
-        return mTableOfContents;
+        return toc;
     }
 
     /**
@@ -152,10 +151,7 @@
                     throw new RuntimeException(e);
                 }
             });
-            if (mTableOfContentsLoaded) {
-                mTableOfContents.add(span.getMetadata());
-                mTableOfContents.sort(PowerStatsSpan.Metadata.COMPARATOR);
-            }
+            mTableOfContents = null;
             removeOldSpansLocked();
         } finally {
             unlockStoreDirectory();
@@ -183,6 +179,41 @@
         return null;
     }
 
+    void storeAggregatedPowerStats(AggregatedPowerStats stats) {
+        PowerStatsSpan span = createPowerStatsSpan(stats);
+        if (span == null) {
+            return;
+        }
+        storePowerStatsSpan(span);
+    }
+
+    static PowerStatsSpan createPowerStatsSpan(AggregatedPowerStats stats) {
+        List<AggregatedPowerStats.ClockUpdate> clockUpdates = stats.getClockUpdates();
+        if (clockUpdates.isEmpty()) {
+            Slog.w(TAG, "No clock updates in aggregated power stats " + stats);
+            return null;
+        }
+
+        long monotonicTime = clockUpdates.get(0).monotonicTime;
+        long durationSum = 0;
+        PowerStatsSpan span = new PowerStatsSpan(monotonicTime);
+        for (int i = 0; i < clockUpdates.size(); i++) {
+            AggregatedPowerStats.ClockUpdate clockUpdate = clockUpdates.get(i);
+            long duration;
+            if (i == clockUpdates.size() - 1) {
+                duration = stats.getDuration() - durationSum;
+            } else {
+                duration = clockUpdate.monotonicTime - monotonicTime;
+            }
+            span.addTimeFrame(clockUpdate.monotonicTime, clockUpdate.currentTime, duration);
+            monotonicTime = clockUpdate.monotonicTime;
+            durationSum += duration;
+        }
+
+        span.addSection(new AggregatedPowerStatsSection(stats));
+        return span;
+    }
+
     /**
      * Stores a {@link PowerStatsSpan} containing a single section for the supplied
      * battery usage stats.
@@ -261,7 +292,7 @@
             }
             totalSize -= entry.getValue();
             mFileSizes.remove(file);
-            mTableOfContentsLoaded = false;
+            mTableOfContents = null;
         }
     }
 
@@ -278,8 +309,7 @@
                     }
                 }
             }
-            mTableOfContents.clear();
-            mTableOfContentsLoaded = false;
+            mTableOfContents = List.of();
         } finally {
             unlockStoreDirectory();
         }
@@ -316,13 +346,26 @@
     }
 
     private static class DefaultSectionReader implements PowerStatsSpan.SectionReader {
+        private final AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
+
+        DefaultSectionReader(AggregatedPowerStatsConfig aggregatedPowerStatsConfig) {
+            mAggregatedPowerStatsConfig = aggregatedPowerStatsConfig;
+        }
+
         @Override
         public PowerStatsSpan.Section read(String sectionType, TypedXmlPullParser parser)
                 throws IOException, XmlPullParserException {
-            if (BatteryUsageStatsSection.TYPE.equals(sectionType)) {
-                return new BatteryUsageStatsSection(BatteryUsageStats.createFromXml(parser));
+            switch (sectionType) {
+                case AggregatedPowerStatsSection.TYPE:
+                    return new AggregatedPowerStatsSection(
+                            AggregatedPowerStats.createFromXml(parser,
+                                    mAggregatedPowerStatsConfig));
+                case BatteryUsageStatsSection.TYPE:
+                    return new BatteryUsageStatsSection(
+                            BatteryUsageStats.createFromXml(parser));
+                default:
+                    return null;
             }
-            return null;
         }
     }
 }
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
new file mode 100644
index 0000000..2003d04
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java
@@ -0,0 +1,201 @@
+/*
+ * 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 android.os.BatteryConsumer;
+import android.os.PersistableBundle;
+import android.util.Xml;
+
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.internal.os.PowerStats;
+import com.android.modules.utils.TypedXmlPullParser;
+import com.android.modules.utils.TypedXmlSerializer;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.text.ParseException;
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class AggregatedPowerStatsTest {
+    private static final int TEST_POWER_COMPONENT = 1077;
+    private static final int APP_1 = 27;
+    private static final int APP_2 = 42;
+
+    private AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
+    private PowerStats.Descriptor mPowerComponentDescriptor;
+
+    @Before
+    public void setup() throws ParseException {
+        mAggregatedPowerStatsConfig = new AggregatedPowerStatsConfig();
+        mAggregatedPowerStatsConfig.trackPowerComponent(TEST_POWER_COMPONENT)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE);
+
+        mPowerComponentDescriptor = new PowerStats.Descriptor(TEST_POWER_COMPONENT, "fan", 2, 3,
+                PersistableBundle.forPair("speed", "fast"));
+    }
+
+    @Test
+    public void aggregation() {
+        AggregatedPowerStats stats = prepareAggregatePowerStats();
+
+        verifyAggregatedPowerStats(stats);
+    }
+
+    @Test
+    public void xmlPersistence() throws Exception {
+        AggregatedPowerStats stats = prepareAggregatePowerStats();
+
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        TypedXmlSerializer serializer = Xml.newFastSerializer();
+        serializer.setOutput(baos, "UTF-8");
+        stats.writeXml(serializer);
+        serializer.flush();
+
+        TypedXmlPullParser parser = Xml.newFastPullParser();
+        parser.setInput(new ByteArrayInputStream(baos.toByteArray()), "UTF-8");
+        AggregatedPowerStats actualStats = AggregatedPowerStats.createFromXml(parser,
+                mAggregatedPowerStatsConfig);
+
+        verifyAggregatedPowerStats(actualStats);
+    }
+
+    private AggregatedPowerStats prepareAggregatePowerStats() {
+        AggregatedPowerStats stats = new AggregatedPowerStats(mAggregatedPowerStatsConfig);
+        stats.addClockUpdate(1000, 456);
+        stats.setDuration(789);
+
+        stats.setDeviceState(AggregatedPowerStatsConfig.STATE_SCREEN,
+                AggregatedPowerStatsConfig.SCREEN_STATE_ON, 2000);
+        stats.setUidState(APP_1, AggregatedPowerStatsConfig.STATE_PROCESS_STATE,
+                BatteryConsumer.PROCESS_STATE_CACHED, 2000);
+        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;
+
+        ps.uidStats.put(APP_1, new long[]{389, 0, 739});
+        ps.uidStats.put(APP_2, new long[]{278, 314, 628});
+
+        stats.addPowerStats(ps, 3000);
+
+        stats.setDeviceState(AggregatedPowerStatsConfig.STATE_SCREEN,
+                AggregatedPowerStatsConfig.SCREEN_STATE_OTHER, 4000);
+        stats.setUidState(APP_2, AggregatedPowerStatsConfig.STATE_PROCESS_STATE,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND, 4000);
+
+        ps.stats[0] = 444;
+        ps.stats[1] = 0;
+
+        ps.uidStats.put(APP_1, new long[]{0, 0, 400});
+        ps.uidStats.put(APP_2, new long[]{100, 200, 300});
+
+        stats.addPowerStats(ps, 5000);
+
+        return stats;
+    }
+
+    private void verifyAggregatedPowerStats(AggregatedPowerStats stats) {
+        PowerStats.Descriptor descriptor = stats.getPowerComponentStats(TEST_POWER_COMPONENT)
+                .getPowerStatsDescriptor();
+        assertThat(descriptor.powerComponentId).isEqualTo(TEST_POWER_COMPONENT);
+        assertThat(descriptor.name).isEqualTo("fan");
+        assertThat(descriptor.statsArrayLength).isEqualTo(2);
+        assertThat(descriptor.uidStatsArrayLength).isEqualTo(3);
+        assertThat(descriptor.extras.getString("speed")).isEqualTo("fast");
+
+        assertThat(getDeviceStats(stats,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_ON))
+                .isEqualTo(new long[]{322, 987});
+
+        assertThat(getDeviceStats(stats,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_OTHER))
+                .isEqualTo(new long[]{222, 0});
+
+        assertThat(getUidDeviceStats(stats,
+                APP_1,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_ON,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED))
+                .isEqualTo(new long[]{259, 0, 492});
+
+        assertThat(getUidDeviceStats(stats,
+                APP_1,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_ON,
+                BatteryConsumer.PROCESS_STATE_CACHED))
+                .isEqualTo(new long[]{129, 0, 446});
+
+        assertThat(getUidDeviceStats(stats,
+                APP_1,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_OTHER,
+                BatteryConsumer.PROCESS_STATE_CACHED))
+                .isEqualTo(new long[]{0, 0, 200});
+
+        assertThat(getUidDeviceStats(stats,
+                APP_2,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_ON,
+                BatteryConsumer.PROCESS_STATE_UNSPECIFIED))
+                .isEqualTo(new long[]{185, 209, 418});
+
+        assertThat(getUidDeviceStats(stats,
+                APP_2,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_ON,
+                BatteryConsumer.PROCESS_STATE_FOREGROUND))
+                .isEqualTo(new long[]{142, 204, 359});
+
+        assertThat(getUidDeviceStats(stats,
+                APP_2,
+                AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                AggregatedPowerStatsConfig.SCREEN_STATE_OTHER,
+                BatteryConsumer.PROCESS_STATE_BACKGROUND))
+                .isEqualTo(new long[]{50, 100, 150});
+    }
+
+    private static long[] getDeviceStats(AggregatedPowerStats stats, int... states) {
+        long[] out = new long[states.length];
+        stats.getPowerComponentStats(TEST_POWER_COMPONENT).getDeviceStats(out, states);
+        return out;
+    }
+
+    private static long[] getUidDeviceStats(AggregatedPowerStats stats, int uid, int... states) {
+        long[] out = new long[states.length];
+        stats.getPowerComponentStats(TEST_POWER_COMPONENT).getUidStats(out, uid, states);
+        return out;
+    }
+}
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 0e14aff..b1da1fc 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
@@ -381,7 +381,7 @@
 
         PowerStatsStore powerStatsStore = new PowerStatsStore(
                 new File(context.getCacheDir(), "BatteryUsageStatsProviderTest"),
-                Integer.MAX_VALUE, new TestHandler());
+                new TestHandler(), null);
 
         BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context,
                 batteryStats, powerStatsStore);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/MultiStateStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/MultiStateStatsTest.java
index 4ecee9f..30a73181 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/MultiStateStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/MultiStateStatsTest.java
@@ -122,7 +122,7 @@
         // 4 * 10 = 40 bits needed to represent the composite state
         MultiStateStats.States[] states = new MultiStateStats.States[10];
         for (int i = 0; i < states.length; i++) {
-            states[i] = new MultiStateStats.States(true, labels);
+            states[i] = new MultiStateStats.States("foo", true, labels);
         }
         IllegalArgumentException e = assertThrows(IllegalArgumentException.class,
                 () -> new MultiStateStats.Factory(DIMENSION_COUNT, states));
@@ -191,8 +191,8 @@
     private static MultiStateStats.Factory makeFactory(boolean trackBatteryState,
             boolean trackProcState, boolean trackScreenState) {
         return new MultiStateStats.Factory(DIMENSION_COUNT,
-                new MultiStateStats.States(trackBatteryState, "plugged-in", "on-battery"),
-                new MultiStateStats.States(trackProcState,
+                new MultiStateStats.States("bs", trackBatteryState, "plugged-in", "on-battery"),
+                new MultiStateStats.States("ps", trackProcState,
                         BatteryConsumer.processStateToString(
                                 BatteryConsumer.PROCESS_STATE_UNSPECIFIED),
                         BatteryConsumer.processStateToString(
@@ -203,7 +203,7 @@
                                 BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE),
                         BatteryConsumer.processStateToString(
                                 BatteryConsumer.PROCESS_STATE_CACHED)),
-                new MultiStateStats.States(trackScreenState, "screen-off", "plugged-in"));
+                new MultiStateStats.States("scr", trackScreenState, "screen-off", "plugged-in"));
     }
 
     private FactorySubject assertThatCpuPerformanceStatsFactory(
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 5bb3885..b52fc8a 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
@@ -62,16 +62,16 @@
                 mock(BatteryStatsHistory.HistoryStepDetailsCalculator.class), mClock,
                 mMonotonicClock);
 
-        PowerStatsAggregator.Builder builder = new PowerStatsAggregator.Builder(mHistory);
-        builder.trackPowerComponent(TEST_POWER_COMPONENT)
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(TEST_POWER_COMPONENT)
                 .trackDeviceStates(
-                        PowerStatsAggregator.STATE_POWER,
-                        PowerStatsAggregator.STATE_SCREEN)
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
                 .trackUidStates(
-                        PowerStatsAggregator.STATE_POWER,
-                        PowerStatsAggregator.STATE_SCREEN,
-                        PowerStatsAggregator.STATE_PROCESS_STATE);
-        mAggregator = builder.build();
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE);
+        mAggregator = new PowerStatsAggregator(config, mHistory);
     }
 
     @Test
@@ -115,7 +115,7 @@
         powerStats.uidStats.put(TEST_UID, new long[]{4444});
         mHistory.recordPowerStats(mClock.realtime, mClock.uptime, powerStats);
 
-        mAggregator.aggregateBatteryStats(0, 0, stats -> {
+        mAggregator.aggregatePowerStats(0, 0, stats -> {
             assertThat(mAggregatedStatsCount++).isEqualTo(0);
             assertThat(stats.getStartTime()).isEqualTo(START_TIME);
 
@@ -138,34 +138,34 @@
                     TEST_POWER_COMPONENT);
 
             assertThat(powerComponentStats.getDeviceStats(values, new int[]{
-                    PowerStatsAggregator.POWER_STATE_OTHER,
-                    PowerStatsAggregator.SCREEN_STATE_ON}))
+                    AggregatedPowerStatsConfig.POWER_STATE_OTHER,
+                    AggregatedPowerStatsConfig.SCREEN_STATE_ON}))
                     .isTrue();
             assertThat(values).isEqualTo(new long[]{10000});
 
             assertThat(powerComponentStats.getDeviceStats(values, new int[]{
-                    PowerStatsAggregator.POWER_STATE_BATTERY,
-                    PowerStatsAggregator.SCREEN_STATE_OTHER}))
+                    AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                    AggregatedPowerStatsConfig.SCREEN_STATE_OTHER}))
                     .isTrue();
             assertThat(values).isEqualTo(new long[]{20000});
 
             assertThat(powerComponentStats.getUidStats(values, TEST_UID, new int[]{
-                    PowerStatsAggregator.POWER_STATE_OTHER,
-                    PowerStatsAggregator.SCREEN_STATE_ON,
+                    AggregatedPowerStatsConfig.POWER_STATE_OTHER,
+                    AggregatedPowerStatsConfig.SCREEN_STATE_ON,
                     BatteryConsumer.PROCESS_STATE_FOREGROUND}))
                     .isTrue();
             assertThat(values).isEqualTo(new long[]{1234});
 
             assertThat(powerComponentStats.getUidStats(values, TEST_UID, new int[]{
-                    PowerStatsAggregator.POWER_STATE_BATTERY,
-                    PowerStatsAggregator.SCREEN_STATE_OTHER,
+                    AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                    AggregatedPowerStatsConfig.SCREEN_STATE_OTHER,
                     BatteryConsumer.PROCESS_STATE_FOREGROUND}))
                     .isTrue();
             assertThat(values).isEqualTo(new long[]{1111});
 
             assertThat(powerComponentStats.getUidStats(values, TEST_UID, new int[]{
-                    PowerStatsAggregator.POWER_STATE_BATTERY,
-                    PowerStatsAggregator.SCREEN_STATE_OTHER,
+                    AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                    AggregatedPowerStatsConfig.SCREEN_STATE_OTHER,
                     BatteryConsumer.PROCESS_STATE_BACKGROUND}))
                     .isTrue();
             assertThat(values).isEqualTo(new long[]{3333});
@@ -211,7 +211,7 @@
 
         mHistory.recordBatteryState(mClock.realtime, mClock.uptime, 50, /* plugged */ true);
 
-        mAggregator.aggregateBatteryStats(0, 0, stats -> {
+        mAggregator.aggregatePowerStats(0, 0, stats -> {
             long[] values = new long[1];
 
             PowerComponentAggregatedPowerStats powerComponentStats =
@@ -222,13 +222,13 @@
                 assertThat(stats.getDuration()).isEqualTo(2000);
 
                 assertThat(powerComponentStats.getDeviceStats(values, new int[]{
-                        PowerStatsAggregator.POWER_STATE_OTHER,
-                        PowerStatsAggregator.SCREEN_STATE_ON}))
+                        AggregatedPowerStatsConfig.POWER_STATE_OTHER,
+                        AggregatedPowerStatsConfig.SCREEN_STATE_ON}))
                         .isTrue();
                 assertThat(values).isEqualTo(new long[]{10000});
                 assertThat(powerComponentStats.getUidStats(values, TEST_UID, new int[]{
-                        PowerStatsAggregator.POWER_STATE_OTHER,
-                        PowerStatsAggregator.SCREEN_STATE_ON,
+                        AggregatedPowerStatsConfig.POWER_STATE_OTHER,
+                        AggregatedPowerStatsConfig.SCREEN_STATE_ON,
                         BatteryConsumer.PROCESS_STATE_FOREGROUND}))
                         .isTrue();
                 assertThat(values).isEqualTo(new long[]{1234});
@@ -237,13 +237,13 @@
                 assertThat(stats.getDuration()).isEqualTo(1000);
 
                 assertThat(powerComponentStats.getDeviceStats(values, new int[]{
-                        PowerStatsAggregator.POWER_STATE_BATTERY,
-                        PowerStatsAggregator.SCREEN_STATE_ON}))
+                        AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                        AggregatedPowerStatsConfig.SCREEN_STATE_ON}))
                         .isTrue();
                 assertThat(values).isEqualTo(new long[]{20000});
                 assertThat(powerComponentStats.getUidStats(values, TEST_UID, new int[]{
-                        PowerStatsAggregator.POWER_STATE_BATTERY,
-                        PowerStatsAggregator.SCREEN_STATE_ON,
+                        AggregatedPowerStatsConfig.POWER_STATE_BATTERY,
+                        AggregatedPowerStatsConfig.SCREEN_STATE_ON,
                         BatteryConsumer.PROCESS_STATE_FOREGROUND}))
                         .isTrue();
                 assertThat(values).isEqualTo(new long[]{4444});
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 87d09f3..0e58787 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
@@ -18,7 +18,12 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
@@ -40,7 +45,12 @@
 import org.junit.runner.RunWith;
 
 import java.io.File;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.List;
+import java.util.TimeZone;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 @RunWith(AndroidJUnit4.class)
 public class PowerStatsSchedulerTest {
@@ -52,25 +62,117 @@
     private BatteryUsageStatsProvider mBatteryUsageStatsProvider;
     private PowerStatsScheduler mPowerStatsScheduler;
     private PowerProfile mPowerProfile;
+    private PowerStatsAggregator mPowerStatsAggregator;
+    private AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
 
     @Before
     public void setup() {
         final Context context = InstrumentationRegistry.getContext();
 
-        mClock.currentTime = 1234567;
+        TimeZone.setDefault(TimeZone.getTimeZone("UTC"));
+
+        mClock.currentTime = Instant.parse("2023-01-02T03:04:05.00Z").toEpochMilli();
         mClock.realtime = 7654321;
 
         HandlerThread bgThread = new HandlerThread("bg thread");
         bgThread.start();
         File systemDir = context.getCacheDir();
         mHandler = new Handler(bgThread.getLooper());
-        mPowerStatsStore = new PowerStatsStore(systemDir, mHandler);
+        mAggregatedPowerStatsConfig = new AggregatedPowerStatsConfig();
+        mPowerStatsStore = new PowerStatsStore(systemDir, mHandler, mAggregatedPowerStatsConfig);
         mPowerProfile = mock(PowerProfile.class);
         when(mPowerProfile.getAveragePower(PowerProfile.POWER_FLASHLIGHT)).thenReturn(1000000.0);
         mBatteryStats = new MockBatteryStatsImpl(mClock).setPowerProfile(mPowerProfile);
         mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context, mBatteryStats);
-        mPowerStatsScheduler = new PowerStatsScheduler(mPowerStatsStore, mMonotonicClock, mHandler,
-                mBatteryStats, mBatteryUsageStatsProvider);
+        mPowerStatsAggregator = mock(PowerStatsAggregator.class);
+        mPowerStatsScheduler = new PowerStatsScheduler(context, mPowerStatsAggregator,
+                TimeUnit.MINUTES.toMillis(30), TimeUnit.HOURS.toMillis(1), mPowerStatsStore, mClock,
+                mMonotonicClock, mHandler, mBatteryStats, mBatteryUsageStatsProvider);
+    }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    public void storeAggregatePowerStats() {
+        mPowerStatsStore.reset();
+
+        assertThat(mPowerStatsStore.getTableOfContents()).isEmpty();
+
+        mPowerStatsStore.storeAggregatedPowerStats(
+                createAggregatedPowerStats(mMonotonicClock.monotonicTime(), mClock.currentTime,
+                        123));
+
+        long delayBeforeAggregating = TimeUnit.MINUTES.toMillis(90);
+        mClock.realtime += delayBeforeAggregating;
+        mClock.currentTime += delayBeforeAggregating;
+
+        doAnswer(invocation -> {
+            // The first span is longer than 30 min, because the end time is being aligned with
+            // the wall clock.  Subsequent spans should be precisely 30 minutes.
+            long startTime = invocation.getArgument(0);
+            long endTime = invocation.getArgument(1);
+            Consumer<AggregatedPowerStats> consumer = invocation.getArgument(2);
+
+            long startTimeWallClock =
+                    mClock.currentTime - (mMonotonicClock.monotonicTime() - startTime);
+            long endTimeWallClock =
+                    mClock.currentTime - (mMonotonicClock.monotonicTime() - endTime);
+
+            assertThat(startTime).isEqualTo(7654321 + 123);
+            assertThat(endTime - startTime).isAtLeast(TimeUnit.MINUTES.toMillis(30));
+            assertThat(Instant.ofEpochMilli(endTimeWallClock))
+                    .isEqualTo(Instant.parse("2023-01-02T04:00:00Z"));
+
+            consumer.accept(
+                    createAggregatedPowerStats(startTime, startTimeWallClock, endTime - startTime));
+            return null;
+        }).doAnswer(invocation -> {
+            long startTime = invocation.getArgument(0);
+            long endTime = invocation.getArgument(1);
+            Consumer<AggregatedPowerStats> consumer = invocation.getArgument(2);
+
+            long startTimeWallClock =
+                    mClock.currentTime - (mMonotonicClock.monotonicTime() - startTime);
+            long endTimeWallClock =
+                    mClock.currentTime - (mMonotonicClock.monotonicTime() - endTime);
+
+            assertThat(Instant.ofEpochMilli(startTimeWallClock))
+                    .isEqualTo(Instant.parse("2023-01-02T04:00:00Z"));
+            assertThat(Instant.ofEpochMilli(endTimeWallClock))
+                    .isEqualTo(Instant.parse("2023-01-02T04:30:00Z"));
+
+            consumer.accept(
+                    createAggregatedPowerStats(startTime, startTimeWallClock, endTime - startTime));
+            return null;
+        }).when(mPowerStatsAggregator).aggregatePowerStats(anyLong(), anyLong(),
+                any(Consumer.class));
+
+        mPowerStatsScheduler.schedulePowerStatsAggregation();
+        ConditionVariable done = new ConditionVariable();
+        mHandler.post(done::open);
+        done.block();
+
+        verify(mPowerStatsAggregator, times(2))
+                .aggregatePowerStats(anyLong(), anyLong(), any(Consumer.class));
+
+        List<PowerStatsSpan.Metadata> contents = mPowerStatsStore.getTableOfContents();
+        assertThat(contents).hasSize(3);
+        // Skip the first entry, which was placed in the store at the beginning of this test
+        PowerStatsSpan.TimeFrame timeFrame1 = contents.get(1).getTimeFrames().get(0);
+        PowerStatsSpan.TimeFrame timeFrame2 = contents.get(2).getTimeFrames().get(0);
+        assertThat(timeFrame1.startMonotonicTime).isEqualTo(7654321 + 123);
+        assertThat(timeFrame2.startMonotonicTime)
+                .isEqualTo(timeFrame1.startMonotonicTime + timeFrame1.duration);
+        assertThat(Instant.ofEpochMilli(timeFrame2.startTime))
+                .isEqualTo(Instant.parse("2023-01-02T04:00:00Z"));
+        assertThat(Duration.ofMillis(timeFrame2.duration)).isEqualTo(Duration.ofMinutes(30));
+    }
+
+    private AggregatedPowerStats createAggregatedPowerStats(long monotonicTime, long currentTime,
+            long duration) {
+        AggregatedPowerStats stats = new AggregatedPowerStats(mAggregatedPowerStatsConfig);
+        stats.addClockUpdate(monotonicTime, currentTime);
+        stats.setDuration(duration);
+        return stats;
     }
 
     @Test
@@ -81,11 +183,11 @@
                     BatteryManager.BATTERY_STATUS_DISCHARGING, 50, 0);
         }
 
-        mPowerStatsStore.reset();
+        mPowerStatsScheduler.start(/* schedulePeriodicPowerStatsCollection */false);
 
         assertThat(mPowerStatsStore.getTableOfContents()).isEmpty();
 
-        mPowerStatsScheduler.start();
+        mPowerStatsScheduler.start(true);
 
         synchronized (mBatteryStats) {
             mBatteryStats.noteFlashlightOnLocked(42, mClock.realtime, mClock.uptime);
@@ -122,7 +224,7 @@
 
         List<PowerStatsSpan.TimeFrame> timeFrames = span.getMetadata().getTimeFrames();
         assertThat(timeFrames).hasSize(1);
-        assertThat(timeFrames.get(0).startTime).isEqualTo(1234567);
+        assertThat(timeFrames.get(0).startMonotonicTime).isEqualTo(7654321);
         assertThat(timeFrames.get(0).duration).isEqualTo(120000);
 
         List<PowerStatsSpan.Section> sections = span.getSections();
@@ -136,4 +238,32 @@
                 .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),
+                123 + TimeUnit.HOURS.toMillis(2),
+                Instant.parse("2007-12-03T10:13:30.00Z").toEpochMilli())).isEqualTo(
+                123 + Duration.parse("PT1M30S").toMillis());
+
+        // Expect the aligned value to be adjusted by 2 min 45 sec - rounded to the next 15 min
+        assertThat(PowerStatsScheduler.alignToWallClock(123, TimeUnit.MINUTES.toMillis(15),
+                123 + TimeUnit.HOURS.toMillis(2),
+                Instant.parse("2007-12-03T10:57:15.00Z").toEpochMilli())).isEqualTo(
+                123 + Duration.parse("PT2M45S").toMillis());
+
+        // Expect the aligned value to be adjusted by 15 sec - rounded to the next 1 min
+        assertThat(PowerStatsScheduler.alignToWallClock(123, TimeUnit.MINUTES.toMillis(1),
+                123 + TimeUnit.HOURS.toMillis(2),
+                Instant.parse("2007-12-03T10:14:45.00Z").toEpochMilli())).isEqualTo(
+                123 + Duration.parse("PT15S").toMillis());
+
+        // Expect the aligned value to be adjusted by 1 hour 46 min 30 sec -
+        // rounded to the next 3 hours
+        assertThat(PowerStatsScheduler.alignToWallClock(123, TimeUnit.HOURS.toMillis(3),
+                123 + TimeUnit.HOURS.toMillis(9),
+                Instant.parse("2007-12-03T10:13:30.00Z").toEpochMilli())).isEqualTo(
+                123 + Duration.parse("PT1H46M30S").toMillis());
+    }
 }