From fdc5369e3e56ce644b8b1c9e36b7899e5e744a2b Mon Sep 17 00:00:00 2001
From: Jeffrey Huang <jeffreyhuang@google.com>
Date: Tue, 14 Jan 2020 15:17:30 -0800
Subject: [PATCH] Migrate pullDiskIo

Test: No cts test. Ran adb shell cmd stats pull-source 10032
Change-Id: I088cfd911e7f097ac37c485e666fe1d2fc6b3663
---
 .../server/stats/StatsCompanionService.java   | 27 -------------
 .../src/external/StatsPullerManager.cpp       |  6 ---
 .../stats/pull/StatsPullAtomService.java      | 38 +++++++++++++++++--
 3 files changed, 35 insertions(+), 36 deletions(-)

diff --git a/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java b/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java
index 7e0333ca1a677..e86fcfb0331cc 100644
--- a/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java
+++ b/apex/statsd/service/java/com/android/server/stats/StatsCompanionService.java
@@ -806,28 +806,6 @@ public class StatsCompanionService extends IStatsCompanionService.Stub {
         }
     }
 
-    private void pullDiskIo(int tagId, long elapsedNanos, final long wallClockNanos,
-            List<StatsLogEventWrapper> pulledData) {
-        mStoragedUidIoStatsReader.readAbsolute((uid, fgCharsRead, fgCharsWrite, fgBytesRead,
-                fgBytesWrite, bgCharsRead, bgCharsWrite, bgBytesRead, bgBytesWrite,
-                fgFsync, bgFsync) -> {
-            StatsLogEventWrapper e = new StatsLogEventWrapper(tagId, elapsedNanos,
-                    wallClockNanos);
-            e.writeInt(uid);
-            e.writeLong(fgCharsRead);
-            e.writeLong(fgCharsWrite);
-            e.writeLong(fgBytesRead);
-            e.writeLong(fgBytesWrite);
-            e.writeLong(bgCharsRead);
-            e.writeLong(bgCharsWrite);
-            e.writeLong(bgBytesRead);
-            e.writeLong(bgBytesWrite);
-            e.writeLong(fgFsync);
-            e.writeLong(bgFsync);
-            pulledData.add(e);
-        });
-    }
-
     private void pullProcessCpuTime(int tagId, long elapsedNanos, final long wallClockNanos,
             List<StatsLogEventWrapper> pulledData) {
         synchronized (this) {
@@ -1161,11 +1139,6 @@ public class StatsCompanionService extends IStatsCompanionService.Stub {
                 break;
             }
 
-            case StatsLog.DISK_IO: {
-                pullDiskIo(tagId, elapsedNanos, wallClockNanos, ret);
-                break;
-            }
-
             case StatsLog.PROCESS_CPU_TIME: {
                 pullProcessCpuTime(tagId, elapsedNanos, wallClockNanos, ret);
                 break;
diff --git a/cmds/statsd/src/external/StatsPullerManager.cpp b/cmds/statsd/src/external/StatsPullerManager.cpp
index d04d1f53e276f..ed86ea5bc56c2 100644
--- a/cmds/statsd/src/external/StatsPullerManager.cpp
+++ b/cmds/statsd/src/external/StatsPullerManager.cpp
@@ -96,12 +96,6 @@ std::map<PullerKey, PullAtomInfo> StatsPullerManager::kAllPullAtomInfo = {
         {{.atomTag = android::util::PROC_STATS_PKG_PROC},
          {.puller = new StatsCompanionServicePuller(android::util::PROC_STATS_PKG_PROC)}},
 
-        // Disk I/O stats per uid.
-        {{.atomTag = android::util::DISK_IO},
-         {.additiveFields = {2, 3, 4, 5, 6, 7, 8, 9, 10, 11},
-          .coolDownNs = 3 * NS_PER_SEC,
-          .puller = new StatsCompanionServicePuller(android::util::DISK_IO)}},
-
         // Process cpu stats. Min cool-down is 5 sec, inline with what AcitivityManagerService uses.
         {{.atomTag = android::util::PROCESS_CPU_TIME},
          {.coolDownNs = 5 * NS_PER_SEC /* min cool-down in seconds*/,
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 c5a957f3a9d67..19dc032b9b8ff 100644
--- a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
+++ b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
@@ -1881,12 +1881,44 @@ public class StatsPullAtomService extends SystemService {
         // No op.
     }
 
+    private StoragedUidIoStatsReader mStoragedUidIoStatsReader =
+            new StoragedUidIoStatsReader();
+
     private void registerDiskIO() {
-        // No op.
+        int tagId = StatsLog.DISK_IO;
+        PullAtomMetadata metadata = PullAtomMetadata.newBuilder()
+                .setAdditiveFields(new int[] {2, 3, 4, 5, 6, 7, 8, 9, 10, 11})
+                .setCoolDownNs(3 * NS_PER_SEC)
+                .build();
+        mStatsManager.registerPullAtomCallback(
+                tagId,
+                metadata,
+                (atomTag, data) -> pullDiskIO(atomTag, data),
+                BackgroundThread.getExecutor()
+        );
     }
 
-    private void pullDiskIO() {
-        // No op.
+    private int pullDiskIO(int atomTag, List<StatsEvent> pulledData) {
+        mStoragedUidIoStatsReader.readAbsolute((uid, fgCharsRead, fgCharsWrite, fgBytesRead,
+                fgBytesWrite, bgCharsRead, bgCharsWrite, bgBytesRead, bgBytesWrite,
+                fgFsync, bgFsync) -> {
+            StatsEvent e = StatsEvent.newBuilder()
+                    .setAtomId(atomTag)
+                    .writeInt(uid)
+                    .writeLong(fgCharsRead)
+                    .writeLong(fgCharsWrite)
+                    .writeLong(fgBytesRead)
+                    .writeLong(fgBytesWrite)
+                    .writeLong(bgCharsRead)
+                    .writeLong(bgCharsWrite)
+                    .writeLong(bgBytesRead)
+                    .writeLong(bgBytesWrite)
+                    .writeLong(fgFsync)
+                    .writeLong(bgFsync)
+                    .build();
+            pulledData.add(e);
+        });
+        return StatsManager.PULL_SUCCESS;
     }
 
     private void registerPowerProfile() {