Merge changes from topic "battery_usage_stats_before_reset" into sc-dev

* changes:
  Implement battery_usage_stats_before_reset pull atom
  Add BatteryUsageStatsQuery parameters: time range
  Implement BatteryUsageStatsStore
  Convert BatteryUsageStats to and from XML
  Add BatteryUsageStats.add(BatteryUsageStats)

core/java/android/os/AggregateBatteryConsumer.java

@@ -17,7 +17,13 @@
 package android.os;
 
 import android.annotation.NonNull;
+import android.util.TypedXmlPullParser;
+import android.util.TypedXmlSerializer;
 
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
 import java.io.PrintWriter;
 
 /**
@@ -72,6 +78,43 @@ public final class AggregateBatteryConsumer extends BatteryConsumer implements P
         return mConsumedPowerMah;
     }
 
+    /** Serializes this object to XML */
+    void writeToXml(TypedXmlSerializer serializer,
+            @BatteryUsageStats.AggregateBatteryConsumerScope int scope) throws IOException {
+        serializer.startTag(null, BatteryUsageStats.XML_TAG_AGGREGATE);
+        serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_SCOPE, scope);
+        serializer.attributeDouble(null, BatteryUsageStats.XML_ATTR_POWER, mConsumedPowerMah);
+        mPowerComponents.writeToXml(serializer);
+        serializer.endTag(null, BatteryUsageStats.XML_TAG_AGGREGATE);
+    }
+
+    /** Parses an XML representation and populates the BatteryUsageStats builder */
+    static void parseXml(TypedXmlPullParser parser, BatteryUsageStats.Builder builder)
+            throws XmlPullParserException, IOException {
+        final int scope = parser.getAttributeInt(null, BatteryUsageStats.XML_ATTR_SCOPE);
+        final Builder consumerBuilder = builder.getAggregateBatteryConsumerBuilder(scope);
+
+        int eventType = parser.getEventType();
+        if (eventType != XmlPullParser.START_TAG || !parser.getName().equals(
+                BatteryUsageStats.XML_TAG_AGGREGATE)) {
+            throw new XmlPullParserException("Invalid XML parser state");
+        }
+
+        consumerBuilder.setConsumedPower(
+                parser.getAttributeDouble(null, BatteryUsageStats.XML_ATTR_POWER));
+
+        while (!(eventType == XmlPullParser.END_TAG && parser.getName().equals(
+                BatteryUsageStats.XML_TAG_AGGREGATE))
+                && eventType != XmlPullParser.END_DOCUMENT) {
+            if (eventType == XmlPullParser.START_TAG) {
+                if (parser.getName().equals(BatteryUsageStats.XML_TAG_POWER_COMPONENTS)) {
+                    PowerComponents.parseXml(parser, consumerBuilder.mPowerComponentsBuilder);
+                }
+            }
+            eventType = parser.next();
+        }
+    }
+
     /**
      * Builder for DeviceBatteryConsumer.
      */
@@ -90,6 +133,14 @@ public final class AggregateBatteryConsumer extends BatteryConsumer implements P
             return this;
         }
 
+        /**
+         * Adds power and usage duration from the supplied AggregateBatteryConsumer.
+         */
+        public void add(AggregateBatteryConsumer aggregateBatteryConsumer) {
+            mConsumedPowerMah += aggregateBatteryConsumer.mConsumedPowerMah;
+            mPowerComponentsBuilder.addPowerAndDuration(aggregateBatteryConsumer.mPowerComponents);
+        }
+
         /**
          * Creates a read-only object out of the Builder values.
          */

core/java/android/os/BatteryUsageStats.java

@@ -20,16 +20,23 @@ import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.util.Range;
 import android.util.SparseArray;
+import android.util.TypedXmlPullParser;
+import android.util.TypedXmlSerializer;
 import android.util.proto.ProtoOutputStream;
 
 import com.android.internal.os.BatteryStatsHistory;
 import com.android.internal.os.BatteryStatsHistoryIterator;
 import com.android.internal.os.PowerCalculator;
 
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
 import java.io.PrintWriter;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Comparator;
 import java.util.List;
 
@@ -77,6 +84,34 @@ public final class BatteryUsageStats implements Parcelable {
 
     private static final int STATSD_PULL_ATOM_MAX_BYTES = 45000;
 
+    // XML tags and attributes for BatteryUsageStats persistence
+    static final String XML_TAG_BATTERY_USAGE_STATS = "battery_usage_stats";
+    static final String XML_TAG_AGGREGATE = "aggregate";
+    static final String XML_TAG_UID = "uid";
+    static final String XML_TAG_USER = "user";
+    static final String XML_TAG_POWER_COMPONENTS = "power_components";
+    static final String XML_TAG_COMPONENT = "component";
+    static final String XML_TAG_CUSTOM_COMPONENT = "custom_component";
+    static final String XML_ATTR_ID = "id";
+    static final String XML_ATTR_UID = "uid";
+    static final String XML_ATTR_USER_ID = "user_id";
+    static final String XML_ATTR_SCOPE = "scope";
+    static final String XML_ATTR_PREFIX_CUSTOM_COMPONENT = "custom_component_";
+    static final String XML_ATTR_START_TIMESTAMP = "start_timestamp";
+    static final String XML_ATTR_END_TIMESTAMP = "end_timestamp";
+    static final String XML_ATTR_POWER = "power";
+    static final String XML_ATTR_DURATION = "duration";
+    static final String XML_ATTR_MODEL = "model";
+    static final String XML_ATTR_BATTERY_CAPACITY = "battery_capacity";
+    static final String XML_ATTR_DISCHARGE_PERCENT = "discharge_pct";
+    static final String XML_ATTR_DISCHARGE_LOWER = "discharge_lower";
+    static final String XML_ATTR_DISCHARGE_UPPER = "discharge_upper";
+    static final String XML_ATTR_BATTERY_REMAINING = "battery_remaining";
+    static final String XML_ATTR_CHARGE_REMAINING = "charge_remaining";
+    static final String XML_ATTR_HIGHEST_DRAIN_PACKAGE = "highest_drain_package";
+    static final String XML_ATTR_TIME_IN_FOREGROUND = "time_in_foreground";
+    static final String XML_ATTR_TIME_IN_BACKGROUND = "time_in_background";
+
     private final int mDischargePercentage;
     private final double mBatteryCapacityMah;
     private final long mStatsStartTimestampMs;
@@ -96,11 +131,7 @@ public final class BatteryUsageStats implements Parcelable {
     private BatteryUsageStats(@NonNull Builder builder) {
         mStatsStartTimestampMs = builder.mStatsStartTimestampMs;
         mStatsEndTimestampMs = builder.mStatsEndTimestampMs;
-        if (builder.mStatsDurationMs != -1) {
-            mStatsDurationMs = builder.mStatsDurationMs;
-        } else {
-            mStatsDurationMs = mStatsEndTimestampMs - mStatsStartTimestampMs;
-        }
+        mStatsDurationMs = builder.getStatsDuration();
         mBatteryCapacityMah = builder.mBatteryCapacityMah;
         mDischargePercentage = builder.mDischargePercentage;
         mDischargedPowerLowerBound = builder.mDischargedPowerLowerBoundMah;
@@ -579,6 +610,109 @@ public final class BatteryUsageStats implements Parcelable {
         }
     }
 
+    /** Serializes this object to XML */
+    public void writeXml(TypedXmlSerializer serializer) throws IOException {
+        serializer.startTag(null, XML_TAG_BATTERY_USAGE_STATS);
+
+        for (int i = 0; i < mCustomPowerComponentNames.length; i++) {
+            serializer.attribute(null, XML_ATTR_PREFIX_CUSTOM_COMPONENT + i,
+                    mCustomPowerComponentNames[i]);
+        }
+
+        serializer.attributeLong(null, XML_ATTR_START_TIMESTAMP, mStatsStartTimestampMs);
+        serializer.attributeLong(null, XML_ATTR_END_TIMESTAMP, mStatsEndTimestampMs);
+        serializer.attributeLong(null, XML_ATTR_DURATION, mStatsDurationMs);
+        serializer.attributeDouble(null, XML_ATTR_BATTERY_CAPACITY, mBatteryCapacityMah);
+        serializer.attributeInt(null, XML_ATTR_DISCHARGE_PERCENT, mDischargePercentage);
+        serializer.attributeDouble(null, XML_ATTR_DISCHARGE_LOWER, mDischargedPowerLowerBound);
+        serializer.attributeDouble(null, XML_ATTR_DISCHARGE_UPPER, mDischargedPowerUpperBound);
+        serializer.attributeLong(null, XML_ATTR_BATTERY_REMAINING, mBatteryTimeRemainingMs);
+        serializer.attributeLong(null, XML_ATTR_CHARGE_REMAINING, mChargeTimeRemainingMs);
+
+        for (int scope = 0; scope < BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_COUNT;
+                scope++) {
+            mAggregateBatteryConsumers[scope].writeToXml(serializer, scope);
+        }
+        for (UidBatteryConsumer consumer : mUidBatteryConsumers) {
+            consumer.writeToXml(serializer);
+        }
+        for (UserBatteryConsumer consumer : mUserBatteryConsumers) {
+            consumer.writeToXml(serializer);
+        }
+        serializer.endTag(null, XML_TAG_BATTERY_USAGE_STATS);
+    }
+
+    /** Parses an XML representation of BatteryUsageStats */
+    public static BatteryUsageStats createFromXml(TypedXmlPullParser parser)
+            throws XmlPullParserException, IOException {
+        Builder builder = null;
+        int eventType = parser.getEventType();
+        while (eventType != XmlPullParser.END_DOCUMENT) {
+            if (eventType == XmlPullParser.START_TAG
+                    && parser.getName().equals(XML_TAG_BATTERY_USAGE_STATS)) {
+                List<String> customComponentNames = new ArrayList<>();
+                int i = 0;
+                while (true) {
+                    int index = parser.getAttributeIndex(null,
+                            XML_ATTR_PREFIX_CUSTOM_COMPONENT + i);
+                    if (index == -1) {
+                        break;
+                    }
+                    customComponentNames.add(parser.getAttributeValue(index));
+                    i++;
+                }
+
+                builder = new Builder(
+                        customComponentNames.toArray(new String[0]), true);
+
+                builder.setStatsStartTimestamp(
+                        parser.getAttributeLong(null, XML_ATTR_START_TIMESTAMP));
+                builder.setStatsEndTimestamp(
+                        parser.getAttributeLong(null, XML_ATTR_END_TIMESTAMP));
+                builder.setStatsDuration(
+                        parser.getAttributeLong(null, XML_ATTR_DURATION));
+                builder.setBatteryCapacity(
+                        parser.getAttributeDouble(null, XML_ATTR_BATTERY_CAPACITY));
+                builder.setDischargePercentage(
+                        parser.getAttributeInt(null, XML_ATTR_DISCHARGE_PERCENT));
+                builder.setDischargedPowerRange(
+                        parser.getAttributeDouble(null, XML_ATTR_DISCHARGE_LOWER),
+                        parser.getAttributeDouble(null, XML_ATTR_DISCHARGE_UPPER));
+                builder.setBatteryTimeRemainingMs(
+                        parser.getAttributeLong(null, XML_ATTR_BATTERY_REMAINING));
+                builder.setChargeTimeRemainingMs(
+                        parser.getAttributeLong(null, XML_ATTR_CHARGE_REMAINING));
+
+                eventType = parser.next();
+                break;
+            }
+            eventType = parser.next();
+        }
+
+        if (builder == null) {
+            throw new XmlPullParserException("No root element");
+        }
+
+        while (eventType != XmlPullParser.END_DOCUMENT) {
+            if (eventType == XmlPullParser.START_TAG) {
+                switch (parser.getName()) {
+                    case XML_TAG_AGGREGATE:
+                        AggregateBatteryConsumer.parseXml(parser, builder);
+                        break;
+                    case XML_TAG_UID:
+                        UidBatteryConsumer.createFromXml(parser, builder);
+                        break;
+                    case XML_TAG_USER:
+                        UserBatteryConsumer.createFromXml(parser, builder);
+                        break;
+                }
+            }
+            eventType = parser.next();
+        }
+
+        return builder.build();
+    }
+
     /**
      * Builder for BatteryUsageStats.
      */
@@ -658,6 +792,14 @@ public final class BatteryUsageStats implements Parcelable {
             return this;
         }
 
+        private long getStatsDuration() {
+            if (mStatsDurationMs != -1) {
+                return mStatsDurationMs;
+            } else {
+                return mStatsEndTimestampMs - mStatsStartTimestampMs;
+            }
+        }
+
         /**
          * Sets the battery discharge amount since BatteryStats reset as percentage of the full
          * charge.
@@ -737,6 +879,22 @@ public final class BatteryUsageStats implements Parcelable {
             return builder;
         }
 
+        /**
+         * Creates or returns a UidBatteryConsumer, which represents battery attribution
+         * data for an individual UID. This version of the method is not suitable for use
+         * with PowerCalculators.
+         */
+        @NonNull
+        public UidBatteryConsumer.Builder getOrCreateUidBatteryConsumerBuilder(int uid) {
+            UidBatteryConsumer.Builder builder = mUidBatteryConsumerBuilders.get(uid);
+            if (builder == null) {
+                builder = new UidBatteryConsumer.Builder(mCustomPowerComponentNames,
+                        mIncludePowerModels, uid);
+                mUidBatteryConsumerBuilders.put(uid, builder);
+            }
+            return builder;
+        }
+
         /**
          * Creates or returns a UserBatteryConsumer, which represents battery attribution
          * data for an individual {@link UserHandle}.
@@ -756,5 +914,59 @@ public final class BatteryUsageStats implements Parcelable {
         public SparseArray<UidBatteryConsumer.Builder> getUidBatteryConsumerBuilders() {
             return mUidBatteryConsumerBuilders;
         }
+
+        /**
+         * Adds battery usage stats from another snapshots. The two snapshots are assumed to be
+         * non-overlapping, meaning that the power consumption estimates and session durations
+         * can be simply summed across the two snapshots.  This remains true even if the timestamps
+         * seem to indicate that the sessions are in fact overlapping: timestamps may be off as a
+         * result of realtime clock adjustments by the user or the system.
+         */
+        @NonNull
+        public Builder add(BatteryUsageStats stats) {
+            if (!Arrays.equals(mCustomPowerComponentNames, stats.mCustomPowerComponentNames)) {
+                throw new IllegalArgumentException(
+                        "BatteryUsageStats have different custom power components");
+            }
+
+            if (mUserBatteryConsumerBuilders.size() != 0
+                    || !stats.getUserBatteryConsumers().isEmpty()) {
+                throw new UnsupportedOperationException(
+                        "Combining UserBatteryConsumers is not supported");
+            }
+
+            mDischargedPowerLowerBoundMah += stats.mDischargedPowerLowerBound;
+            mDischargedPowerUpperBoundMah += stats.mDischargedPowerUpperBound;
+            mDischargePercentage += stats.mDischargePercentage;
+
+            mStatsDurationMs = getStatsDuration() + stats.getStatsDuration();
+
+            if (mStatsStartTimestampMs == 0
+                    || stats.mStatsStartTimestampMs < mStatsStartTimestampMs) {
+                mStatsStartTimestampMs = stats.mStatsStartTimestampMs;
+            }
+
+            final boolean addingLaterSnapshot = stats.mStatsEndTimestampMs > mStatsEndTimestampMs;
+            if (addingLaterSnapshot) {
+                mStatsEndTimestampMs = stats.mStatsEndTimestampMs;
+            }
+
+            for (int scope = 0; scope < AGGREGATE_BATTERY_CONSUMER_SCOPE_COUNT; scope++) {
+                getAggregateBatteryConsumerBuilder(scope)
+                        .add(stats.mAggregateBatteryConsumers[scope]);
+            }
+
+            for (UidBatteryConsumer consumer : stats.getUidBatteryConsumers()) {
+                getOrCreateUidBatteryConsumerBuilder(consumer.getUid()).add(consumer);
+            }
+
+            if (addingLaterSnapshot) {
+                mBatteryCapacityMah = stats.mBatteryCapacityMah;
+                mBatteryTimeRemainingMs = stats.mBatteryTimeRemainingMs;
+                mChargeTimeRemainingMs = stats.mChargeTimeRemainingMs;
+            }
+
+            return this;
+        }
     }
 }

core/java/android/os/BatteryUsageStatsQuery.java

@@ -72,12 +72,16 @@ public final class BatteryUsageStatsQuery implements Parcelable {
     @NonNull
     private final int[] mUserIds;
     private final long mMaxStatsAgeMs;
+    private long mFromTimestamp;
+    private long mToTimestamp;
 
     private BatteryUsageStatsQuery(@NonNull Builder builder) {
         mFlags = builder.mFlags;
         mUserIds = builder.mUserIds != null ? builder.mUserIds.toArray()
                 : new int[]{UserHandle.USER_ALL};
         mMaxStatsAgeMs = builder.mMaxStatsAgeMs;
+        mFromTimestamp = builder.mFromTimestamp;
+        mToTimestamp = builder.mToTimestamp;
     }
 
     @BatteryUsageStatsFlags
@@ -112,11 +116,30 @@ public final class BatteryUsageStatsQuery implements Parcelable {
         return mMaxStatsAgeMs;
     }
 
+    /**
+     * Returns the exclusive lower bound of the stored snapshot timestamps that should be included
+     * in the aggregation.  Ignored if {@link #getToTimestamp()} is zero.
+     */
+    public long getFromTimestamp() {
+        return mFromTimestamp;
+    }
+
+    /**
+     * Returns the inclusive upper bound of the stored snapshot timestamps that should
+     * be included in the aggregation.  The default is to include only the current stats
+     * accumulated since the latest battery reset.
+     */
+    public long getToTimestamp() {
+        return mToTimestamp;
+    }
+
     private BatteryUsageStatsQuery(Parcel in) {
         mFlags = in.readInt();
         mUserIds = new int[in.readInt()];
         in.readIntArray(mUserIds);
         mMaxStatsAgeMs = in.readLong();
+        mFromTimestamp = in.readLong();
+        mToTimestamp = in.readLong();
     }
 
     @Override
@@ -125,6 +148,8 @@ public final class BatteryUsageStatsQuery implements Parcelable {
         dest.writeInt(mUserIds.length);
         dest.writeIntArray(mUserIds);
         dest.writeLong(mMaxStatsAgeMs);
+        dest.writeLong(mFromTimestamp);
+        dest.writeLong(mToTimestamp);
     }
 
     @Override
@@ -153,6 +178,8 @@ public final class BatteryUsageStatsQuery implements Parcelable {
         private int mFlags;
         private IntArray mUserIds;
         private long mMaxStatsAgeMs = DEFAULT_MAX_STATS_AGE_MS;
+        private long mFromTimestamp;
+        private long mToTimestamp;
 
         /**
          * Builds a read-only BatteryUsageStatsQuery object.
@@ -203,6 +230,17 @@ public final class BatteryUsageStatsQuery implements Parcelable {
             return this;
         }
 
+        /**
+         * Requests to aggregate stored snapshots between the two supplied timestamps
+         * @param fromTimestamp Exclusive starting timestamp, as per System.currentTimeMillis()
+         * @param toTimestamp Inclusive ending timestamp, as per System.currentTimeMillis()
+         */
+        public Builder aggregateSnapshots(long fromTimestamp, long toTimestamp) {
+            mFromTimestamp = fromTimestamp;
+            mToTimestamp = toTimestamp;
+            return this;
+        }
+
         /**
          * Set the client's tolerance for stale battery stats. The data may be up to
          * this many milliseconds out-of-date.

core/java/android/os/PowerComponents.java

@@ -19,11 +19,18 @@ import static android.os.BatteryConsumer.convertMahToDeciCoulombs;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.util.TypedXmlPullParser;
+import android.util.TypedXmlSerializer;
 import android.util.proto.ProtoOutputStream;
 
 import com.android.internal.os.PowerCalculator;
 
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
 import java.io.PrintWriter;
+import java.util.Arrays;
 
 /**
  * Contains details of battery attribution data broken down to individual power drain types
@@ -36,9 +43,12 @@ class PowerComponents {
             - BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID;
 
     private final double mConsumedPowerMah;
+    @NonNull
     private final double[] mPowerComponentsMah;
+    @NonNull
     private final long[] mUsageDurationsMs;
     private final int mCustomPowerComponentCount;
+    @Nullable
     private final byte[] mPowerModels;
     // Not written to Parcel and must be explicitly restored during the parent object's unparceling
     private String[] mCustomPowerComponentNames;
@@ -49,7 +59,7 @@ class PowerComponents {
         mPowerComponentsMah = builder.mPowerComponentsMah;
         mUsageDurationsMs = builder.mUsageDurationsMs;
         mConsumedPowerMah = builder.getTotalPower();
-        mPowerModels = builder.mPowerModels;
+        mPowerModels = builder.getPowerModels();
     }
 
     PowerComponents(@NonNull Parcel source) {
@@ -146,9 +156,13 @@ class PowerComponents {
         }
     }
 
+    public boolean hasPowerModels() {
+        return mPowerModels != null;
+    }
+
     @BatteryConsumer.PowerModel
     int getPowerModel(@BatteryConsumer.PowerComponent int component) {
-        if (mPowerModels == null) {
+        if (!hasPowerModels()) {
             throw new IllegalStateException(
                     "Power model IDs were not requested in the BatteryUsageStatsQuery");
         }
@@ -294,10 +308,128 @@ class PowerComponents {
         return interestingData;
     }
 
+    void writeToXml(TypedXmlSerializer serializer) throws IOException {
+        serializer.startTag(null, BatteryUsageStats.XML_TAG_POWER_COMPONENTS);
+        for (int componentId = 0; componentId < BatteryConsumer.POWER_COMPONENT_COUNT;
+                componentId++) {
+            final double powerMah = getConsumedPower(componentId);
+            final long durationMs = getUsageDurationMillis(componentId);
+            if (powerMah == 0 && durationMs == 0) {
+                continue;
+            }
+
+            serializer.startTag(null, BatteryUsageStats.XML_TAG_COMPONENT);
+            serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_ID, componentId);
+            if (powerMah != 0) {
+                serializer.attributeDouble(null, BatteryUsageStats.XML_ATTR_POWER, powerMah);
+            }
+            if (durationMs != 0) {
+                serializer.attributeLong(null, BatteryUsageStats.XML_ATTR_DURATION, durationMs);
+            }
+            if (mPowerModels != null) {
+                serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_MODEL,
+                        mPowerModels[componentId]);
+            }
+            serializer.endTag(null, BatteryUsageStats.XML_TAG_COMPONENT);
+        }
+
+        final int customComponentEnd =
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID + mCustomPowerComponentCount;
+        for (int componentId = BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID;
+                componentId < customComponentEnd;
+                componentId++) {
+            final double powerMah = getConsumedPowerForCustomComponent(componentId);
+            final long durationMs = getUsageDurationForCustomComponentMillis(componentId);
+            if (powerMah == 0 && durationMs == 0) {
+                continue;
+            }
+
+            serializer.startTag(null, BatteryUsageStats.XML_TAG_CUSTOM_COMPONENT);
+            serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_ID, componentId);
+            if (powerMah != 0) {
+                serializer.attributeDouble(null, BatteryUsageStats.XML_ATTR_POWER, powerMah);
+            }
+            if (durationMs != 0) {
+                serializer.attributeLong(null, BatteryUsageStats.XML_ATTR_DURATION, durationMs);
+            }
+            serializer.endTag(null, BatteryUsageStats.XML_TAG_CUSTOM_COMPONENT);
+        }
+
+        serializer.endTag(null, BatteryUsageStats.XML_TAG_POWER_COMPONENTS);
+    }
+
+
+    static void parseXml(TypedXmlPullParser parser, PowerComponents.Builder builder)
+            throws XmlPullParserException, IOException {
+        int eventType = parser.getEventType();
+        if (eventType != XmlPullParser.START_TAG || !parser.getName().equals(
+                BatteryUsageStats.XML_TAG_POWER_COMPONENTS)) {
+            throw new XmlPullParserException("Invalid XML parser state");
+        }
+
+        while (!(eventType == XmlPullParser.END_TAG && parser.getName().equals(
+                BatteryUsageStats.XML_TAG_POWER_COMPONENTS))
+                && eventType != XmlPullParser.END_DOCUMENT) {
+            if (eventType == XmlPullParser.START_TAG) {
+                switch (parser.getName()) {
+                    case BatteryUsageStats.XML_TAG_COMPONENT: {
+                        int componentId = -1;
+                        double powerMah = 0;
+                        long durationMs = 0;
+                        int model = BatteryConsumer.POWER_MODEL_UNDEFINED;
+                        for (int i = 0; i < parser.getAttributeCount(); i++) {
+                            switch (parser.getAttributeName(i)) {
+                                case BatteryUsageStats.XML_ATTR_ID:
+                                    componentId = parser.getAttributeInt(i);
+                                    break;
+                                case BatteryUsageStats.XML_ATTR_POWER:
+                                    powerMah = parser.getAttributeDouble(i);
+                                    break;
+                                case BatteryUsageStats.XML_ATTR_DURATION:
+                                    durationMs = parser.getAttributeLong(i);
+                                    break;
+                                case BatteryUsageStats.XML_ATTR_MODEL:
+                                    model = parser.getAttributeInt(i);
+                                    break;
+                            }
+                        }
+                        builder.setConsumedPower(componentId, powerMah, model);
+                        builder.setUsageDurationMillis(componentId, durationMs);
+                        break;
+                    }
+                    case BatteryUsageStats.XML_TAG_CUSTOM_COMPONENT: {
+                        int componentId = -1;
+                        double powerMah = 0;
+                        long durationMs = 0;
+                        for (int i = 0; i < parser.getAttributeCount(); i++) {
+                            switch (parser.getAttributeName(i)) {
+                                case BatteryUsageStats.XML_ATTR_ID:
+                                    componentId = parser.getAttributeInt(i);
+                                    break;
+                                case BatteryUsageStats.XML_ATTR_POWER:
+                                    powerMah = parser.getAttributeDouble(i);
+                                    break;
+                                case BatteryUsageStats.XML_ATTR_DURATION:
+                                    durationMs = parser.getAttributeLong(i);
+                                    break;
+                            }
+                        }
+                        builder.setConsumedPowerForCustomComponent(componentId, powerMah);
+                        builder.setUsageDurationForCustomComponentMillis(componentId, durationMs);
+                        break;
+                    }
+                }
+            }
+            eventType = parser.next();
+        }
+    }
+
     /**
      * Builder for PowerComponents.
      */
     static final class Builder {
+        private static final byte POWER_MODEL_UNINITIALIZED = -1;
+
         private final double[] mPowerComponentsMah;
         private final String[] mCustomPowerComponentNames;
         private final long[] mUsageDurationsMs;
@@ -311,6 +443,7 @@ class PowerComponents {
             mUsageDurationsMs = new long[powerComponentCount];
             if (includePowerModels) {
                 mPowerModels = new byte[BatteryConsumer.POWER_COMPONENT_COUNT];
+                Arrays.fill(mPowerModels, POWER_MODEL_UNINITIALIZED);
             } else {
                 mPowerModels = null;
             }
@@ -412,12 +545,39 @@ class PowerComponents {
             return this;
         }
 
-        public void addPowerAndDuration(Builder other) {
+        public void addPowerAndDuration(PowerComponents.Builder other) {
+            addPowerAndDuration(other.mPowerComponentsMah, other.mUsageDurationsMs,
+                    other.mPowerModels);
+        }
+
+        public void addPowerAndDuration(PowerComponents other) {
+            addPowerAndDuration(other.mPowerComponentsMah, other.mUsageDurationsMs,
+                    other.mPowerModels);
+        }
+
+        private void addPowerAndDuration(double[] powerComponentsMah,
+                long[] usageDurationsMs, byte[] powerModels) {
+            if (mPowerComponentsMah.length != powerComponentsMah.length) {
+                throw new IllegalArgumentException(
+                        "Number of power components does not match: " + powerComponentsMah.length
+                                + ", expected: " + mPowerComponentsMah.length);
+            }
+
             for (int i = mPowerComponentsMah.length - 1; i >= 0; i--) {
-                mPowerComponentsMah[i] += other.mPowerComponentsMah[i];
+                mPowerComponentsMah[i] += powerComponentsMah[i];
             }
             for (int i = mUsageDurationsMs.length - 1; i >= 0; i--) {
-                mUsageDurationsMs[i] += other.mUsageDurationsMs[i];
+                mUsageDurationsMs[i] += usageDurationsMs[i];
+            }
+            if (mPowerModels != null && powerModels != null) {
+                for (int i = mPowerModels.length - 1; i >= 0; i--) {
+                    if (mPowerModels[i] == POWER_MODEL_UNINITIALIZED) {
+                        mPowerModels[i] = powerModels[i];
+                    } else if (mPowerModels[i] != powerModels[i]
+                            && powerModels[i] != POWER_MODEL_UNINITIALIZED) {
+                        mPowerModels[i] = BatteryConsumer.POWER_MODEL_UNDEFINED;
+                    }
+                }
             }
         }
 
@@ -433,6 +593,19 @@ class PowerComponents {
             return totalPowerMah;
         }
 
+        private byte[] getPowerModels() {
+            if (mPowerModels == null) {
+                return null;
+            }
+
+            byte[] powerModels = new byte[mPowerModels.length];
+            for (int i = mPowerModels.length - 1; i >= 0; i--) {
+                powerModels[i] = mPowerModels[i] != POWER_MODEL_UNINITIALIZED ? mPowerModels[i]
+                        : BatteryConsumer.POWER_MODEL_UNDEFINED;
+            }
+            return powerModels;
+        }
+
         /**
          * Creates a read-only object out of the Builder values.
          */

core/java/android/os/UidBatteryConsumer.java

@@ -19,9 +19,16 @@ package android.os;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.text.TextUtils;
+import android.util.TypedXmlPullParser;
+import android.util.TypedXmlSerializer;
 
 import com.android.internal.os.PowerCalculator;
 
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
 import java.io.PrintWriter;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -143,13 +150,65 @@ public final class UidBatteryConsumer extends BatteryConsumer implements Parcela
         return 0;
     }
 
+    /** Serializes this object to XML */
+    void writeToXml(TypedXmlSerializer serializer) throws IOException {
+        if (getConsumedPower() == 0) {
+            return;
+        }
+
+        serializer.startTag(null, BatteryUsageStats.XML_TAG_UID);
+        serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_UID, getUid());
+        if (!TextUtils.isEmpty(mPackageWithHighestDrain)) {
+            serializer.attribute(null, BatteryUsageStats.XML_ATTR_HIGHEST_DRAIN_PACKAGE,
+                    mPackageWithHighestDrain);
+        }
+        serializer.attributeLong(null, BatteryUsageStats.XML_ATTR_TIME_IN_FOREGROUND,
+                mTimeInForegroundMs);
+        serializer.attributeLong(null, BatteryUsageStats.XML_ATTR_TIME_IN_BACKGROUND,
+                mTimeInBackgroundMs);
+        mPowerComponents.writeToXml(serializer);
+        serializer.endTag(null, BatteryUsageStats.XML_TAG_UID);
+    }
+
+    /** Parses an XML representation and populates the BatteryUsageStats builder */
+    static void createFromXml(TypedXmlPullParser parser, BatteryUsageStats.Builder builder)
+            throws XmlPullParserException, IOException {
+        final int uid = parser.getAttributeInt(null, BatteryUsageStats.XML_ATTR_UID);
+        final UidBatteryConsumer.Builder consumerBuilder =
+                builder.getOrCreateUidBatteryConsumerBuilder(uid);
+
+        int eventType = parser.getEventType();
+        if (eventType != XmlPullParser.START_TAG
+                || !parser.getName().equals(BatteryUsageStats.XML_TAG_UID)) {
+            throw new XmlPullParserException("Invalid XML parser state");
+        }
+
+        consumerBuilder.setPackageWithHighestDrain(
+                parser.getAttributeValue(null, BatteryUsageStats.XML_ATTR_HIGHEST_DRAIN_PACKAGE));
+        consumerBuilder.setTimeInStateMs(STATE_FOREGROUND,
+                parser.getAttributeLong(null, BatteryUsageStats.XML_ATTR_TIME_IN_FOREGROUND));
+        consumerBuilder.setTimeInStateMs(STATE_BACKGROUND,
+                parser.getAttributeLong(null, BatteryUsageStats.XML_ATTR_TIME_IN_BACKGROUND));
+        while (!(eventType == XmlPullParser.END_TAG
+                && parser.getName().equals(BatteryUsageStats.XML_TAG_UID))
+                && eventType != XmlPullParser.END_DOCUMENT) {
+            if (eventType == XmlPullParser.START_TAG) {
+                if (parser.getName().equals(BatteryUsageStats.XML_TAG_POWER_COMPONENTS)) {
+                    PowerComponents.parseXml(parser, consumerBuilder.mPowerComponentsBuilder);
+                }
+            }
+            eventType = parser.next();
+        }
+    }
+
     /**
      * Builder for UidBatteryConsumer.
      */
     public static final class Builder extends BaseBuilder<Builder> {
+        private static final String PACKAGE_NAME_UNINITIALIZED = "";
         private final BatteryStats.Uid mBatteryStatsUid;
         private final int mUid;
-        private String mPackageWithHighestDrain;
+        private String mPackageWithHighestDrain = PACKAGE_NAME_UNINITIALIZED;
         public long mTimeInForegroundMs;
         public long mTimeInBackgroundMs;
         private boolean mExcludeFromBatteryUsageStats;
@@ -161,8 +220,19 @@ public final class UidBatteryConsumer extends BatteryConsumer implements Parcela
             mUid = batteryStatsUid.getUid();
         }
 
+        public Builder(@NonNull String[] customPowerComponentNames, boolean includePowerModels,
+                int uid) {
+            super(customPowerComponentNames, includePowerModels);
+            mBatteryStatsUid = null;
+            mUid = uid;
+        }
+
         @NonNull
         public BatteryStats.Uid getBatteryStatsUid() {
+            if (mBatteryStatsUid == null) {
+                throw new IllegalStateException(
+                        "UidBatteryConsumer.Builder was initialized without a BatteryStats.Uid");
+            }
             return mBatteryStatsUid;
         }
 
@@ -176,7 +246,7 @@ public final class UidBatteryConsumer extends BatteryConsumer implements Parcela
          */
         @NonNull
         public Builder setPackageWithHighestDrain(@Nullable String packageName) {
-            mPackageWithHighestDrain = packageName;
+            mPackageWithHighestDrain = TextUtils.nullIfEmpty(packageName);
             return this;
         }
 
@@ -207,6 +277,30 @@ public final class UidBatteryConsumer extends BatteryConsumer implements Parcela
             return this;
         }
 
+        /**
+         * Adds power and usage duration from the supplied UidBatteryConsumer.
+         */
+        public Builder add(UidBatteryConsumer consumer) {
+            mPowerComponentsBuilder.addPowerAndDuration(consumer.mPowerComponents);
+            mTimeInBackgroundMs += consumer.mTimeInBackgroundMs;
+            mTimeInForegroundMs += consumer.mTimeInForegroundMs;
+
+            if (mPackageWithHighestDrain == PACKAGE_NAME_UNINITIALIZED) {
+                mPackageWithHighestDrain = consumer.mPackageWithHighestDrain;
+            } else if (!TextUtils.equals(mPackageWithHighestDrain,
+                    consumer.mPackageWithHighestDrain)) {
+                // Consider combining two UidBatteryConsumers with this distribution
+                // of power drain between packages:
+                // (package1=100, package2=10) and (package1=100, package2=101).
+                // Since we don't know the actual power distribution between packages at this
+                // point, we have no way to correctly declare package1 as the winner.
+                // The naive logic of picking the consumer with the higher total consumed
+                // power would produce an incorrect result.
+                mPackageWithHighestDrain = null;
+            }
+            return this;
+        }
+
         /**
          * Returns true if this UidBatteryConsumer must be excluded from the
          * BatteryUsageStats.
@@ -220,6 +314,9 @@ public final class UidBatteryConsumer extends BatteryConsumer implements Parcela
          */
         @NonNull
         public UidBatteryConsumer build() {
+            if (mPackageWithHighestDrain == PACKAGE_NAME_UNINITIALIZED) {
+                mPackageWithHighestDrain = null;
+            }
             return new UidBatteryConsumer(this);
         }
     }

core/java/android/os/UserBatteryConsumer.java

@@ -17,9 +17,15 @@
 package android.os;
 
 import android.annotation.NonNull;
+import android.util.TypedXmlPullParser;
+import android.util.TypedXmlSerializer;
 
 import com.android.internal.os.PowerCalculator;
 
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.List;
@@ -85,6 +91,42 @@ public class UserBatteryConsumer extends BatteryConsumer implements Parcelable {
         return 0;
     }
 
+    /** Serializes this object to XML */
+    void writeToXml(TypedXmlSerializer serializer) throws IOException {
+        if (getConsumedPower() == 0) {
+            return;
+        }
+
+        serializer.startTag(null, BatteryUsageStats.XML_TAG_USER);
+        serializer.attributeInt(null, BatteryUsageStats.XML_ATTR_USER_ID, getUserId());
+        mPowerComponents.writeToXml(serializer);
+        serializer.endTag(null, BatteryUsageStats.XML_TAG_USER);
+    }
+
+    /** Parses an XML representation and populates the BatteryUsageStats builder */
+    static void createFromXml(TypedXmlPullParser parser, BatteryUsageStats.Builder builder)
+            throws XmlPullParserException, IOException {
+        final int userId = parser.getAttributeInt(null, BatteryUsageStats.XML_ATTR_USER_ID);
+        final UserBatteryConsumer.Builder consumerBuilder =
+                builder.getOrCreateUserBatteryConsumerBuilder(userId);
+
+        int eventType = parser.getEventType();
+        if (eventType != XmlPullParser.START_TAG
+                || !parser.getName().equals(BatteryUsageStats.XML_TAG_USER)) {
+            throw new XmlPullParserException("Invalid XML parser state");
+        }
+        while (!(eventType == XmlPullParser.END_TAG
+                && parser.getName().equals(BatteryUsageStats.XML_TAG_USER))
+                && eventType != XmlPullParser.END_DOCUMENT) {
+            if (eventType == XmlPullParser.START_TAG) {
+                if (parser.getName().equals(BatteryUsageStats.XML_TAG_POWER_COMPONENTS)) {
+                    PowerComponents.parseXml(parser, consumerBuilder.mPowerComponentsBuilder);
+                }
+            }
+            eventType = parser.next();
+        }
+    }
+
     /**
      * Builder for UserBatteryConsumer.
      */

core/java/com/android/internal/os/BatteryStatsImpl.java

@@ -341,6 +341,19 @@ public class BatteryStatsImpl extends BatteryStats {
         }
     }
 
+    /**
+     * Listener for the battery stats reset.
+     */
+    public interface BatteryResetListener {
+
+        /**
+         * Callback invoked immediately prior to resetting battery stats.
+         */
+        void prepareForBatteryStatsReset();
+    }
+
+    private BatteryResetListener mBatteryResetListener;
+
     public interface BatteryCallback {
         public void batteryNeedsCpuUpdate();
         public void batteryPowerChanged(boolean onBattery);
@@ -10736,6 +10749,10 @@ public class BatteryStatsImpl extends BatteryStats {
         }
     }
 
+    PowerProfile getPowerProfile() {
+        return mPowerProfile;
+    }
+
     /**
      * Starts tracking CPU time-in-state for threads of the system server process,
      * keeping a separate account of threads receiving incoming binder calls.
@@ -11184,6 +11201,10 @@ public class BatteryStatsImpl extends BatteryStats {
         mDischargeCounter.reset(false, elapsedRealtimeUs);
     }
 
+    public void setBatteryResetListener(BatteryResetListener batteryResetListener) {
+        mBatteryResetListener = batteryResetListener;
+    }
+
     public void resetAllStatsCmdLocked() {
         final long mSecUptime = mClocks.uptimeMillis();
         long uptimeUs = mSecUptime * 1000;
@@ -11219,6 +11240,10 @@ public class BatteryStatsImpl extends BatteryStats {
     }
 
     private void resetAllStatsLocked(long uptimeMillis, long elapsedRealtimeMillis) {
+        if (mBatteryResetListener != null) {
+            mBatteryResetListener.prepareForBatteryStatsReset();
+        }
+
         final long uptimeUs = uptimeMillis * 1000;
         final long elapsedRealtimeUs = elapsedRealtimeMillis * 1000;
         mStartCount = 0;

core/java/com/android/internal/os/BatteryUsageStatsProvider.java

@@ -38,14 +38,24 @@ import java.util.Map;
 public class BatteryUsageStatsProvider {
     private final Context mContext;
     private final BatteryStats mStats;
+    private final BatteryUsageStatsStore mBatteryUsageStatsStore;
     private final PowerProfile mPowerProfile;
     private final Object mLock = new Object();
     private List<PowerCalculator> mPowerCalculators;
 
     public BatteryUsageStatsProvider(Context context, BatteryStats stats) {
+        this(context, stats, null);
+    }
+
+    @VisibleForTesting
+    public BatteryUsageStatsProvider(Context context, BatteryStats stats,
+            BatteryUsageStatsStore batteryUsageStatsStore) {
         mContext = context;
         mStats = stats;
-        mPowerProfile = new PowerProfile(mContext);
+        mBatteryUsageStatsStore = batteryUsageStatsStore;
+        mPowerProfile = stats instanceof BatteryStatsImpl
+                ? ((BatteryStatsImpl) stats).getPowerProfile()
+                : new PowerProfile(context);
     }
 
     private List<PowerCalculator> getPowerCalculators() {
@@ -126,6 +136,15 @@ public class BatteryUsageStatsProvider {
 
     private BatteryUsageStats getBatteryUsageStats(BatteryUsageStatsQuery query,
             long currentTimeMs) {
+        if (query.getToTimestamp() == 0) {
+            return getCurrentBatteryUsageStats(query, currentTimeMs);
+        } else {
+            return getAggregatedBatteryUsageStats(query);
+        }
+    }
+
+    private BatteryUsageStats getCurrentBatteryUsageStats(BatteryUsageStatsQuery query,
+            long currentTimeMs) {
         final long realtimeUs = elapsedRealtime() * 1000;
         final long uptimeUs = uptimeMillis() * 1000;
 
@@ -209,6 +228,25 @@ public class BatteryUsageStatsProvider {
                 BatteryStats.STATS_SINCE_CHARGED) / 1000;
     }
 
+    private BatteryUsageStats getAggregatedBatteryUsageStats(BatteryUsageStatsQuery query) {
+        final boolean includePowerModels = (query.getFlags()
+                & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_POWER_MODELS) != 0;
+
+        final BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(
+                mStats.getCustomEnergyConsumerNames(), includePowerModels);
+        final long[] timestamps = mBatteryUsageStatsStore.listBatteryUsageStatsTimestamps();
+        for (long timestamp : timestamps) {
+            if (timestamp > query.getFromTimestamp() && timestamp <= query.getToTimestamp()) {
+                final BatteryUsageStats snapshot =
+                        mBatteryUsageStatsStore.loadBatteryUsageStats(timestamp);
+                if (snapshot != null) {
+                    builder.add(snapshot);
+                }
+            }
+        }
+        return builder.build();
+    }
+
     private long elapsedRealtime() {
         if (mStats instanceof BatteryStatsImpl) {
             return ((BatteryStatsImpl) mStats).mClocks.elapsedRealtime();

core/java/com/android/internal/os/BatteryUsageStatsStore.java

@@ -0,0 +1,285 @@
+/*
+ * Copyright (C) 2021 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.internal.os;
+
+import android.annotation.Nullable;
+import android.content.Context;
+import android.os.BatteryUsageStats;
+import android.os.BatteryUsageStatsQuery;
+import android.os.Handler;
+import android.util.AtomicFile;
+import android.util.LongArray;
+import android.util.Slog;
+import android.util.TypedXmlPullParser;
+import android.util.TypedXmlSerializer;
+import android.util.Xml;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.channels.FileChannel;
+import java.nio.channels.FileLock;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.StandardOpenOption;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Properties;
+import java.util.TreeMap;
+
+/**
+ * A storage mechanism for BatteryUsageStats snapshots.
+ */
+public class BatteryUsageStatsStore {
+    private static final String TAG = "BatteryUsageStatsStore";
+
+    private static final List<BatteryUsageStatsQuery> BATTERY_USAGE_STATS_QUERY = List.of(
+            new BatteryUsageStatsQuery.Builder()
+                    .setMaxStatsAgeMs(0)
+                    .includePowerModels()
+                    .build());
+    private static final String BATTERY_USAGE_STATS_DIR = "battery-usage-stats";
+    private static final String SNAPSHOT_FILE_EXTENSION = ".bus";
+    private static final String DIR_LOCK_FILENAME = ".lock";
+    private static final String CONFIG_FILENAME = "config";
+    private static final String BATTERY_USAGE_STATS_BEFORE_RESET_TIMESTAMP_PROPERTY =
+            "BATTERY_USAGE_STATS_BEFORE_RESET_TIMESTAMP";
+    private static final long MAX_BATTERY_STATS_SNAPSHOT_STORAGE_BYTES = 100 * 1024;
+
+    private final Context mContext;
+    private final BatteryStatsImpl mBatteryStats;
+    private final File mStoreDir;
+    private final File mLockFile;
+    private final AtomicFile mConfigFile;
+    private final long mMaxStorageBytes;
+    private final Handler mHandler;
+    private final BatteryUsageStatsProvider mBatteryUsageStatsProvider;
+
+    public BatteryUsageStatsStore(Context context, BatteryStatsImpl stats, File systemDir,
+            Handler handler) {
+        this(context, stats, systemDir, handler, MAX_BATTERY_STATS_SNAPSHOT_STORAGE_BYTES);
+    }
+
+    @VisibleForTesting
+    public BatteryUsageStatsStore(Context context, BatteryStatsImpl batteryStats, File systemDir,
+            Handler handler, long maxStorageBytes) {
+        mContext = context;
+        mBatteryStats = batteryStats;
+        mStoreDir = new File(systemDir, BATTERY_USAGE_STATS_DIR);
+        mLockFile = new File(mStoreDir, DIR_LOCK_FILENAME);
+        mConfigFile = new AtomicFile(new File(mStoreDir, CONFIG_FILENAME));
+        mHandler = handler;
+        mMaxStorageBytes = maxStorageBytes;
+        mBatteryStats.setBatteryResetListener(this::prepareForBatteryStatsReset);
+        mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(mContext, mBatteryStats);
+    }
+
+    private void prepareForBatteryStatsReset() {
+        final List<BatteryUsageStats> stats =
+                mBatteryUsageStatsProvider.getBatteryUsageStats(BATTERY_USAGE_STATS_QUERY);
+        if (stats.isEmpty()) {
+            Slog.wtf(TAG, "No battery usage stats generated");
+            return;
+        }
+
+        mHandler.post(() -> storeBatteryUsageStats(stats.get(0)));
+    }
+
+    private void storeBatteryUsageStats(BatteryUsageStats stats) {
+        try (FileLock lock = lockSnapshotDirectory()) {
+            if (!mStoreDir.exists()) {
+                if (!mStoreDir.mkdirs()) {
+                    Slog.e(TAG,
+                            "Could not create a directory for battery usage stats snapshots");
+                    return;
+                }
+            }
+            File file = makeSnapshotFilename(stats.getStatsEndTimestamp());
+            try {
+                writeXmlFileLocked(stats, file);
+            } catch (Exception e) {
+                Slog.e(TAG, "Cannot save battery usage stats", e);
+            }
+
+            removeOldSnapshotsLocked();
+        } catch (IOException e) {
+            Slog.e(TAG, "Cannot lock battery usage stats directory", e);
+        }
+    }
+
+    /**
+     * Returns the timestamps of the stored BatteryUsageStats snapshots. The timestamp corresponds
+     * to the time the snapshot was taken {@link BatteryUsageStats#getStatsEndTimestamp()}.
+     */
+    public long[] listBatteryUsageStatsTimestamps() {
+        LongArray timestamps = new LongArray(100);
+        try (FileLock lock = lockSnapshotDirectory()) {
+            for (File file : mStoreDir.listFiles()) {
+                String fileName = file.getName();
+                if (fileName.endsWith(SNAPSHOT_FILE_EXTENSION)) {
+                    try {
+                        String fileNameWithoutExtension = fileName.substring(0,
+                                fileName.length() - SNAPSHOT_FILE_EXTENSION.length());
+                        timestamps.add(Long.parseLong(fileNameWithoutExtension));
+                    } catch (NumberFormatException e) {
+                        Slog.wtf(TAG, "Invalid format of BatteryUsageStats snapshot file name: "
+                                + fileName);
+                    }
+                }
+            }
+        } catch (IOException e) {
+            Slog.e(TAG, "Cannot lock battery usage stats directory", e);
+        }
+        return timestamps.toArray();
+    }
+
+    /**
+     * Reads the specified snapshot of BatteryUsageStats.  Returns null if the snapshot
+     * does not exist.
+     */
+    @Nullable
+    public BatteryUsageStats loadBatteryUsageStats(long timestamp) {
+        try (FileLock lock = lockSnapshotDirectory()) {
+            File file = makeSnapshotFilename(timestamp);
+            try {
+                return readXmlFileLocked(file);
+            } catch (Exception e) {
+                Slog.e(TAG, "Cannot read battery usage stats", e);
+            }
+        } catch (IOException e) {
+            Slog.e(TAG, "Cannot lock battery usage stats directory", e);
+        }
+        return null;
+    }
+
+    /**
+     * Saves the supplied timestamp of the BATTERY_USAGE_STATS_BEFORE_RESET statsd atom pull
+     * in persistent file.
+     */
+    public void setLastBatteryUsageStatsBeforeResetAtomPullTimestamp(long timestamp) {
+        Properties props = new Properties();
+        try (FileLock lock = lockSnapshotDirectory()) {
+            try (InputStream in = mConfigFile.openRead()) {
+                props.load(in);
+            } catch (IOException e) {
+                Slog.e(TAG, "Cannot load config file " + mConfigFile, e);
+            }
+            props.put(BATTERY_USAGE_STATS_BEFORE_RESET_TIMESTAMP_PROPERTY,
+                    String.valueOf(timestamp));
+            FileOutputStream out = null;
+            try {
+                out = mConfigFile.startWrite();
+                props.store(out, "Statsd atom pull timestamps");
+                mConfigFile.finishWrite(out);
+            } catch (IOException e) {
+                mConfigFile.failWrite(out);
+                Slog.e(TAG, "Cannot save config file " + mConfigFile, e);
+            }
+        } catch (IOException e) {
+            Slog.e(TAG, "Cannot lock battery usage stats directory", e);
+        }
+    }
+
+    /**
+     * Retrieves the previously saved timestamp of the last BATTERY_USAGE_STATS_BEFORE_RESET
+     * statsd atom pull.
+     */
+    public long getLastBatteryUsageStatsBeforeResetAtomPullTimestamp() {
+        Properties props = new Properties();
+        try (FileLock lock = lockSnapshotDirectory()) {
+            try (InputStream in = mConfigFile.openRead()) {
+                props.load(in);
+            } catch (IOException e) {
+                Slog.e(TAG, "Cannot load config file " + mConfigFile, e);
+            }
+        } catch (IOException e) {
+            Slog.e(TAG, "Cannot lock battery usage stats directory", e);
+        }
+        return Long.parseLong(
+                props.getProperty(BATTERY_USAGE_STATS_BEFORE_RESET_TIMESTAMP_PROPERTY, "0"));
+    }
+
+    private FileLock lockSnapshotDirectory() throws IOException {
+        mLockFile.getParentFile().mkdirs();
+        mLockFile.createNewFile();
+        return FileChannel.open(mLockFile.toPath(), StandardOpenOption.WRITE).lock();
+    }
+
+    /**
+     * Creates a file name by formatting the timestamp as 19-digit zero-padded number.
+     * This ensures that sorted directory list follows the chronological order.
+     */
+    private File makeSnapshotFilename(long statsEndTimestamp) {
+        return new File(mStoreDir, String.format(Locale.ENGLISH, "%019d", statsEndTimestamp)
+                + SNAPSHOT_FILE_EXTENSION);
+    }
+
+    private void writeXmlFileLocked(BatteryUsageStats stats, File file) throws IOException {
+        try (OutputStream out = new FileOutputStream(file)) {
+            TypedXmlSerializer serializer = Xml.newBinarySerializer();
+            serializer.setOutput(out, StandardCharsets.UTF_8.name());
+            serializer.startDocument(null, true);
+            stats.writeXml(serializer);
+            serializer.endDocument();
+        }
+    }
+
+    private BatteryUsageStats readXmlFileLocked(File file)
+            throws IOException, XmlPullParserException {
+        try (InputStream in = new FileInputStream(file)) {
+            TypedXmlPullParser parser = Xml.newBinaryPullParser();
+            parser.setInput(in, StandardCharsets.UTF_8.name());
+            return BatteryUsageStats.createFromXml(parser);
+        }
+    }
+
+    private void removeOldSnapshotsLocked() {
+        // Read the directory list into a _sorted_ map.  The alphanumeric ordering
+        // corresponds to the historical order of snapshots because the file names
+        // are timestamps zero-padded to the same length.
+        long totalSize = 0;
+        TreeMap<File, Long> mFileSizes = new TreeMap<>();
+        for (File file : mStoreDir.listFiles()) {
+            final long fileSize = file.length();
+            totalSize += fileSize;
+            if (file.getName().endsWith(SNAPSHOT_FILE_EXTENSION)) {
+                mFileSizes.put(file, fileSize);
+            }
+        }
+
+        while (totalSize > mMaxStorageBytes) {
+            final Map.Entry<File, Long> entry = mFileSizes.firstEntry();
+            if (entry == null) {
+                break;
+            }
+
+            File file = entry.getKey();
+            if (!file.delete()) {
+                Slog.e(TAG, "Cannot delete battery usage stats " + file);
+            }
+            totalSize -= entry.getValue();
+            mFileSizes.remove(file);
+        }
+    }
+}

core/tests/coretests/src/com/android/internal/os/BatteryStatsTests.java

@@ -44,6 +44,7 @@ import org.junit.runners.Suite;
         BatteryStatsUidTest.class,
         BatteryUsageStatsProviderTest.class,
         BatteryUsageStatsTest.class,
+        BatteryUsageStatsStoreTest.class,
         BatteryStatsUserLifecycleTests.class,
         BluetoothPowerCalculatorTest.class,
         BstatsCpuTimesValidationTest.class,

core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsProviderTest.java

@@ -20,10 +20,14 @@ import static com.google.common.truth.Truth.assertThat;
 
 import android.app.ActivityManager;
 import android.content.Context;
+import android.os.BatteryConsumer;
 import android.os.BatteryManager;
 import android.os.BatteryStats;
 import android.os.BatteryUsageStats;
 import android.os.BatteryUsageStatsQuery;
+import android.os.Handler;
+import android.os.Looper;
+import android.os.Message;
 import android.os.Parcel;
 import android.os.Process;
 import android.os.UidBatteryConsumer;
@@ -36,6 +40,7 @@ import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.io.File;
 import java.util.List;
 
 @SmallTest
@@ -45,7 +50,8 @@ public class BatteryUsageStatsProviderTest {
     private static final long MINUTE_IN_MS = 60 * 1000;
 
     @Rule
-    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule(12345);
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule(12345)
+            .setAveragePower(PowerProfile.POWER_FLASHLIGHT, 360.0);
 
     @Test
     public void test_getBatteryUsageStats() {
@@ -187,4 +193,84 @@ public class BatteryUsageStatsProviderTest {
         mStatsRule.setTime(11500, 0);
         assertThat(provider.shouldUpdateStats(queries, 10000)).isTrue();
     }
+
+    @Test
+    public void testAggregateBatteryStats() {
+        Context context = InstrumentationRegistry.getContext();
+        BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
+        mStatsRule.setCurrentTime(5 * MINUTE_IN_MS);
+        batteryStats.resetAllStatsCmdLocked();
+
+        BatteryUsageStatsStore batteryUsageStatsStore = new BatteryUsageStatsStore(context,
+                batteryStats, new File(context.getCacheDir(), "BatteryUsageStatsProviderTest"),
+                new TestHandler(), Integer.MAX_VALUE);
+
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(context,
+                batteryStats, batteryUsageStatsStore);
+
+        batteryStats.noteFlashlightOnLocked(APP_UID,
+                10 * MINUTE_IN_MS, 10 * MINUTE_IN_MS);
+        batteryStats.noteFlashlightOffLocked(APP_UID,
+                20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
+        mStatsRule.setCurrentTime(25 * MINUTE_IN_MS);
+        batteryStats.resetAllStatsCmdLocked();
+
+        batteryStats.noteFlashlightOnLocked(APP_UID,
+                30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS);
+        batteryStats.noteFlashlightOffLocked(APP_UID,
+                50 * MINUTE_IN_MS, 50 * MINUTE_IN_MS);
+        mStatsRule.setCurrentTime(55 * MINUTE_IN_MS);
+        batteryStats.resetAllStatsCmdLocked();
+
+        // This section should be ignored because the timestamp is out or range
+        batteryStats.noteFlashlightOnLocked(APP_UID,
+                60 * MINUTE_IN_MS, 60 * MINUTE_IN_MS);
+        batteryStats.noteFlashlightOffLocked(APP_UID,
+                70 * MINUTE_IN_MS, 70 * MINUTE_IN_MS);
+        mStatsRule.setCurrentTime(75 * MINUTE_IN_MS);
+        batteryStats.resetAllStatsCmdLocked();
+
+        // This section should be ignored because it represents the current stats session
+        batteryStats.noteFlashlightOnLocked(APP_UID,
+                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.noteFlashlightOffLocked(APP_UID,
+                90 * MINUTE_IN_MS, 90 * MINUTE_IN_MS);
+        mStatsRule.setCurrentTime(95 * MINUTE_IN_MS);
+
+        // Include the first and the second snapshot, but not the third or current
+        BatteryUsageStatsQuery query = new BatteryUsageStatsQuery.Builder()
+                .aggregateSnapshots(20 * MINUTE_IN_MS, 60 * MINUTE_IN_MS)
+                .build();
+        final BatteryUsageStats stats = provider.getBatteryUsageStats(query);
+
+        assertThat(stats.getStatsStartTimestamp()).isEqualTo(5 * MINUTE_IN_MS);
+        assertThat(stats.getStatsEndTimestamp()).isEqualTo(55 * MINUTE_IN_MS);
+        assertThat(stats.getAggregateBatteryConsumer(
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
+                .getConsumedPower(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
+                .isWithin(0.0001)
+                .of(180.0);  // 360 mA * 0.5 hours
+        assertThat(stats.getAggregateBatteryConsumer(
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
+                .getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
+                .isEqualTo((10 + 20) * MINUTE_IN_MS);
+        final UidBatteryConsumer uidBatteryConsumer = stats.getUidBatteryConsumers().stream()
+                .filter(uid -> uid.getUid() == APP_UID).findFirst().get();
+        assertThat(uidBatteryConsumer
+                .getConsumedPower(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
+                .isWithin(0.1)
+                .of(180.0);
+    }
+
+    private static class TestHandler extends Handler {
+        TestHandler() {
+            super(Looper.getMainLooper());
+        }
+
+        @Override
+        public boolean sendMessageAtTime(Message msg, long uptimeMillis) {
+            msg.getCallback().run();
+            return true;
+        }
+    }
 }

core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsStoreTest.java

@@ -0,0 +1,196 @@
+/*
+ * Copyright (C) 2021 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.internal.os;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.mock;
+
+import android.content.Context;
+import android.os.BatteryManager;
+import android.os.BatteryUsageStats;
+import android.os.BatteryUsageStatsQuery;
+import android.os.Handler;
+import android.os.Looper;
+import android.os.Message;
+import android.util.TypedXmlSerializer;
+import android.util.Xml;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+
+@RunWith(AndroidJUnit4.class)
+public class BatteryUsageStatsStoreTest {
+    private static final long MAX_BATTERY_STATS_SNAPSHOT_STORAGE_BYTES = 2 * 1024;
+
+    private final MockClocks mMockClocks = new MockClocks();
+    private MockBatteryStatsImpl mBatteryStats;
+    private BatteryUsageStatsStore mBatteryUsageStatsStore;
+    private BatteryUsageStatsProvider mBatteryUsageStatsProvider;
+    private File mStoreDirectory;
+
+    @Before
+    public void setup() {
+        mMockClocks.currentTime = 123;
+        mBatteryStats = new MockBatteryStatsImpl(mMockClocks);
+        mBatteryStats.setNoAutoReset(true);
+        mBatteryStats.setPowerProfile(mock(PowerProfile.class));
+
+        Context context = InstrumentationRegistry.getContext();
+
+        mStoreDirectory = new File(context.getCacheDir(), "BatteryUsageStatsStoreTest");
+        clearDirectory(mStoreDirectory);
+
+        mBatteryUsageStatsStore = new BatteryUsageStatsStore(context, mBatteryStats,
+                mStoreDirectory, new TestHandler(), MAX_BATTERY_STATS_SNAPSHOT_STORAGE_BYTES);
+
+        mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context, mBatteryStats);
+    }
+
+    @Test
+    public void testStoreSnapshot() {
+        mMockClocks.currentTime = 1_600_000;
+
+        prepareBatteryStats();
+        mBatteryStats.resetAllStatsCmdLocked();
+
+        final long[] timestamps = mBatteryUsageStatsStore.listBatteryUsageStatsTimestamps();
+        assertThat(timestamps).hasLength(1);
+        assertThat(timestamps[0]).isEqualTo(1_600_000);
+
+        final BatteryUsageStats batteryUsageStats = mBatteryUsageStatsStore.loadBatteryUsageStats(
+                1_600_000);
+        assertThat(batteryUsageStats.getStatsStartTimestamp()).isEqualTo(123);
+        assertThat(batteryUsageStats.getStatsEndTimestamp()).isEqualTo(1_600_000);
+        assertThat(batteryUsageStats.getBatteryCapacity()).isEqualTo(4000);
+        assertThat(batteryUsageStats.getDischargePercentage()).isEqualTo(5);
+        assertThat(batteryUsageStats.getAggregateBatteryConsumer(
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE).getConsumedPower())
+                .isEqualTo(600);  // (3_600_000 - 3_000_000) / 1000
+    }
+
+    @Test
+    public void testGarbageCollectOldSnapshots() throws Exception {
+        prepareBatteryStats();
+
+        mMockClocks.realtime = 10_000_000;
+        mMockClocks.uptime = 10_000_000;
+        mMockClocks.currentTime = 10_000_000;
+
+        final int snapshotFileSize = getSnapshotFileSize();
+        final int numberOfSnapshots =
+                (int) (MAX_BATTERY_STATS_SNAPSHOT_STORAGE_BYTES / snapshotFileSize);
+        for (int i = 0; i < numberOfSnapshots + 2; i++) {
+            mBatteryStats.resetAllStatsCmdLocked();
+
+            mMockClocks.realtime += 10_000_000;
+            mMockClocks.uptime += 10_000_000;
+            mMockClocks.currentTime += 10_000_000;
+            prepareBatteryStats();
+        }
+
+        final long[] timestamps = mBatteryUsageStatsStore.listBatteryUsageStatsTimestamps();
+        Arrays.sort(timestamps);
+        assertThat(timestamps).hasLength(numberOfSnapshots);
+        // Two snapshots (10_000_000 and 20_000_000) should have been discarded
+        assertThat(timestamps[0]).isEqualTo(30_000_000);
+        assertThat(getDirectorySize(mStoreDirectory))
+                .isAtMost(MAX_BATTERY_STATS_SNAPSHOT_STORAGE_BYTES);
+    }
+
+    @Test
+    public void testSavingStatsdAtomPullTimestamp() {
+        mBatteryUsageStatsStore.setLastBatteryUsageStatsBeforeResetAtomPullTimestamp(1234);
+        assertThat(mBatteryUsageStatsStore.getLastBatteryUsageStatsBeforeResetAtomPullTimestamp())
+                .isEqualTo(1234);
+        mBatteryUsageStatsStore.setLastBatteryUsageStatsBeforeResetAtomPullTimestamp(5478);
+        assertThat(mBatteryUsageStatsStore.getLastBatteryUsageStatsBeforeResetAtomPullTimestamp())
+                .isEqualTo(5478);
+    }
+
+    private void prepareBatteryStats() {
+        mBatteryStats.setBatteryStateLocked(BatteryManager.BATTERY_STATUS_DISCHARGING, 100,
+                /* plugType */ 0, 90, 72, 3700, 3_600_000, 4_000_000, 0,
+                mMockClocks.realtime, mMockClocks.uptime, mMockClocks.currentTime);
+        mBatteryStats.setBatteryStateLocked(BatteryManager.BATTERY_STATUS_DISCHARGING, 100,
+                /* plugType */ 0, 85, 72, 3700, 3_000_000, 4_000_000, 0,
+                mMockClocks.realtime + 500_000, mMockClocks.uptime + 500_000,
+                mMockClocks.currentTime + 500_000);
+    }
+
+    private void clearDirectory(File dir) {
+        if (dir.exists()) {
+            for (File child : dir.listFiles()) {
+                if (child.isDirectory()) {
+                    clearDirectory(child);
+                }
+                child.delete();
+            }
+        }
+    }
+
+    private long getDirectorySize(File dir) {
+        long size = 0;
+        if (dir.exists()) {
+            for (File child : dir.listFiles()) {
+                if (child.isDirectory()) {
+                    size += getDirectorySize(child);
+                } else {
+                    size += child.length();
+                }
+            }
+        }
+        return size;
+    }
+
+    private int getSnapshotFileSize() throws IOException {
+        BatteryUsageStats stats = mBatteryUsageStatsProvider.getBatteryUsageStats(
+                new BatteryUsageStatsQuery.Builder()
+                        .setMaxStatsAgeMs(0)
+                        .includePowerModels()
+                        .build());
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        TypedXmlSerializer serializer = Xml.newBinarySerializer();
+        serializer.setOutput(out, StandardCharsets.UTF_8.name());
+        serializer.startDocument(null, true);
+        stats.writeXml(serializer);
+        serializer.endDocument();
+        return out.toByteArray().length;
+    }
+
+    private static class TestHandler extends Handler {
+        TestHandler() {
+            super(Looper.getMainLooper());
+        }
+
+        @Override
+        public boolean sendMessageAtTime(Message msg, long uptimeMillis) {
+            msg.getCallback().run();
+            return true;
+        }
+    }
+}

core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsTest.java

@@ -16,8 +16,13 @@
 
 package com.android.internal.os;
 
+import static android.os.BatteryConsumer.POWER_MODEL_MEASURED_ENERGY;
+import static android.os.BatteryConsumer.POWER_MODEL_POWER_PROFILE;
+import static android.os.BatteryConsumer.POWER_MODEL_UNDEFINED;
+
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.fail;
 
 import android.os.BatteryConsumer;
@@ -25,6 +30,9 @@ import android.os.BatteryUsageStats;
 import android.os.Parcel;
 import android.os.UidBatteryConsumer;
 import android.os.UserBatteryConsumer;
+import android.util.TypedXmlPullParser;
+import android.util.TypedXmlSerializer;
+import android.util.Xml;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
@@ -32,8 +40,11 @@ import androidx.test.runner.AndroidJUnit4;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.PrintWriter;
 import java.io.StringWriter;
+import java.nio.charset.StandardCharsets;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -42,15 +53,19 @@ import java.util.Set;
 @RunWith(AndroidJUnit4.class)
 public class BatteryUsageStatsTest {
 
+    private static final int USER_ID = 42;
+    private static final int APP_UID1 = 271;
+    private static final int APP_UID2 = 314;
+
     @Test
     public void testBuilder() {
-        BatteryUsageStats batteryUsageStats = buildBatteryUsageStats().build();
-        validateBatteryUsageStats(batteryUsageStats);
+        BatteryUsageStats batteryUsageStats = buildBatteryUsageStats1(true).build();
+        assertBatteryUsageStats1(batteryUsageStats, true);
     }
 
     @Test
     public void testParcelability() {
-        final BatteryUsageStats outBatteryUsageStats = buildBatteryUsageStats().build();
+        final BatteryUsageStats outBatteryUsageStats = buildBatteryUsageStats1(true).build();
         final Parcel outParcel = Parcel.obtain();
         outParcel.writeParcelable(outBatteryUsageStats, 0);
         final byte[] bytes = outParcel.marshall();
@@ -62,20 +77,20 @@ public class BatteryUsageStatsTest {
         final BatteryUsageStats inBatteryUsageStats =
                 inParcel.readParcelable(getClass().getClassLoader());
         assertThat(inBatteryUsageStats).isNotNull();
-        validateBatteryUsageStats(inBatteryUsageStats);
+        assertBatteryUsageStats1(inBatteryUsageStats, true);
     }
 
 
     @Test
     public void testDefaultSessionDuration() {
         final BatteryUsageStats stats =
-                buildBatteryUsageStats().setStatsDuration(10000).build();
+                buildBatteryUsageStats1(true).setStatsDuration(10000).build();
         assertThat(stats.getStatsDuration()).isEqualTo(10000);
     }
 
     @Test
     public void testDump() {
-        final BatteryUsageStats stats = buildBatteryUsageStats().build();
+        final BatteryUsageStats stats = buildBatteryUsageStats1(true).build();
         final StringWriter out = new StringWriter();
         try (PrintWriter pw = new PrintWriter(out)) {
             stats.dump(pw, "  ");
@@ -87,7 +102,7 @@ public class BatteryUsageStatsTest {
         assertThat(dump).contains("actual drain: 1000-2000");
         assertThat(dump).contains("cpu: 20100 apps: 10100 duration: 20s 300ms");
         assertThat(dump).contains("FOO: 20200 apps: 10200 duration: 20s 400ms");
-        assertThat(dump).contains("UID 2000: 1200 ( screen=300 cpu=400 FOO=500 )");
+        assertThat(dump).contains("UID 271: 1200 ( screen=300 cpu=400 FOO=500 )");
         assertThat(dump).contains("User 42: 30.0 ( cpu=10.0 FOO=20.0 )");
     }
 
@@ -101,154 +116,297 @@ public class BatteryUsageStatsTest {
         assertThat(allNames).hasSize(BatteryConsumer.POWER_COMPONENT_COUNT);
     }
 
-    private BatteryUsageStats.Builder buildBatteryUsageStats() {
-        final MockClocks clocks = new MockClocks();
-        final MockBatteryStatsImpl batteryStats = new MockBatteryStatsImpl(clocks);
-        final BatteryStatsImpl.Uid batteryStatsUid = batteryStats.getUidStatsLocked(2000);
+    @Test
+    public void testAdd() {
+        final BatteryUsageStats stats1 = buildBatteryUsageStats1(false).build();
+        final BatteryUsageStats stats2 = buildBatteryUsageStats2(new String[] {"FOO"}).build();
 
-        final BatteryUsageStats.Builder builder =
-                new BatteryUsageStats.Builder(new String[]{"FOO"}, true)
-                        .setBatteryCapacity(4000)
-                        .setDischargePercentage(20)
-                        .setDischargedPowerRange(1000, 2000)
-                        .setStatsStartTimestamp(1000)
-                        .setStatsEndTimestamp(3000);
-        builder.getOrCreateUidBatteryConsumerBuilder(batteryStatsUid)
-                .setPackageWithHighestDrain("foo")
-                .setTimeInStateMs(UidBatteryConsumer.STATE_FOREGROUND, 1000)
-                .setTimeInStateMs(UidBatteryConsumer.STATE_BACKGROUND, 2000)
-                .setConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_SCREEN, 300)
-                .setConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 400)
-                .setConsumedPowerForCustomComponent(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 500)
-                .setUsageDurationMillis(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 600)
-                .setUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 800);
+        final BatteryUsageStats sum =
+                new BatteryUsageStats.Builder(new String[] {"FOO"}, true)
+                        .add(stats1)
+                        .add(stats2)
+                        .build();
 
-        builder.getAggregateBatteryConsumerBuilder(
-                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS)
-                .setConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 10100)
-                .setConsumedPowerForCustomComponent(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 10200)
-                .setUsageDurationMillis(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 10300)
-                .setUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 10400);
-
-        builder.getAggregateBatteryConsumerBuilder(
-                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
-                .setConsumedPower(30000)
-                .setConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 20100)
-                .setConsumedPowerForCustomComponent(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 20200)
-                .setUsageDurationMillis(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 20300)
-                .setUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 20400);
-
-        builder.getOrCreateUserBatteryConsumerBuilder(42)
-                .setConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 10)
-                .setConsumedPowerForCustomComponent(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 20)
-                .setUsageDurationMillis(
-                        BatteryConsumer.POWER_COMPONENT_CPU, 30)
-                .setUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 40);
-
-        return builder;
-    }
-
-    public void validateBatteryUsageStats(BatteryUsageStats batteryUsageStats) {
-        assertThat(batteryUsageStats.getConsumedPower()).isEqualTo(30000);
-        assertThat(batteryUsageStats.getBatteryCapacity()).isEqualTo(4000);
-        assertThat(batteryUsageStats.getDischargePercentage()).isEqualTo(20);
-        assertThat(batteryUsageStats.getDischargedPowerRange().getLower()).isEqualTo(1000);
-        assertThat(batteryUsageStats.getDischargedPowerRange().getUpper()).isEqualTo(2000);
-        assertThat(batteryUsageStats.getStatsStartTimestamp()).isEqualTo(1000);
-        assertThat(batteryUsageStats.getStatsEndTimestamp()).isEqualTo(3000);
-        assertThat(batteryUsageStats.getStatsDuration()).isEqualTo(2000);
+        assertBatteryUsageStats(sum, 42345, 50, 2234, 4345, 1000, 5000, 5000);
 
         final List<UidBatteryConsumer> uidBatteryConsumers =
-                batteryUsageStats.getUidBatteryConsumers();
+                sum.getUidBatteryConsumers();
         for (UidBatteryConsumer uidBatteryConsumer : uidBatteryConsumers) {
-            if (uidBatteryConsumer.getUid() == 2000) {
-                assertThat(uidBatteryConsumer.getPackageWithHighestDrain()).isEqualTo("foo");
-                assertThat(uidBatteryConsumer.getTimeInStateMs(
-                        UidBatteryConsumer.STATE_FOREGROUND)).isEqualTo(1000);
-                assertThat(uidBatteryConsumer.getTimeInStateMs(
-                        UidBatteryConsumer.STATE_BACKGROUND)).isEqualTo(2000);
-                assertThat(uidBatteryConsumer.getConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_SCREEN)).isEqualTo(300);
-                assertThat(uidBatteryConsumer.getConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(400);
-                assertThat(uidBatteryConsumer.getConsumedPowerForCustomComponent(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(500);
-                assertThat(uidBatteryConsumer.getUsageDurationMillis(
-                        BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(600);
-                assertThat(uidBatteryConsumer.getUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(800);
-                assertThat(uidBatteryConsumer.getConsumedPower()).isEqualTo(1200);
-                assertThat(uidBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
-                assertThat(uidBatteryConsumer.getCustomPowerComponentName(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
+            if (uidBatteryConsumer.getUid() == APP_UID1) {
+                assertUidBatteryConsumer(uidBatteryConsumer, 2124, null,
+                        5321, 7432, 423, POWER_MODEL_POWER_PROFILE, 745, POWER_MODEL_UNDEFINED,
+                        956, 1167, 1478);
+            } else if (uidBatteryConsumer.getUid() == APP_UID2) {
+                assertUidBatteryConsumer(uidBatteryConsumer, 1332, "bar",
+                        1111, 2222, 333, POWER_MODEL_POWER_PROFILE, 444, POWER_MODEL_POWER_PROFILE,
+                        555, 666, 777);
             } else {
                 fail("Unexpected UID " + uidBatteryConsumer.getUid());
             }
         }
 
-        final BatteryConsumer appsBatteryConsumer = batteryUsageStats.getAggregateBatteryConsumer(
-                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS);
-        assertThat(appsBatteryConsumer.getConsumedPower(
-                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(10100);
-        assertThat(appsBatteryConsumer.getConsumedPowerForCustomComponent(
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(10200);
-        assertThat(appsBatteryConsumer.getUsageDurationMillis(
-                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(10300);
-        assertThat(appsBatteryConsumer.getUsageDurationForCustomComponentMillis(
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(10400);
-        assertThat(appsBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
-        assertThat(appsBatteryConsumer.getCustomPowerComponentName(
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
+        assertAggregateBatteryConsumer(sum,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS,
+                20223, 20434, 20645, 20856);
 
-        final BatteryConsumer deviceBatteryConsumer = batteryUsageStats.getAggregateBatteryConsumer(
-                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE);
-        assertThat(deviceBatteryConsumer.getConsumedPower(
-                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(20100);
-        assertThat(deviceBatteryConsumer.getConsumedPowerForCustomComponent(
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(20200);
-        assertThat(deviceBatteryConsumer.getUsageDurationMillis(
-                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(20300);
-        assertThat(deviceBatteryConsumer.getUsageDurationForCustomComponentMillis(
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(20400);
-        assertThat(deviceBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
-        assertThat(deviceBatteryConsumer.getCustomPowerComponentName(
-                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
+        assertAggregateBatteryConsumer(sum,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE,
+                40211, 40422, 40633, 40844);
+    }
+
+    @Test
+    public void testAdd_customComponentMismatch() {
+        final BatteryUsageStats.Builder builder =
+                new BatteryUsageStats.Builder(new String[] {"FOO"}, true);
+        final BatteryUsageStats stats = buildBatteryUsageStats2(new String[] {"BAR"}).build();
+
+        assertThrows(IllegalArgumentException.class, () -> builder.add(stats));
+    }
+
+    @Test
+    public void testXml() throws Exception {
+        ByteArrayOutputStream out = new ByteArrayOutputStream();
+        TypedXmlSerializer serializer = Xml.newBinarySerializer();
+        serializer.setOutput(out, StandardCharsets.UTF_8.name());
+        serializer.startDocument(null, true);
+        final BatteryUsageStats stats = buildBatteryUsageStats1(true).build();
+        stats.writeXml(serializer);
+        serializer.endDocument();
+
+        ByteArrayInputStream in = new ByteArrayInputStream(out.toByteArray());
+        TypedXmlPullParser parser = Xml.newBinaryPullParser();
+        parser.setInput(in, StandardCharsets.UTF_8.name());
+        final BatteryUsageStats fromXml = BatteryUsageStats.createFromXml(parser);
+
+        assertBatteryUsageStats1(fromXml, true);
+    }
+
+    private BatteryUsageStats.Builder buildBatteryUsageStats1(boolean includeUserBatteryConsumer) {
+        final MockClocks clocks = new MockClocks();
+        final MockBatteryStatsImpl batteryStats = new MockBatteryStatsImpl(clocks);
+
+        final BatteryUsageStats.Builder builder =
+                new BatteryUsageStats.Builder(new String[] {"FOO"}, true)
+                        .setBatteryCapacity(4000)
+                        .setDischargePercentage(20)
+                        .setDischargedPowerRange(1000, 2000)
+                        .setStatsStartTimestamp(1000)
+                        .setStatsEndTimestamp(3000);
+
+        addUidBatteryConsumer(builder, batteryStats, APP_UID1, "foo",
+                1000, 2000,
+                300, POWER_MODEL_POWER_PROFILE, 400, POWER_MODEL_POWER_PROFILE, 500, 600, 800);
+
+        addAggregateBatteryConsumer(builder,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS, 0,
+                10100, 10200, 10300, 10400);
+
+        addAggregateBatteryConsumer(builder,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE, 30000,
+                20100, 20200, 20300, 20400);
+
+
+        if (includeUserBatteryConsumer) {
+            builder.getOrCreateUserBatteryConsumerBuilder(USER_ID)
+                    .setConsumedPower(
+                            BatteryConsumer.POWER_COMPONENT_CPU, 10)
+                    .setConsumedPowerForCustomComponent(
+                            BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 20)
+                    .setUsageDurationMillis(
+                            BatteryConsumer.POWER_COMPONENT_CPU, 30)
+                    .setUsageDurationForCustomComponentMillis(
+                            BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 40);
+        }
+        return builder;
+    }
+
+    private BatteryUsageStats.Builder buildBatteryUsageStats2(String[] customPowerComponentNames) {
+        final MockClocks clocks = new MockClocks();
+        final MockBatteryStatsImpl batteryStats = new MockBatteryStatsImpl(clocks);
+
+        final BatteryUsageStats.Builder builder =
+                new BatteryUsageStats.Builder(customPowerComponentNames, true)
+                        .setDischargePercentage(30)
+                        .setDischargedPowerRange(1234, 2345)
+                        .setStatsStartTimestamp(2000)
+                        .setStatsEndTimestamp(5000);
+
+        addUidBatteryConsumer(builder, batteryStats, APP_UID1, null,
+                4321, 5432,
+                123, POWER_MODEL_POWER_PROFILE, 345, POWER_MODEL_MEASURED_ENERGY, 456, 567, 678);
+
+        addUidBatteryConsumer(builder, batteryStats, APP_UID2, "bar",
+                1111, 2222,
+                333, POWER_MODEL_POWER_PROFILE, 444, POWER_MODEL_POWER_PROFILE, 555, 666, 777);
+
+        addAggregateBatteryConsumer(builder,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS, 0,
+                10123, 10234, 10345, 10456);
+
+        addAggregateBatteryConsumer(builder,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE, 12345,
+                20111, 20222, 20333, 20444);
+
+        return builder;
+    }
+
+    private void addUidBatteryConsumer(BatteryUsageStats.Builder builder,
+            MockBatteryStatsImpl batteryStats, int uid, String packageWithHighestDrain,
+            int timeInStateForeground, int timeInStateBackground, int screenPower,
+            int screenPowerModel, int cpuPower, int cpuPowerModel, int customComponentPower,
+            int cpuDuration, int customComponentDuration) {
+        final BatteryStatsImpl.Uid batteryStatsUid = batteryStats.getUidStatsLocked(uid);
+        builder.getOrCreateUidBatteryConsumerBuilder(batteryStatsUid)
+                .setPackageWithHighestDrain(packageWithHighestDrain)
+                .setTimeInStateMs(UidBatteryConsumer.STATE_FOREGROUND, timeInStateForeground)
+                .setTimeInStateMs(UidBatteryConsumer.STATE_BACKGROUND, timeInStateBackground)
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_SCREEN, screenPower, screenPowerModel)
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_CPU, cpuPower, cpuPowerModel)
+                .setConsumedPowerForCustomComponent(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, customComponentPower)
+                .setUsageDurationMillis(
+                        BatteryConsumer.POWER_COMPONENT_CPU, cpuDuration)
+                .setUsageDurationForCustomComponentMillis(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, customComponentDuration);
+    }
+
+    private void addAggregateBatteryConsumer(BatteryUsageStats.Builder builder, int scope,
+            double consumedPower, int cpuPower, int customComponentPower, int cpuDuration,
+            int customComponentDuration) {
+        builder.getAggregateBatteryConsumerBuilder(scope)
+                .setConsumedPower(consumedPower)
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_CPU, cpuPower)
+                .setConsumedPowerForCustomComponent(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, customComponentPower)
+                .setUsageDurationMillis(
+                        BatteryConsumer.POWER_COMPONENT_CPU, cpuDuration)
+                .setUsageDurationForCustomComponentMillis(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, customComponentDuration);
+    }
+
+    public void assertBatteryUsageStats1(BatteryUsageStats batteryUsageStats,
+            boolean includesUserBatteryConsumers) {
+        assertBatteryUsageStats(batteryUsageStats, 30000, 20, 1000, 2000, 1000, 3000, 2000);
+
+        final List<UidBatteryConsumer> uidBatteryConsumers =
+                batteryUsageStats.getUidBatteryConsumers();
+        assertThat(uidBatteryConsumers).hasSize(1);
+        for (UidBatteryConsumer uidBatteryConsumer : uidBatteryConsumers) {
+            if (uidBatteryConsumer.getUid() == APP_UID1) {
+                assertUidBatteryConsumer(uidBatteryConsumer, 1200, "foo",
+                        1000, 2000, 300, POWER_MODEL_POWER_PROFILE, 400, POWER_MODEL_POWER_PROFILE,
+                        500, 600, 800);
+            } else {
+                fail("Unexpected UID " + uidBatteryConsumer.getUid());
+            }
+        }
 
         final List<UserBatteryConsumer> userBatteryConsumers =
                 batteryUsageStats.getUserBatteryConsumers();
-        for (UserBatteryConsumer userBatteryConsumer : userBatteryConsumers) {
-            if (userBatteryConsumer.getUserId() == 42) {
-                assertThat(userBatteryConsumer.getConsumedPower(
-                        BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(10);
-                assertThat(userBatteryConsumer.getConsumedPowerForCustomComponent(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(20);
-                assertThat(userBatteryConsumer.getUsageDurationMillis(
-                        BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(30);
-                assertThat(userBatteryConsumer.getUsageDurationForCustomComponentMillis(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(40);
-                assertThat(userBatteryConsumer.getConsumedPower()).isEqualTo(30);
-                assertThat(userBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
-                assertThat(userBatteryConsumer.getCustomPowerComponentName(
-                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
-            } else {
-                fail("Unexpected user ID " + userBatteryConsumer.getUserId());
+        if (includesUserBatteryConsumers) {
+            assertThat(userBatteryConsumers).hasSize(1);
+            for (UserBatteryConsumer userBatteryConsumer : userBatteryConsumers) {
+                if (userBatteryConsumer.getUserId() == USER_ID) {
+                    assertUserBatteryConsumer(userBatteryConsumer, 42, 10, 20, 30, 40);
+                } else {
+                    fail("Unexpected User ID " + userBatteryConsumer.getUserId());
+                }
             }
+        } else {
+            assertThat(userBatteryConsumers).isEmpty();
         }
+
+        assertAggregateBatteryConsumer(batteryUsageStats,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS,
+                10100, 10200, 10300, 10400);
+
+        assertAggregateBatteryConsumer(batteryUsageStats,
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE,
+                20100, 20200, 20300, 20400);
+    }
+
+    private void assertBatteryUsageStats(BatteryUsageStats batteryUsageStats, int consumedPower,
+            int dischargePercentage, int dischagePowerLower, int dischargePowerUpper,
+            int statsStartTimestamp, int statsEndTimestamp, int statsDuration) {
+        assertThat(batteryUsageStats.getConsumedPower()).isEqualTo(consumedPower);
+        assertThat(batteryUsageStats.getDischargePercentage()).isEqualTo(dischargePercentage);
+        assertThat(batteryUsageStats.getDischargedPowerRange().getLower()).isEqualTo(
+                dischagePowerLower);
+        assertThat(batteryUsageStats.getDischargedPowerRange().getUpper()).isEqualTo(
+                dischargePowerUpper);
+        assertThat(batteryUsageStats.getStatsStartTimestamp()).isEqualTo(statsStartTimestamp);
+        assertThat(batteryUsageStats.getStatsEndTimestamp()).isEqualTo(statsEndTimestamp);
+        assertThat(batteryUsageStats.getStatsDuration()).isEqualTo(statsDuration);
+    }
+
+    private void assertUidBatteryConsumer(UidBatteryConsumer uidBatteryConsumer,
+            int consumedPower, String packageWithHighestDrain, int timeInStateForeground,
+            int timeInStateBackground, int screenPower, int screenPowerModel, int cpuPower,
+            int cpuPowerModel, int customComponentPower, int cpuDuration,
+            int customComponentDuration) {
+        assertThat(uidBatteryConsumer.getConsumedPower()).isEqualTo(consumedPower);
+        assertThat(uidBatteryConsumer.getPackageWithHighestDrain()).isEqualTo(
+                packageWithHighestDrain);
+        assertThat(uidBatteryConsumer.getTimeInStateMs(
+                UidBatteryConsumer.STATE_FOREGROUND)).isEqualTo(timeInStateForeground);
+        assertThat(uidBatteryConsumer.getTimeInStateMs(
+                UidBatteryConsumer.STATE_BACKGROUND)).isEqualTo(timeInStateBackground);
+        assertThat(uidBatteryConsumer.getConsumedPower(
+                BatteryConsumer.POWER_COMPONENT_SCREEN)).isEqualTo(screenPower);
+        assertThat(uidBatteryConsumer.getPowerModel(
+                BatteryConsumer.POWER_COMPONENT_SCREEN)).isEqualTo(screenPowerModel);
+        assertThat(uidBatteryConsumer.getConsumedPower(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuPower);
+        assertThat(uidBatteryConsumer.getPowerModel(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuPowerModel);
+        assertThat(uidBatteryConsumer.getConsumedPowerForCustomComponent(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(customComponentPower);
+        assertThat(uidBatteryConsumer.getUsageDurationMillis(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuDuration);
+        assertThat(uidBatteryConsumer.getUsageDurationForCustomComponentMillis(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(
+                customComponentDuration);
+        assertThat(uidBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
+        assertThat(uidBatteryConsumer.getCustomPowerComponentName(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
+    }
+
+    private void assertUserBatteryConsumer(UserBatteryConsumer userBatteryConsumer,
+            int userId, int cpuPower, int customComponentPower,
+            int cpuDuration, int customComponentDuration) {
+        assertThat(userBatteryConsumer.getConsumedPower(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuPower);
+        assertThat(userBatteryConsumer.getConsumedPowerForCustomComponent(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(customComponentPower);
+        assertThat(userBatteryConsumer.getUsageDurationMillis(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuDuration);
+        assertThat(userBatteryConsumer.getUsageDurationForCustomComponentMillis(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(
+                customComponentDuration);
+        assertThat(userBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
+        assertThat(userBatteryConsumer.getCustomPowerComponentName(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
+    }
+
+    private void assertAggregateBatteryConsumer(BatteryUsageStats batteryUsageStats,
+            int aggregateBatteryConsumerScopeAllApps, int cpuPower, int customComponentPower,
+            int cpuDuration, int customComponentDuration) {
+        final BatteryConsumer appsBatteryConsumer = batteryUsageStats.getAggregateBatteryConsumer(
+                aggregateBatteryConsumerScopeAllApps);
+        assertThat(appsBatteryConsumer.getConsumedPower(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuPower);
+        assertThat(appsBatteryConsumer.getConsumedPowerForCustomComponent(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(customComponentPower);
+        assertThat(appsBatteryConsumer.getUsageDurationMillis(
+                BatteryConsumer.POWER_COMPONENT_CPU)).isEqualTo(cpuDuration);
+        assertThat(appsBatteryConsumer.getUsageDurationForCustomComponentMillis(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo(
+                customComponentDuration);
+        assertThat(appsBatteryConsumer.getCustomPowerComponentCount()).isEqualTo(1);
+        assertThat(appsBatteryConsumer.getCustomPowerComponentName(
+                BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID)).isEqualTo("FOO");
     }
 }

services/core/java/com/android/server/am/BatteryStatsService.java

@@ -78,6 +78,7 @@ import com.android.internal.os.BackgroundThread;
 import com.android.internal.os.BatteryStatsHelper;
 import com.android.internal.os.BatteryStatsImpl;
 import com.android.internal.os.BatteryUsageStatsProvider;
+import com.android.internal.os.BatteryUsageStatsStore;
 import com.android.internal.os.BinderCallsStats;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.RailStats;
@@ -124,10 +125,12 @@ public final class BatteryStatsService extends IBatteryStats.Stub
         Watchdog.Monitor {
     static final String TAG = "BatteryStatsService";
     static final boolean DBG = false;
+    private static final boolean BATTERY_USAGE_STORE_ENABLED = true;
 
     private static IBatteryStats sService;
 
     final BatteryStatsImpl mStats;
+    private final BatteryUsageStatsStore mBatteryUsageStatsStore;
     private final BatteryStatsImpl.UserInfoProvider mUserManagerUserInfoProvider;
     private final Context mContext;
     private final BatteryExternalStatsWorker mWorker;
@@ -348,7 +351,14 @@ public final class BatteryStatsService extends IBatteryStats.Stub
         mStats.setPowerProfileLocked(new PowerProfile(context));
         mStats.startTrackingSystemServerCpuTime();
 
-        mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context, mStats);
+        if (BATTERY_USAGE_STORE_ENABLED) {
+            mBatteryUsageStatsStore =
+                    new BatteryUsageStatsStore(context, mStats, systemDir, mHandler);
+        } else {
+            mBatteryUsageStatsStore = null;
+        }
+        mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context, mStats,
+                mBatteryUsageStatsStore);
     }
 
     public void publish() {
@@ -752,6 +762,10 @@ public final class BatteryStatsService extends IBatteryStats.Stub
                 FrameworkStatsLog.BATTERY_USAGE_STATS_SINCE_RESET_USING_POWER_PROFILE_MODEL,
                 null, // use default PullAtomMetadata values
                 BackgroundThread.getExecutor(), pullAtomCallback);
+        statsManager.setPullAtomCallback(
+                FrameworkStatsLog.BATTERY_USAGE_STATS_BEFORE_RESET,
+                null, // use default PullAtomMetadata values
+                BackgroundThread.getExecutor(), pullAtomCallback);
     }
 
     /** StatsPullAtomCallback for pulling BatteryUsageStats data. */
@@ -768,6 +782,17 @@ public final class BatteryStatsService extends IBatteryStats.Stub
                             new BatteryUsageStatsQuery.Builder().powerProfileModeledOnly().build();
                     bus = getBatteryUsageStats(List.of(powerProfileQuery)).get(0);
                     break;
+                case FrameworkStatsLog.BATTERY_USAGE_STATS_BEFORE_RESET:
+                    final long sessionStart = mBatteryUsageStatsStore
+                            .getLastBatteryUsageStatsBeforeResetAtomPullTimestamp();
+                    final long sessionEnd = mStats.getStartClockTime();
+                    final BatteryUsageStatsQuery query = new BatteryUsageStatsQuery.Builder()
+                            .aggregateSnapshots(sessionStart, sessionEnd)
+                            .build();
+                    bus = getBatteryUsageStats(List.of(query)).get(0);
+                    mBatteryUsageStatsStore
+                            .setLastBatteryUsageStatsBeforeResetAtomPullTimestamp(sessionEnd);
+                    break;
                 default:
                     throw new UnsupportedOperationException("Unknown tagId=" + atomTag);
             }