Merge "Add shard ID to statsd atom for procstats collection" into rvc-dev am: b0f26f32d8

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/11971040

Change-Id: I1bf803475b106b476b6a941af975b09ad05c32ab

2 files changed, 13 insertions, 3 deletions

diff --git a/cmds/statsd/src/atoms.proto b/cmds/statsd/src/atoms.proto
index a5f0ac97cebc..ca03343ddb80 100644
--- a/cmds/statsd/src/atoms.proto
+++ b/cmds/statsd/src/atoms.proto
@@ -6131,6 +6131,10 @@ message ProcessStatsAvailablePagesProto {
  */
 message ProcStats {
     optional ProcessStatsSectionProto proc_stats_section = 1;
+    // Data pulled from device into this is sometimes sharded across multiple atoms to work around
+    // a size limit. When this happens, this shard ID will contain an increasing 1-indexed integer
+    // with the number of this shard.
+    optional int32 shard_id = 2;
 }
 
 /**
diff --git a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
index 09fd33d5b4ed..dbdef2368c7c 100644
--- a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
+++ b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
@@ -2575,11 +2575,17 @@ public class StatsPullAtomService extends SystemService {
                     lastHighWaterMark, section, true, statsFiles, procStats);
             procStats.dumpAggregatedProtoForStatsd(protoStreams, MAX_PROCSTATS_RAW_SHARD_SIZE);
 
-            for (ProtoOutputStream proto : protoStreams) {
-                if (proto.getBytes().length > 0) {
+            for (int i = 0; i < protoStreams.length; i++) {
+                byte[] bytes = protoStreams[i].getBytes(); // cache the value
+                if (bytes.length > 0) {
                     StatsEvent e = StatsEvent.newBuilder()
                             .setAtomId(atomTag)
-                            .writeByteArray(proto.getBytes())
+                            .writeByteArray(bytes)
+                            // This is a shard ID, and is specified in the metric definition to be
+                            // a dimension. This will result in statsd using RANDOM_ONE_SAMPLE to
+                            // keep all the shards, as it thinks each shard is a different dimension
+                            // of data.
+                            .writeInt(i)
                             .build();
                     pulledData.add(e);
                 }