Reduce size of BatteryUsageStats Parcel

Bug: 191237968
Test: atest FrameworksCoreTests:BatteryUsageStatsTest

Change-Id: I89ec1b7f31935d65b170f180ff6cab4352e7ccbf

2 files changed, 73 insertions, 17 deletions

diff --git a/core/java/android/os/BatteryUsageStats.java b/core/java/android/os/BatteryUsageStats.java
index 77f8a87cdcd1..f48375246616 100644
--- a/core/java/android/os/BatteryUsageStats.java
+++ b/core/java/android/os/BatteryUsageStats.java
@@ -306,14 +306,23 @@ public final class BatteryUsageStats implements Parcelable {
                     AggregateBatteryConsumer.CREATOR.createFromParcel(source);
             mAggregateBatteryConsumers[i].setCustomPowerComponentNames(mCustomPowerComponentNames);
         }
-        int uidCount = source.readInt();
+
+        // UidBatteryConsumers are included as a blob to avoid a TransactionTooLargeException
+        final Parcel blob = Parcel.obtain();
+        final byte[] bytes = source.readBlob();
+        blob.unmarshall(bytes, 0, bytes.length);
+        blob.setDataPosition(0);
+
+        final int uidCount = blob.readInt();
         mUidBatteryConsumers = new ArrayList<>(uidCount);
         for (int i = 0; i < uidCount; i++) {
             final UidBatteryConsumer consumer =
-                    UidBatteryConsumer.CREATOR.createFromParcel(source);
+                    UidBatteryConsumer.CREATOR.createFromParcel(blob);
             consumer.setCustomPowerComponentNames(mCustomPowerComponentNames);
             mUidBatteryConsumers.add(consumer);
         }
+        blob.recycle();
+
         int userCount = source.readInt();
         mUserBatteryConsumers = new ArrayList<>(userCount);
         for (int i = 0; i < userCount; i++) {
@@ -323,14 +332,10 @@ public final class BatteryUsageStats implements Parcelable {
             mUserBatteryConsumers.add(consumer);
         }
         if (source.readBoolean()) {
-            mHistoryBuffer = Parcel.obtain();
-            mHistoryBuffer.setDataSize(0);
-            mHistoryBuffer.setDataPosition(0);
+            final byte[] historyBlob = source.readBlob();
 
-            int historyBufferSize = source.readInt();
-            int curPos = source.dataPosition();
-            mHistoryBuffer.appendFrom(source, curPos, historyBufferSize);
-            source.setDataPosition(curPos + historyBufferSize);
+            mHistoryBuffer = Parcel.obtain();
+            mHistoryBuffer.unmarshall(historyBlob, 0, historyBlob.length);
 
             int historyTagCount = source.readInt();
             mHistoryTagPool = new ArrayList<>(historyTagCount);
@@ -362,21 +367,26 @@ public final class BatteryUsageStats implements Parcelable {
         for (int i = 0; i < AGGREGATE_BATTERY_CONSUMER_SCOPE_COUNT; i++) {
             mAggregateBatteryConsumers[i].writeToParcel(dest, flags);
         }
-        dest.writeInt(mUidBatteryConsumers.size());
+
+        // UidBatteryConsumers are included as a blob, because each UidBatteryConsumer
+        // takes > 300 bytes, so a typical number of UIDs in the system, 300 would result
+        // in a 90 kB Parcel, which is not safe to pass in a binder transaction because
+        // of the possibility of TransactionTooLargeException
+        final Parcel blob = Parcel.obtain();
+        blob.writeInt(mUidBatteryConsumers.size());
         for (int i = mUidBatteryConsumers.size() - 1; i >= 0; i--) {
-            mUidBatteryConsumers.get(i).writeToParcel(dest, flags);
+            mUidBatteryConsumers.get(i).writeToParcel(blob, flags);
         }
+        dest.writeBlob(blob.marshall());
+        blob.recycle();
+
         dest.writeInt(mUserBatteryConsumers.size());
         for (int i = mUserBatteryConsumers.size() - 1; i >= 0; i--) {
             mUserBatteryConsumers.get(i).writeToParcel(dest, flags);
         }
         if (mHistoryBuffer != null) {
             dest.writeBoolean(true);
-
-            final int historyBufferSize = mHistoryBuffer.dataSize();
-            dest.writeInt(historyBufferSize);
-            dest.appendFrom(mHistoryBuffer, 0, historyBufferSize);
-
+            dest.writeBlob(mHistoryBuffer.marshall());
             dest.writeInt(mHistoryTagPool.size());
             for (int i = mHistoryTagPool.size() - 1; i >= 0; i--) {
                 final BatteryStats.HistoryTag tag = mHistoryTagPool.get(i);
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsTest.java
index 3e620c2bbec6..fedbf7a9868e 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsTest.java
@@ -24,6 +24,8 @@ import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertThrows;
 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;
@@ -47,7 +49,9 @@ import java.io.StringWriter;
 import java.nio.charset.StandardCharsets;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
@@ -64,13 +68,15 @@ public class BatteryUsageStatsTest {
     }
 
     @Test
-    public void testParcelability() {
+    public void testParcelability_smallNumberOfUids() {
         final BatteryUsageStats outBatteryUsageStats = buildBatteryUsageStats1(true).build();
         final Parcel outParcel = Parcel.obtain();
         outParcel.writeParcelable(outBatteryUsageStats, 0);
         final byte[] bytes = outParcel.marshall();
         outParcel.recycle();
 
+        assertThat(bytes.length).isLessThan(2000);
+
         final Parcel inParcel = Parcel.obtain();
         inParcel.unmarshall(bytes, 0, bytes.length);
         inParcel.setDataPosition(0);
@@ -80,6 +86,46 @@ public class BatteryUsageStatsTest {
         assertBatteryUsageStats1(inBatteryUsageStats, true);
     }
 
+    @Test
+    public void testParcelability_largeNumberOfUids() {
+        final BatteryUsageStats.Builder builder =
+                new BatteryUsageStats.Builder(new String[0]);
+
+        // Without the use of a blob, this BatteryUsageStats object would generate a Parcel
+        // larger than 64 Kb
+        final int uidCount = 200;
+        for (int i = 0; i < uidCount; i++) {
+            BatteryStatsImpl.Uid mockUid = mock(BatteryStatsImpl.Uid.class);
+            when(mockUid.getUid()).thenReturn(i);
+            builder.getOrCreateUidBatteryConsumerBuilder(mockUid)
+                    .setConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN, i * 100);
+        }
+
+        BatteryUsageStats outBatteryUsageStats = builder.build();
+
+        final Parcel parcel = Parcel.obtain();
+        parcel.writeParcelable(outBatteryUsageStats, 0);
+
+        assertThat(parcel.dataSize()).isLessThan(2000);
+
+        // This parcel cannot be marshaled because it contains a file descriptor.
+        // Assuming that parcel marshaling works fine, let's just rewind the parcel.
+        parcel.setDataPosition(0);
+
+        final BatteryUsageStats inBatteryUsageStats =
+                parcel.readParcelable(getClass().getClassLoader());
+        parcel.recycle();
+
+        assertThat(inBatteryUsageStats.getUidBatteryConsumers()).hasSize(uidCount);
+        final Map<Integer, UidBatteryConsumer> consumersByUid =
+                inBatteryUsageStats.getUidBatteryConsumers().stream().collect(
+                        Collectors.toMap(UidBatteryConsumer::getUid, c -> c));
+        for (int i = 0; i < uidCount; i++) {
+            final UidBatteryConsumer uidBatteryConsumer = consumersByUid.get(i);
+            assertThat(uidBatteryConsumer).isNotNull();
+            assertThat(uidBatteryConsumer.getConsumedPower()).isEqualTo(i * 100);
+        }
+    }
 
     @Test
     public void testDefaultSessionDuration() {