BatteryUsageStats atom - frameworks/base

Writes the BatteryUsageStats atoms.proto atoms
based on the current BatteryUsageStats data in BatteryStats.

Does NOT write the past pre-reset snapshot atoms; that is
an adventure for a future cl.

Bug: 184095105
Test: atest BatteryUsageStatsProtoTests
Test: statsd_testdrive <atomId>
Change-Id: I2fc5a983deb58d7d393c0696db2165b124c94dc2

8 files changed, 611 insertions, 20 deletions

diff --git a/core/java/android/os/BatteryConsumer.java b/core/java/android/os/BatteryConsumer.java
index edb30b0f26ca..da94d74e5ae8 100644
--- a/core/java/android/os/BatteryConsumer.java
+++ b/core/java/android/os/BatteryConsumer.java
@@ -18,6 +18,8 @@ package android.os;
 
 import android.annotation.IntDef;
 import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.util.proto.ProtoOutputStream;
 
 import java.io.PrintWriter;
 import java.lang.annotation.Retention;
@@ -59,27 +61,32 @@ public abstract class BatteryConsumer {
     public static @interface PowerComponent {
     }
 
-    public static final int POWER_COMPONENT_SCREEN = 0;
-    public static final int POWER_COMPONENT_CPU = 1;
-    public static final int POWER_COMPONENT_BLUETOOTH = 2;
-    public static final int POWER_COMPONENT_CAMERA = 3;
-    public static final int POWER_COMPONENT_AUDIO = 4;
-    public static final int POWER_COMPONENT_VIDEO = 5;
-    public static final int POWER_COMPONENT_FLASHLIGHT = 6;
-    public static final int POWER_COMPONENT_SYSTEM_SERVICES = 7;
-    public static final int POWER_COMPONENT_MOBILE_RADIO = 8;
-    public static final int POWER_COMPONENT_SENSORS = 9;
-    public static final int POWER_COMPONENT_GNSS = 10;
-    public static final int POWER_COMPONENT_WIFI = 11;
-    public static final int POWER_COMPONENT_WAKELOCK = 12;
-    public static final int POWER_COMPONENT_MEMORY = 13;
-    public static final int POWER_COMPONENT_PHONE = 14;
-    public static final int POWER_COMPONENT_AMBIENT_DISPLAY = 15;
-    public static final int POWER_COMPONENT_IDLE = 16;
+    public static final int POWER_COMPONENT_SCREEN = OsProtoEnums.POWER_COMPONENT_SCREEN; // 0
+    public static final int POWER_COMPONENT_CPU = OsProtoEnums.POWER_COMPONENT_CPU; // 1
+    public static final int POWER_COMPONENT_BLUETOOTH = OsProtoEnums.POWER_COMPONENT_BLUETOOTH; // 2
+    public static final int POWER_COMPONENT_CAMERA = OsProtoEnums.POWER_COMPONENT_CAMERA; // 3
+    public static final int POWER_COMPONENT_AUDIO = OsProtoEnums.POWER_COMPONENT_AUDIO; // 4
+    public static final int POWER_COMPONENT_VIDEO = OsProtoEnums.POWER_COMPONENT_VIDEO; // 5
+    public static final int POWER_COMPONENT_FLASHLIGHT =
+            OsProtoEnums.POWER_COMPONENT_FLASHLIGHT; // 6
+    public static final int POWER_COMPONENT_SYSTEM_SERVICES =
+            OsProtoEnums.POWER_COMPONENT_SYSTEM_SERVICES; // 7
+    public static final int POWER_COMPONENT_MOBILE_RADIO =
+            OsProtoEnums.POWER_COMPONENT_MOBILE_RADIO; // 8
+    public static final int POWER_COMPONENT_SENSORS = OsProtoEnums.POWER_COMPONENT_SENSORS; // 9
+    public static final int POWER_COMPONENT_GNSS = OsProtoEnums.POWER_COMPONENT_GNSS; // 10
+    public static final int POWER_COMPONENT_WIFI = OsProtoEnums.POWER_COMPONENT_WIFI; // 11
+    public static final int POWER_COMPONENT_WAKELOCK = OsProtoEnums.POWER_COMPONENT_WAKELOCK; // 12
+    public static final int POWER_COMPONENT_MEMORY = OsProtoEnums.POWER_COMPONENT_MEMORY; // 13
+    public static final int POWER_COMPONENT_PHONE = OsProtoEnums.POWER_COMPONENT_PHONE; // 14
+    public static final int POWER_COMPONENT_AMBIENT_DISPLAY =
+            OsProtoEnums.POWER_COMPONENT_AMBIENT_DISPLAY; // 15
+    public static final int POWER_COMPONENT_IDLE = OsProtoEnums.POWER_COMPONENT_IDLE; // 16
     // Power that is re-attributed to other battery consumers. For example, for System Server
     // this represents the power attributed to apps requesting system services.
     // The value should be negative or zero.
-    public static final int POWER_COMPONENT_REATTRIBUTED_TO_OTHER_CONSUMERS = 17;
+    public static final int POWER_COMPONENT_REATTRIBUTED_TO_OTHER_CONSUMERS =
+            OsProtoEnums.POWER_COMPONENT_REATTRIBUTED_TO_OTHER_CONSUMERS; // 17
 
     public static final int POWER_COMPONENT_COUNT = 18;
 
@@ -263,6 +270,54 @@ public abstract class BatteryConsumer {
      */
     public abstract void dump(PrintWriter pw, boolean skipEmptyComponents);
 
+    /** Returns whether there are any atoms.proto BATTERY_CONSUMER_DATA data to write to a proto. */
+    boolean hasStatsProtoData() {
+        return writeStatsProtoImpl(null, /* Irrelevant fieldId: */ 0);
+    }
+
+    /** Writes the atoms.proto BATTERY_CONSUMER_DATA for this BatteryConsumer to the given proto. */
+    void writeStatsProto(@NonNull ProtoOutputStream proto, long fieldId) {
+        writeStatsProtoImpl(proto, fieldId);
+    }
+
+    /**
+     * Returns whether there are any atoms.proto BATTERY_CONSUMER_DATA data to write to a proto,
+     * and writes it to the given proto if it is non-null.
+     */
+    private boolean writeStatsProtoImpl(@Nullable ProtoOutputStream proto, long fieldId) {
+        final long totalConsumedPowerDeciCoulombs = convertMahToDeciCoulombs(getConsumedPower());
+
+        if (totalConsumedPowerDeciCoulombs == 0) {
+            // NOTE: Strictly speaking we should also check !mPowerComponents.hasStatsProtoData().
+            // However, that call is a bit expensive (a for loop). And the only way that
+            // totalConsumedPower can be 0 while mPowerComponents.hasStatsProtoData() is true is
+            // if POWER_COMPONENT_REATTRIBUTED_TO_OTHER_CONSUMERS (which is the only negative
+            // allowed) happens to exactly equal the sum of all other components, which
+            // can't really happen in practice.
+            // So we'll just adopt the rule "if total==0, don't write any details".
+            // If negative values are used for other things in the future, this can be revisited.
+            return false;
+        }
+        if (proto == null) {
+            // We're just asked whether there is data, not to actually write it. And there is.
+            return true;
+        }
+
+        final long token = proto.start(fieldId);
+        proto.write(
+                BatteryUsageStatsAtomsProto.BatteryConsumerData.TOTAL_CONSUMED_POWER_DECI_COULOMBS,
+                totalConsumedPowerDeciCoulombs);
+        mPowerComponents.writeStatsProto(proto);
+        proto.end(token);
+
+        return true;
+    }
+
+    /** Converts charge from milliamp hours (mAh) to decicoulombs (dC). */
+    static long convertMahToDeciCoulombs(double powerMah) {
+        return (long) (powerMah * (10 * 3600 / 1000) + 0.5);
+    }
+
     protected abstract static class BaseBuilder<T extends BaseBuilder<?>> {
         final PowerComponents.Builder mPowerComponentsBuilder;
 
diff --git a/core/java/android/os/BatteryUsageStats.java b/core/java/android/os/BatteryUsageStats.java
index 32256677c7f1..bd319925bbca 100644
--- a/core/java/android/os/BatteryUsageStats.java
+++ b/core/java/android/os/BatteryUsageStats.java
@@ -20,6 +20,7 @@ import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.util.Range;
 import android.util.SparseArray;
+import android.util.proto.ProtoOutputStream;
 
 import com.android.internal.os.BatteryStatsHistory;
 import com.android.internal.os.BatteryStatsHistoryIterator;
@@ -158,7 +159,8 @@ public final class BatteryUsageStats implements Parcelable {
 
     /**
      * Portion of battery charge drained since BatteryStats reset (e.g. due to being fully
-     * charged), as percentage of the full charge in the range [0:100]
+     * charged), as percentage of the full charge in the range [0:100]. May exceed 100 if
+     * the device repeatedly charged and discharged prior to the reset.
      */
     public int getDischargePercentage() {
         return mDischargePercentage;
@@ -335,6 +337,70 @@ public final class BatteryUsageStats implements Parcelable {
         }
     };
 
+    /** Returns a proto (as used for atoms.proto) corresponding to this BatteryUsageStats. */
+    public byte[] getStatsProto(long sessionEndTimestampMs) {
+
+        final long sessionStartMillis = getStatsStartTimestamp();
+        // TODO(b/187223764): Use the getStatsEndTimestamp() instead, once that is added.
+        final long sessionEndMillis = sessionEndTimestampMs;
+        final long sessionDurationMillis = sessionEndTimestampMs - getStatsStartTimestamp();
+
+        final BatteryConsumer deviceBatteryConsumer = getAggregateBatteryConsumer(
+                AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE);
+
+        final int sessionDischargePercentage = getDischargePercentage();
+
+        final ProtoOutputStream proto = new ProtoOutputStream();
+        proto.write(BatteryUsageStatsAtomsProto.SESSION_START_MILLIS, sessionStartMillis);
+        proto.write(BatteryUsageStatsAtomsProto.SESSION_END_MILLIS, sessionEndMillis);
+        proto.write(BatteryUsageStatsAtomsProto.SESSION_DURATION_MILLIS, sessionDurationMillis);
+        deviceBatteryConsumer.writeStatsProto(proto,
+                BatteryUsageStatsAtomsProto.DEVICE_BATTERY_CONSUMER);
+        writeUidBatteryConsumersProto(proto);
+        proto.write(BatteryUsageStatsAtomsProto.SESSION_DISCHARGE_PERCENTAGE,
+                sessionDischargePercentage);
+        return proto.getBytes();
+    }
+
+    /**
+     * Writes the UidBatteryConsumers data, held by this BatteryUsageStats, to the proto (as used
+     * for atoms.proto).
+     */
+    private void writeUidBatteryConsumersProto(ProtoOutputStream proto) {
+        final List<UidBatteryConsumer> consumers = getUidBatteryConsumers();
+
+        // TODO: Sort the list by power consumption. If during the for, proto.getRawSize() > 45kb,
+        //       truncate the remainder of the list.
+        final int size = consumers.size();
+        for (int i = 0; i < size; i++) {
+            final UidBatteryConsumer consumer = consumers.get(i);
+
+            final long fgMs = consumer.getTimeInStateMs(UidBatteryConsumer.STATE_FOREGROUND);
+            final long bgMs = consumer.getTimeInStateMs(UidBatteryConsumer.STATE_BACKGROUND);
+            final boolean hasBaseData = consumer.hasStatsProtoData();
+
+            if (fgMs == 0 && bgMs == 0 && !hasBaseData) {
+                continue;
+            }
+
+            final long token = proto.start(BatteryUsageStatsAtomsProto.UID_BATTERY_CONSUMERS);
+            proto.write(
+                    BatteryUsageStatsAtomsProto.UidBatteryConsumer.UID,
+                    consumer.getUid());
+            if (hasBaseData) {
+                consumer.writeStatsProto(proto,
+                        BatteryUsageStatsAtomsProto.UidBatteryConsumer.BATTERY_CONSUMER_DATA);
+            }
+            proto.write(
+                    BatteryUsageStatsAtomsProto.UidBatteryConsumer.TIME_IN_FOREGROUND_MILLIS,
+                    fgMs);
+            proto.write(
+                    BatteryUsageStatsAtomsProto.UidBatteryConsumer.TIME_IN_BACKGROUND_MILLIS,
+                    bgMs);
+            proto.end(token);
+        }
+    }
+
     /**
      * Prints the stats in a human-readable format.
      */
diff --git a/core/java/android/os/PowerComponents.java b/core/java/android/os/PowerComponents.java
index 964f1b681866..a90ed20d54fc 100644
--- a/core/java/android/os/PowerComponents.java
+++ b/core/java/android/os/PowerComponents.java
@@ -15,7 +15,11 @@
  */
 package android.os;
 
+import static android.os.BatteryConsumer.convertMahToDeciCoulombs;
+
 import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.util.proto.ProtoOutputStream;
 
 import com.android.internal.os.PowerCalculator;
 
@@ -237,6 +241,59 @@ class PowerComponents {
         }
     }
 
+    /** Returns whether there are any atoms.proto POWER_COMPONENTS data to write to a proto. */
+    boolean hasStatsProtoData() {
+        return writeStatsProtoImpl(null);
+    }
+
+    /** Writes all atoms.proto POWER_COMPONENTS for this PowerComponents to the given proto. */
+    void writeStatsProto(@NonNull ProtoOutputStream proto) {
+        writeStatsProtoImpl(proto);
+    }
+
+    /**
+     * Returns whether there are any atoms.proto POWER_COMPONENTS data to write to a proto,
+     * and writes it to the given proto if it is non-null.
+     */
+    private boolean writeStatsProtoImpl(@Nullable ProtoOutputStream proto) {
+        boolean interestingData = false;
+
+        for (int idx = 0; idx < mPowerComponentsMah.length; idx++) {
+            final int componentId = idx < BatteryConsumer.POWER_COMPONENT_COUNT ?
+                    idx : idx - CUSTOM_POWER_COMPONENT_OFFSET;
+            final long powerDeciCoulombs = convertMahToDeciCoulombs(mPowerComponentsMah[idx]);
+            final long durationMs = mUsageDurationsMs[idx];
+
+            if (powerDeciCoulombs == 0 && durationMs == 0) {
+                // No interesting data. Make sure not to even write the COMPONENT int.
+                continue;
+            }
+
+            interestingData = true;
+            if (proto == null) {
+                // We're just asked whether there is data, not to actually write it. And there is.
+                return true;
+            }
+
+            final long token =
+                    proto.start(BatteryUsageStatsAtomsProto.BatteryConsumerData.POWER_COMPONENTS);
+            proto.write(
+                    BatteryUsageStatsAtomsProto.BatteryConsumerData.PowerComponentUsage
+                            .COMPONENT,
+                    componentId);
+            proto.write(
+                    BatteryUsageStatsAtomsProto.BatteryConsumerData.PowerComponentUsage
+                            .POWER_DECI_COULOMBS,
+                    powerDeciCoulombs);
+            proto.write(
+                    BatteryUsageStatsAtomsProto.BatteryConsumerData.PowerComponentUsage
+                            .DURATION_MILLIS,
+                    durationMs);
+            proto.end(token);
+        }
+        return interestingData;
+    }
+
     /**
      * Builder for PowerComponents.
      */
diff --git a/core/proto/android/os/batteryusagestats.proto b/core/proto/android/os/batteryusagestats.proto
new file mode 100644
index 000000000000..bcce784ae0be
--- /dev/null
+++ b/core/proto/android/os/batteryusagestats.proto
@@ -0,0 +1,81 @@
+/*
+ * 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.
+ */
+
+syntax = "proto2";
+package android.os;
+
+option java_multiple_files = true;
+
+import "frameworks/proto_logging/stats/enums/os/enums.proto";
+
+// This message is used for statsd logging and should be kept in sync with
+// frameworks/proto_logging/stats/atoms.proto
+/**
+ * Represents a device's BatteryUsageStats, with power usage information about the device
+ * and each app.
+ */
+message BatteryUsageStatsAtomsProto {
+
+    // The session start timestamp in UTC milliseconds since January 1, 1970, per Date#getTime().
+    // All data is no older than this time.
+    optional int64 session_start_millis = 1;
+
+    // The session end timestamp in UTC milliseconds since January 1, 1970, per Date#getTime().
+    // All data is no more recent than this time.
+    optional int64 session_end_millis = 2;
+
+    // Length that the reported data covered. This usually will be equal to the entire session,
+    // session_end_millis - session_start_millis, but may not be if some data during this time frame
+    // is missing.
+    optional int64 session_duration_millis = 3;
+
+    // Represents usage of a consumer, storing all of its power component usage.
+    message BatteryConsumerData {
+        // Total power consumed by this BatteryConsumer (including all of its PowerComponents).
+        // May not equal the sum of the PowerComponentUsage due to under- or over-estimations.
+        // Multiply by 1/36 to obtain mAh.
+        optional int64 total_consumed_power_deci_coulombs = 1;
+
+        // Represents power and time usage of a particular power component.
+        message PowerComponentUsage {
+            // Holds android.os.PowerComponentEnum, or custom component value between 1000 and 9999.
+            // Evidently, if one attempts to write an int to an enum field that is out of range, it
+            // is treated as 0, so we must make this an int32.
+            optional int32 component = 1;
+
+            // Power consumed by this component. Multiply by 1/36 to obtain mAh.
+            optional int64 power_deci_coulombs = 2;
+
+            optional int64 duration_millis = 3;
+        }
+        repeated PowerComponentUsage power_components = 2;
+    }
+
+    // Total power usage for the device during this session.
+    optional BatteryConsumerData device_battery_consumer = 4;
+
+    // Power usage by a uid during this session.
+    message UidBatteryConsumer {
+        optional int32 uid = 1;
+        optional BatteryConsumerData battery_consumer_data = 2;
+        optional int64 time_in_foreground_millis = 3;
+        optional int64 time_in_background_millis = 4;
+    }
+    repeated UidBatteryConsumer uid_battery_consumers = 5;
+
+    // Sum of all discharge percentage point drops during the reported session.
+    optional int32 session_discharge_percentage = 6;
+}
\ No newline at end of file
diff --git a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp
new file mode 100644
index 000000000000..36cb55407d8d
--- /dev/null
+++ b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp
@@ -0,0 +1,28 @@
+package {
+    // See: http://go/android-license-faq
+    // A large-scale-change added 'default_applicable_licenses' to import
+    // all of the 'license_kinds' from "frameworks_base_license"
+    // to get the below license kinds:
+    //   SPDX-license-identifier-Apache-2.0
+    default_applicable_licenses: ["frameworks_base_license"],
+}
+
+android_test {
+    name: "BatteryUsageStatsProtoTests",
+    srcs: ["src/**/*.java"],
+
+    static_libs: [
+        "androidx.test.rules",
+        "junit",
+        "platform-test-annotations",
+        "platformprotosnano",
+        "statsdprotolite",
+    ],
+
+    libs: ["android.test.runner"],
+
+    platform_apis: true,
+    certificate: "platform",
+
+    test_suites: ["device-tests"],
+}
diff --git a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/AndroidManifest.xml b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/AndroidManifest.xml
new file mode 100644
index 000000000000..9128dca2080b
--- /dev/null
+++ b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/AndroidManifest.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ 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.
+  -->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+          package="com.android.frameworks.core.batteryusagestatsprototests">
+
+    <uses-permission android:name="android.permission.BATTERY_STATS"/>
+
+    <instrumentation
+        android:name="androidx.test.runner.AndroidJUnitRunner"
+        android:targetPackage="com.android.frameworks.core.batteryusagestatsprototests"
+        android:label="BatteryUsageStats Proto Tests" />
+
+</manifest>
diff --git a/core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java
new file mode 100644
index 000000000000..bee0a0bf1fd6
--- /dev/null
+++ b/core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java
@@ -0,0 +1,229 @@
+/*
+ * 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 android.os.BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
+
+import android.os.BatteryConsumer;
+import android.os.BatteryUsageStats;
+import android.os.nano.BatteryUsageStatsAtomsProto;
+import android.os.nano.BatteryUsageStatsAtomsProto.BatteryConsumerData.PowerComponentUsage;
+
+import androidx.test.filters.SmallTest;
+
+import com.google.protobuf.nano.InvalidProtocolBufferNanoException;
+
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.List;
+
+
+
+@SmallTest
+public class BatteryUsageStatsPulledTest {
+
+    private static final int UID_0 = 1000;
+    private static final int UID_1 = 2000;
+    private static final int UID_2 = 3000;
+    private static final int UID_3 = 4000;
+
+    @Test
+    public void testGetStatsProto() {
+        final long sessionEndTimestampMs = 1050;
+        final BatteryUsageStats bus = buildBatteryUsageStats();
+        final byte[] bytes = bus.getStatsProto(sessionEndTimestampMs);
+        BatteryUsageStatsAtomsProto proto;
+        try {
+            proto = BatteryUsageStatsAtomsProto.parseFrom(bytes);
+        } catch (InvalidProtocolBufferNanoException e) {
+            fail("Invalid proto: " + e);
+            return;
+        }
+
+        assertEquals(bus.getStatsStartTimestamp(), proto.sessionStartMillis);
+        assertEquals(sessionEndTimestampMs, proto.sessionEndMillis);
+        assertEquals(
+                sessionEndTimestampMs - bus.getStatsStartTimestamp(),
+                proto.sessionDurationMillis);
+        assertEquals(bus.getDischargePercentage(), proto.sessionDischargePercentage);
+
+        assertEquals(3, proto.deviceBatteryConsumer.powerComponents.length); // Only 3 are non-empty
+        assertSameBatteryConsumer("For deviceBatteryConsumer",
+                bus.getAggregateBatteryConsumer(AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE),
+                proto.deviceBatteryConsumer);
+
+        // Now for the UidBatteryConsumers.
+        final List<android.os.UidBatteryConsumer> uidConsumers = bus.getUidBatteryConsumers();
+        uidConsumers.sort((a, b) -> a.getUid() - b.getUid());
+
+        final BatteryUsageStatsAtomsProto.UidBatteryConsumer[] uidConsumersProto
+                = proto.uidBatteryConsumers;
+        Arrays.sort(uidConsumersProto, (a, b) -> a.uid - b.uid);
+
+        // UID_0 - After sorting, UID_0 should be in position 0 for both data structures
+        assertEquals(UID_0, bus.getUidBatteryConsumers().get(0).getUid());
+        assertEquals(UID_0, proto.uidBatteryConsumers[0].uid);
+        assertSameUidBatteryConsumer(
+                bus.getUidBatteryConsumers().get(0),
+                proto.uidBatteryConsumers[0],
+                false);
+
+        // UID_1 - After sorting, UID_1 should be in position 1 for both data structures
+        assertEquals(UID_1, bus.getUidBatteryConsumers().get(1).getUid());
+        assertEquals(UID_1, proto.uidBatteryConsumers[1].uid);
+        assertSameUidBatteryConsumer(
+                bus.getUidBatteryConsumers().get(1),
+                proto.uidBatteryConsumers[1],
+                true);
+
+        // UID_2 - After sorting, UID_2 should be in position 2 for both data structures
+        assertEquals(UID_2, bus.getUidBatteryConsumers().get(2).getUid());
+        assertEquals(UID_2, proto.uidBatteryConsumers[2].uid);
+        assertSameUidBatteryConsumer(
+                bus.getUidBatteryConsumers().get(2),
+                proto.uidBatteryConsumers[2],
+                false);
+
+        // UID_3 - Should be none, since no interesting data (done last for debugging convenience).
+        assertEquals(3, proto.uidBatteryConsumers.length);
+    }
+
+    private void assertSameBatteryConsumer(String message, BatteryConsumer consumer,
+            android.os.nano.BatteryUsageStatsAtomsProto.BatteryConsumerData consumerProto) {
+        assertNotNull(message, consumerProto);
+        assertEquals(
+                convertMahToDc(consumer.getConsumedPower()),
+                consumerProto.totalConsumedPowerDeciCoulombs);
+
+        for (PowerComponentUsage componentProto : consumerProto.powerComponents) {
+            final int componentId = componentProto.component;
+            if (componentId < BatteryConsumer.POWER_COMPONENT_COUNT) {
+                assertEquals(message + " for component " + componentId,
+                        convertMahToDc(consumer.getConsumedPower(componentId)),
+                        componentProto.powerDeciCoulombs);
+                assertEquals(message + " for component " + componentId,
+                        consumer.getUsageDurationMillis(componentId),
+                        componentProto.durationMillis);
+            } else {
+                assertEquals(message + " for custom component " + componentId,
+                        convertMahToDc(consumer.getConsumedPowerForCustomComponent(componentId)),
+                        componentProto.powerDeciCoulombs);
+                assertEquals(message + " for custom component " + componentId,
+                        consumer.getUsageDurationForCustomComponentMillis(componentId),
+                        componentProto.durationMillis);
+            }
+        }
+    }
+
+    private void assertSameUidBatteryConsumer(
+            android.os.UidBatteryConsumer uidConsumer,
+            BatteryUsageStatsAtomsProto.UidBatteryConsumer uidConsumerProto,
+            boolean expectNullBatteryConsumerData) {
+
+        final int uid = uidConsumerProto.uid;
+        assertEquals("Uid consumers had mismatched uids", uid, uidConsumer.getUid());
+
+        assertEquals("For uid " + uid,
+                uidConsumer.getTimeInStateMs(android.os.UidBatteryConsumer.STATE_FOREGROUND),
+                uidConsumerProto.timeInForegroundMillis);
+        assertEquals("For uid " + uid,
+                uidConsumer.getTimeInStateMs(android.os.UidBatteryConsumer.STATE_BACKGROUND),
+                uidConsumerProto.timeInBackgroundMillis);
+        if (expectNullBatteryConsumerData) {
+            assertNull("For uid " + uid, uidConsumerProto.batteryConsumerData);
+        } else {
+            assertSameBatteryConsumer("For uid " + uid,
+                    uidConsumer,
+                    uidConsumerProto.batteryConsumerData);
+        }
+    }
+
+    /** Converts charge from milliamp hours (mAh) to decicoulombs (dC). */
+    private long convertMahToDc(double powerMah) {
+        return (long) (powerMah * 36 + 0.5);
+    }
+
+    private BatteryUsageStats buildBatteryUsageStats() {
+        final BatteryStatsImpl batteryStats = new BatteryStatsImpl();
+        final BatteryStatsImpl.Uid batteryStatsUid0 = batteryStats.getUidStatsLocked(UID_0);
+        final BatteryStatsImpl.Uid batteryStatsUid1 = batteryStats.getUidStatsLocked(UID_1);
+        final BatteryStatsImpl.Uid batteryStatsUid2 = batteryStats.getUidStatsLocked(UID_2);
+        final BatteryStatsImpl.Uid batteryStatsUid3 = batteryStats.getUidStatsLocked(UID_3);
+
+        final BatteryUsageStats.Builder builder =
+                new BatteryUsageStats.Builder(new String[]{"CustomConsumer1", "CustomConsumer2"})
+                        .setDischargePercentage(20)
+                        .setDischargedPowerRange(1000, 2000)
+                        .setStatsStartTimestamp(1000);
+        builder.getOrCreateUidBatteryConsumerBuilder(batteryStatsUid0)
+                .setPackageWithHighestDrain("myPackage0")
+                .setTimeInStateMs(android.os.UidBatteryConsumer.STATE_FOREGROUND, 1000)
+                .setTimeInStateMs(android.os.UidBatteryConsumer.STATE_BACKGROUND, 2000)
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_SCREEN, 300)
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_CPU, 400)
+                .setConsumedPowerForCustomComponent(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 450)
+                .setConsumedPowerForCustomComponent(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID + 1, 500)
+                .setUsageDurationMillis(
+                        BatteryConsumer.POWER_COMPONENT_CPU, 600)
+                .setUsageDurationForCustomComponentMillis(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID + 1, 800);
+
+        builder.getOrCreateUidBatteryConsumerBuilder(batteryStatsUid1)
+                .setPackageWithHighestDrain("myPackage1")
+                .setTimeInStateMs(android.os.UidBatteryConsumer.STATE_FOREGROUND, 1234);
+
+        builder.getOrCreateUidBatteryConsumerBuilder(batteryStatsUid2)
+                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN,
+                        766);
+
+        builder.getOrCreateUidBatteryConsumerBuilder(batteryStatsUid3);
+
+        builder.getAggregateBatteryConsumerBuilder(AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
+                .setConsumedPower(30000)
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_CPU, 20100)
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_AUDIO, 0) // Empty
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_CAMERA, 20150)
+                .setConsumedPowerForCustomComponent(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 20200)
+                .setUsageDurationMillis(
+                        BatteryConsumer.POWER_COMPONENT_CPU, 20300)
+                .setUsageDurationForCustomComponentMillis(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 20400);
+
+        // Not used; just to make sure extraneous data doesn't mess things up.
+        builder.getAggregateBatteryConsumerBuilder(
+                BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_ALL_APPS)
+                .setConsumedPower(
+                        BatteryConsumer.POWER_COMPONENT_CPU, 10100)
+                .setConsumedPowerForCustomComponent(
+                        BatteryConsumer.FIRST_CUSTOM_POWER_COMPONENT_ID, 10200);
+
+        return builder.build();
+    }
+}
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 5937a18673d9..5da8f30c4dfb 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -20,6 +20,7 @@ import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_SUSPENDED;
 import static android.os.BatteryStats.POWER_DATA_UNAVAILABLE;
 
 import android.annotation.NonNull;
+import android.app.StatsManager;
 import android.bluetooth.BluetoothActivityEnergyInfo;
 import android.content.ContentResolver;
 import android.content.Context;
@@ -69,9 +70,11 @@ import android.telephony.ModemActivityInfo;
 import android.telephony.SignalStrength;
 import android.telephony.TelephonyManager;
 import android.util.Slog;
+import android.util.StatsEvent;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.app.IBatteryStats;
+import com.android.internal.os.BackgroundThread;
 import com.android.internal.os.BatteryStatsHelper;
 import com.android.internal.os.BatteryStatsImpl;
 import com.android.internal.os.BatteryUsageStatsProvider;
@@ -380,6 +383,8 @@ public final class BatteryStatsService extends IBatteryStats.Stub
 
         final DataConnectionStats dataConnectionStats = new DataConnectionStats(mContext, mHandler);
         dataConnectionStats.startMonitoring();
+
+        registerStatsCallbacks();
     }
 
     private final class LocalService extends BatteryStatsInternal {
@@ -672,7 +677,7 @@ public final class BatteryStatsService extends IBatteryStats.Stub
      * and per-UID basis.
      */
     public List<BatteryUsageStats> getBatteryUsageStats(List<BatteryUsageStatsQuery> queries) {
-        mContext.enforceCallingPermission(
+        mContext.enforceCallingOrSelfPermission(
                 android.Manifest.permission.BATTERY_STATS, null);
         awaitCompletion();
 
@@ -734,6 +739,48 @@ public final class BatteryStatsService extends IBatteryStats.Stub
         }
     }
 
+    /** Register callbacks for statsd pulled atoms. */
+    private void registerStatsCallbacks() {
+        final StatsManager statsManager = mContext.getSystemService(StatsManager.class);
+        final StatsPullAtomCallbackImpl pullAtomCallback = new StatsPullAtomCallbackImpl();
+
+        statsManager.setPullAtomCallback(
+                FrameworkStatsLog.BATTERY_USAGE_STATS_SINCE_RESET,
+                null, // use default PullAtomMetadata values
+                BackgroundThread.getExecutor(), pullAtomCallback);
+        statsManager.setPullAtomCallback(
+                FrameworkStatsLog.BATTERY_USAGE_STATS_SINCE_RESET_USING_POWER_PROFILE_MODEL,
+                null, // use default PullAtomMetadata values
+                BackgroundThread.getExecutor(), pullAtomCallback);
+    }
+
+    /** StatsPullAtomCallback for pulling BatteryUsageStats data. */
+    private class StatsPullAtomCallbackImpl implements StatsManager.StatsPullAtomCallback {
+        @Override
+        public int onPullAtom(int atomTag, List<StatsEvent> data) {
+            final BatteryUsageStats bus;
+            switch (atomTag) {
+                case FrameworkStatsLog.BATTERY_USAGE_STATS_SINCE_RESET:
+                    bus = getBatteryUsageStats(List.of(BatteryUsageStatsQuery.DEFAULT)).get(0);
+                    break;
+                case FrameworkStatsLog.BATTERY_USAGE_STATS_SINCE_RESET_USING_POWER_PROFILE_MODEL:
+                    final BatteryUsageStatsQuery powerProfileQuery =
+                            new BatteryUsageStatsQuery.Builder().powerProfileModeledOnly().build();
+                    bus = getBatteryUsageStats(List.of(powerProfileQuery)).get(0);
+                    break;
+                default:
+                    throw new UnsupportedOperationException("Unknown tagId=" + atomTag);
+            }
+            // TODO(b/187223764): busTime won't be needed once end_session is a field in BUS.
+            final long busTime = System.currentTimeMillis();
+            final byte[] statsProto = bus.getStatsProto(busTime);
+
+            data.add(FrameworkStatsLog.buildStatsEvent(atomTag, statsProto));
+
+            return StatsManager.PULL_SUCCESS;
+        }
+    }
+
     public boolean isCharging() {
         synchronized (mStats) {
             return mStats.isCharging();