Merge "Limit the number of UidBatteryConsumers included in the BatteryUsageStats atom" into sc-dev

core/java/android/os/BatteryUsageStats.java

@@ -75,6 +75,8 @@ public final class BatteryUsageStats implements Parcelable {
 
     public static final int AGGREGATE_BATTERY_CONSUMER_SCOPE_COUNT = 2;
 
+    private static final int STATSD_PULL_ATOM_MAX_BYTES = 45000;
+
     private final int mDischargePercentage;
     private final double mBatteryCapacityMah;
     private final long mStatsStartTimestampMs;
@@ -369,30 +371,60 @@ public final class BatteryUsageStats implements Parcelable {
 
     /** Returns a proto (as used for atoms.proto) corresponding to this BatteryUsageStats. */
     public byte[] getStatsProto() {
+        // ProtoOutputStream.getRawSize() returns the buffer size before compaction.
+        // BatteryUsageStats contains a lot of integers, so compaction of integers to
+        // varint reduces the size of the proto buffer by as much as 50%.
+        int maxRawSize = (int) (STATSD_PULL_ATOM_MAX_BYTES * 1.75);
+        // Limit the number of attempts in order to prevent an infinite loop
+        for (int i = 0; i < 3; i++) {
+            final ProtoOutputStream proto = new ProtoOutputStream();
+            writeStatsProto(proto, maxRawSize);
+
+            final int rawSize = proto.getRawSize();
+            final byte[] protoOutput = proto.getBytes();
+
+            if (protoOutput.length <= STATSD_PULL_ATOM_MAX_BYTES) {
+                return protoOutput;
+            }
+
+            // Adjust maxRawSize proportionately and try again.
+            maxRawSize =
+                    (int) ((long) STATSD_PULL_ATOM_MAX_BYTES * rawSize / protoOutput.length - 1024);
+        }
+
+        // Fallback: if we have failed to generate a proto smaller than STATSD_PULL_ATOM_MAX_BYTES,
+        // just generate a proto with the _rawSize_ of STATSD_PULL_ATOM_MAX_BYTES, which is
+        // guaranteed to produce a compacted proto (significantly) smaller than
+        // STATSD_PULL_ATOM_MAX_BYTES.
+        final ProtoOutputStream proto = new ProtoOutputStream();
+        writeStatsProto(proto, STATSD_PULL_ATOM_MAX_BYTES);
+        return proto.getBytes();
+    }
+
+    @NonNull
+    private void writeStatsProto(ProtoOutputStream proto, int maxRawSize) {
         final BatteryConsumer deviceBatteryConsumer = getAggregateBatteryConsumer(
                 AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE);
 
-        final ProtoOutputStream proto = new ProtoOutputStream();
         proto.write(BatteryUsageStatsAtomsProto.SESSION_START_MILLIS, getStatsStartTimestamp());
         proto.write(BatteryUsageStatsAtomsProto.SESSION_END_MILLIS, getStatsEndTimestamp());
         proto.write(BatteryUsageStatsAtomsProto.SESSION_DURATION_MILLIS, getStatsDuration());
-        deviceBatteryConsumer.writeStatsProto(proto,
-                BatteryUsageStatsAtomsProto.DEVICE_BATTERY_CONSUMER);
-        writeUidBatteryConsumersProto(proto);
         proto.write(BatteryUsageStatsAtomsProto.SESSION_DISCHARGE_PERCENTAGE,
                 getDischargePercentage());
-        return proto.getBytes();
+        deviceBatteryConsumer.writeStatsProto(proto,
+                BatteryUsageStatsAtomsProto.DEVICE_BATTERY_CONSUMER);
+        writeUidBatteryConsumersProto(proto, maxRawSize);
     }
 
     /**
      * Writes the UidBatteryConsumers data, held by this BatteryUsageStats, to the proto (as used
      * for atoms.proto).
      */
-    private void writeUidBatteryConsumersProto(ProtoOutputStream proto) {
+    private void writeUidBatteryConsumersProto(ProtoOutputStream proto, int maxRawSize) {
         final List<UidBatteryConsumer> consumers = getUidBatteryConsumers();
+        // Order consumers by descending weight (a combination of consumed power and usage time)
+        consumers.sort(Comparator.comparingDouble(this::getUidBatteryConsumerWeight).reversed());
 
-        // TODO(b/189225426): 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);
@@ -420,9 +452,35 @@ public final class BatteryUsageStats implements Parcelable {
                     BatteryUsageStatsAtomsProto.UidBatteryConsumer.TIME_IN_BACKGROUND_MILLIS,
                     bgMs);
             proto.end(token);
+
+            if (proto.getRawSize() >= maxRawSize) {
+                break;
+            }
         }
     }
 
+    private static final double WEIGHT_CONSUMED_POWER = 1;
+    // Weight one hour in foreground the same as 100 mAh of power drain
+    private static final double WEIGHT_FOREGROUND_STATE = 100.0 / (1 * 60 * 60 * 1000);
+    // Weight one hour in background the same as 300 mAh of power drain
+    private static final double WEIGHT_BACKGROUND_STATE = 300.0 / (1 * 60 * 60 * 1000);
+
+    /**
+     * Computes the weight associated with a UidBatteryConsumer, which is used for sorting.
+     * We want applications with the largest consumed power as well as applications
+     * with the highest usage time to be included in the statsd atom.
+     */
+    private double getUidBatteryConsumerWeight(UidBatteryConsumer uidBatteryConsumer) {
+        final double consumedPower = uidBatteryConsumer.getConsumedPower();
+        final long timeInForeground =
+                uidBatteryConsumer.getTimeInStateMs(UidBatteryConsumer.STATE_FOREGROUND);
+        final long timeInBackground =
+                uidBatteryConsumer.getTimeInStateMs(UidBatteryConsumer.STATE_BACKGROUND);
+        return consumedPower * WEIGHT_CONSUMED_POWER
+                + timeInForeground * WEIGHT_FOREGROUND_STATE
+                + timeInBackground * WEIGHT_BACKGROUND_STATE;
+    }
+
     /**
      * Prints the stats in a human-readable format.
      */

core/tests/batterystatstests/BatteryUsageStatsProtoTests/Android.bp

@@ -14,9 +14,11 @@ android_test {
     static_libs: [
         "androidx.test.rules",
         "junit",
+        "mockito-target-minus-junit4",
         "platform-test-annotations",
         "platformprotosnano",
         "statsdprotolite",
+        "truth-prebuilt",
     ],
 
     libs: ["android.test.runner"],

core/tests/batterystatstests/BatteryUsageStatsProtoTests/src/com/android/internal/os/BatteryUsageStatsPulledTest.java

@@ -17,10 +17,14 @@ package com.android.internal.os;
 
 import static android.os.BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE;
 
+import static com.google.common.truth.Truth.assertThat;
+
 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 static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import android.os.BatteryConsumer;
 import android.os.BatteryUsageStats;
@@ -37,7 +41,6 @@ import java.util.Arrays;
 import java.util.List;
 
 
-
 @SmallTest
 public class BatteryUsageStatsPulledTest {
 
@@ -225,4 +228,75 @@ public class BatteryUsageStatsPulledTest {
 
         return builder.build();
     }
+
+    @Test
+    public void testLargeAtomTruncated() {
+        final BatteryUsageStats.Builder builder =
+                new BatteryUsageStats.Builder(new String[0]);
+        // If not truncated, this BatteryUsageStats object would generate a proto buffer
+        // larger than 70 Kb
+        for (int i = 0; i < 20000; i++) {
+            BatteryStatsImpl.Uid mockUid = mock(BatteryStatsImpl.Uid.class);
+            when(mockUid.getUid()).thenReturn(i);
+            builder.getOrCreateUidBatteryConsumerBuilder(mockUid)
+                    .setTimeInStateMs(android.os.UidBatteryConsumer.STATE_FOREGROUND, 1 * 60 * 1000)
+                    .setTimeInStateMs(android.os.UidBatteryConsumer.STATE_BACKGROUND, 2 * 60 * 1000)
+                    .setConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN, 30)
+                    .setConsumedPower(BatteryConsumer.POWER_COMPONENT_CPU, 40);
+        }
+
+        // Add a UID with much larger battery footprint
+        final int largeConsumerUid = 20001;
+        BatteryStatsImpl.Uid largeConsumerMockUid = mock(BatteryStatsImpl.Uid.class);
+        when(largeConsumerMockUid.getUid()).thenReturn(largeConsumerUid);
+        builder.getOrCreateUidBatteryConsumerBuilder(largeConsumerMockUid)
+                .setTimeInStateMs(android.os.UidBatteryConsumer.STATE_FOREGROUND, 10 * 60 * 1000)
+                .setTimeInStateMs(android.os.UidBatteryConsumer.STATE_BACKGROUND, 20 * 60 * 1000)
+                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN, 300)
+                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_CPU, 400);
+
+        // Add a UID with much larger usage duration
+        final int highUsageUid = 20002;
+        BatteryStatsImpl.Uid highUsageMockUid = mock(BatteryStatsImpl.Uid.class);
+        when(highUsageMockUid.getUid()).thenReturn(highUsageUid);
+        builder.getOrCreateUidBatteryConsumerBuilder(highUsageMockUid)
+                .setTimeInStateMs(android.os.UidBatteryConsumer.STATE_FOREGROUND, 60 * 60 * 1000)
+                .setTimeInStateMs(android.os.UidBatteryConsumer.STATE_BACKGROUND, 120 * 60 * 1000)
+                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN, 3)
+                .setConsumedPower(BatteryConsumer.POWER_COMPONENT_CPU, 4);
+
+        BatteryUsageStats batteryUsageStats = builder.build();
+        final byte[] bytes = batteryUsageStats.getStatsProto();
+        assertThat(bytes.length).isGreaterThan(40000);
+        assertThat(bytes.length).isLessThan(50000);
+
+        BatteryUsageStatsAtomsProto proto;
+        try {
+            proto = BatteryUsageStatsAtomsProto.parseFrom(bytes);
+        } catch (InvalidProtocolBufferNanoException e) {
+            fail("Invalid proto: " + e);
+            return;
+        }
+
+        boolean largeConsumerIncluded = false;
+        boolean highUsageAppIncluded = false;
+        for (int i = 0; i < proto.uidBatteryConsumers.length; i++) {
+            if (proto.uidBatteryConsumers[i].uid == largeConsumerUid) {
+                largeConsumerIncluded = true;
+                BatteryUsageStatsAtomsProto.BatteryConsumerData consumerData =
+                        proto.uidBatteryConsumers[i].batteryConsumerData;
+                assertThat(consumerData.totalConsumedPowerDeciCoulombs / 36)
+                        .isEqualTo(300 + 400);
+            } else if (proto.uidBatteryConsumers[i].uid == highUsageUid) {
+                highUsageAppIncluded = true;
+                BatteryUsageStatsAtomsProto.BatteryConsumerData consumerData =
+                        proto.uidBatteryConsumers[i].batteryConsumerData;
+                assertThat(consumerData.totalConsumedPowerDeciCoulombs / 36)
+                        .isEqualTo(3 + 4);
+            }
+        }
+
+        assertThat(largeConsumerIncluded).isTrue();
+        assertThat(highUsageAppIncluded).isTrue();
+    }
 }