Adds pulling for network bytes for statsd.

Adds atoms to collect network bytes transmitted and received via wifi
and mobile data. We need to get the list of correct ifaces from
BatteryStats since NetworkStatsService only tracks the mobile ifaces.
We split the atoms so that we can also capture metrics sliced on
foreground and background since they are available from the kernel.
Also adds an atom for the Kernel Wakelocks.

Test: Manually tested with adb shell cmd stats pull-source after
modifying the StatsPullerManager.cpp

Change-Id: I9467afad799c6d97560e868f8625fefae5c2b5e1
diff --git a/cmds/statsd/src/external/KernelWakelockPuller.cpp b/cmds/statsd/src/external/KernelWakelockPuller.cpp
index 521d3f6..00259a8 100644
--- a/cmds/statsd/src/external/KernelWakelockPuller.cpp
+++ b/cmds/statsd/src/external/KernelWakelockPuller.cpp
@@ -33,7 +33,7 @@
 namespace os {
 namespace statsd {
 
-const int KernelWakelockPuller::PULL_CODE_KERNEL_WAKELOCKS = 20;
+const int KernelWakelockPuller::PULL_CODE_KERNEL_WAKELOCKS = 1004;
 
 // The reading and parsing are implemented in Java. It is not difficult to port over. But for now
 // let StatsCompanionService handle that and send the data back.
diff --git a/cmds/statsd/src/stats_events.proto b/cmds/statsd/src/stats_events.proto
index e4e18de..51244c6 100644
--- a/cmds/statsd/src/stats_events.proto
+++ b/cmds/statsd/src/stats_events.proto
@@ -35,7 +35,8 @@
  * in the format defined here and in stats_log.proto.
  */
 message StatsEvent {
-    oneof event {
+    // Pushed events start at 2.
+    oneof pushed {
         // For StatsLog reasons, 1 is illegal and will not work. Must start at 2.
         BleScanStateChanged ble_scan_state_changed = 2;
         BleUnoptimizedScanStateChanged ble_unoptimized_scan_state_changed = 3;
@@ -70,10 +71,18 @@
         WifiScanStateChanged wifi_scan_state_changed = 39;
         PhoneSignalStrengthChanged phone_signal_strength_changed = 40;
         SettingChanged setting_changed = 41;
-        KernelWakelockPulled kernel_wakelock_pulled = 42;
-        ActivityForegroundStateChanged activity_foreground_state_changed = 43;
+        ActivityForegroundStateChanged activity_foreground_state_changed = 42;
         // TODO: Reorder the numbering so that the most frequent occur events occur in the first 15.
     }
+
+    // Pulled events will start at field 1000.
+    oneof pulled {
+        WifiBytesTransferred wifi_bytes_transferred = 1000;
+        WifiBytesTransferredByFgBg wifi_bytes_transferred_by_fg_bg = 1001;
+        MobileBytesTransferred mobile_bytes_transferred = 1002;
+        MobileBytesTransferredByFgBg mobile_bytes_transferred_by_fg_bg = 1003;
+        KernelWakelocksReported kernel_wakelocks_reported = 1004;
+    }
 }
 
 /**
@@ -682,18 +691,6 @@
     optional int32 user = 7;
 }
 
-/*
- * Pulls kernel wakelock changes.
- *
- * Pulled from:
-  *   frameworks/base/services/core/java/com/android/server/stats/StatsCompanionService.java
- */
-message KernelWakelockPulled {
-    optional int32 count = 1;
-    optional int32 version = 2;
-    optional int64 total_time = 3;
-    optional string name = 4;
-}
 
 /*
  * Logs activity going to foreground or background
@@ -711,3 +708,98 @@
     optional string class_name = 3;
     optional Activity activity = 4;
 }
+
+/**
+ * Pulls bytes transferred via wifi (Sum of foreground and background usage).
+ *
+ * Pulled from:
+ *   StatsCompanionService (using BatteryStats to get which interfaces are wifi)
+ */
+message WifiBytesTransferred {
+    optional int32 uid = 1;
+
+    optional int64 rx_bytes = 2;
+
+    optional int64 rx_packets = 3;
+
+    optional int64 tx_bytes = 4;
+
+    optional int64 tx_packets = 5;
+}
+
+/**
+ * Pulls bytes transferred via wifi (separated by foreground and background usage).
+ *
+ * Pulled from:
+ *   StatsCompanionService (using BatteryStats to get which interfaces are wifi)
+ */
+message WifiBytesTransferredByFgBg {
+    optional int32 uid = 1;
+
+    // 1 denotes foreground and 0 denotes background. This is called Set in NetworkStats.
+    optional int32 is_foreground = 2;
+
+    optional int64 rx_bytes = 3;
+
+    optional int64 rx_packets = 4;
+
+    optional int64 tx_bytes = 5;
+
+    optional int64 tx_packets = 6;
+}
+
+/**
+ * Pulls bytes transferred via mobile networks (Sum of foreground and background usage).
+ *
+ * Pulled from:
+ *   StatsCompanionService (using BatteryStats to get which interfaces are mobile data)
+ */
+message MobileBytesTransferred {
+    optional int32 uid = 1;
+
+    optional int64 rx_bytes = 2;
+
+    optional int64 rx_packets = 3;
+
+    optional int64 tx_bytes = 4;
+
+    optional int64 tx_packets = 5;
+}
+
+/**
+ * Pulls bytes transferred via mobile networks (separated by foreground and background usage).
+ *
+ * Pulled from:
+ *   StatsCompanionService (using BatteryStats to get which interfaces are mobile data)
+ */
+message MobileBytesTransferredByFgBg {
+    optional int32 uid = 1;
+
+    // 1 denotes foreground and 0 denotes background. This is called Set in NetworkStats.
+    optional int32 is_foreground = 2;
+
+    optional int64 rx_bytes = 3;
+
+    optional int64 rx_packets = 4;
+
+    optional int64 tx_bytes = 5;
+
+    optional int64 tx_packets = 6;
+}
+
+/**
+ * Pulls the kernel wakelock durations. This atom is adapted from
+ * android/internal/os/KernelWakelockStats.java
+ *
+ * Pulled from:
+ *   StatsCompanionService using KernelWakelockReader.
+ */
+message KernelWakelocksReported {
+    optional string name = 1;
+
+    optional int32 count = 2;
+
+    optional int32 version = 3;
+
+    optional int64 time = 4;
+}
diff --git a/core/java/android/os/BatteryStatsInternal.java b/core/java/android/os/BatteryStatsInternal.java
new file mode 100644
index 0000000..b0436eb
--- /dev/null
+++ b/core/java/android/os/BatteryStatsInternal.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright 2017 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 android.os;
+
+/**
+ * Battery stats local system service interface. This is used to pass internal data out of
+ * BatteryStatsImpl.
+ *
+ * @hide Only for use within Android OS.
+ */
+public abstract class BatteryStatsInternal {
+    /**
+     * Returns the wifi interfaces.
+     */
+    public abstract String[] getWifiIfaces();
+
+    /**
+     * Returns the mobile data interfaces.
+     */
+    public abstract String[] getMobileIfaces();
+}
diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index f0d05da..0535ebe 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -5412,6 +5412,18 @@
         }
     }
 
+    public String[] getWifiIfaces() {
+        synchronized (mWifiNetworkLock) {
+            return mWifiIfaces;
+        }
+    }
+
+    public String[] getMobileIfaces() {
+        synchronized (mModemNetworkLock) {
+            return mModemIfaces;
+        }
+    }
+
     @Override public long getScreenOnTime(long elapsedRealtimeUs, int which) {
         return mScreenOnTimer.getTotalTimeLocked(elapsedRealtimeUs, which);
     }
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index db12ae2..52f9702 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -23,6 +23,7 @@
 import android.net.wifi.WifiActivityEnergyInfo;
 import android.os.PowerSaveState;
 import android.os.BatteryStats;
+import android.os.BatteryStatsInternal;
 import android.os.Binder;
 import android.os.Handler;
 import android.os.IBinder;
@@ -177,9 +178,22 @@
     }
 
     public void publish() {
+        LocalServices.addService(BatteryStatsInternal.class, new LocalService());
         ServiceManager.addService(BatteryStats.SERVICE_NAME, asBinder());
     }
 
+    private final class LocalService extends BatteryStatsInternal {
+        @Override
+        public String[] getWifiIfaces() {
+            return mStats.getWifiIfaces().clone();
+        }
+
+        @Override
+        public String[] getMobileIfaces() {
+            return mStats.getMobileIfaces().clone();
+        }
+    }
+
     private static void awaitUninterruptibly(Future<?> future) {
         while (true) {
             try {
diff --git a/services/core/java/com/android/server/stats/StatsCompanionService.java b/services/core/java/com/android/server/stats/StatsCompanionService.java
index f50939f..41534cb 100644
--- a/services/core/java/com/android/server/stats/StatsCompanionService.java
+++ b/services/core/java/com/android/server/stats/StatsCompanionService.java
@@ -24,7 +24,8 @@
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.UserInfo;
-import android.content.IntentFilter;
+import android.net.NetworkStats;
+import android.os.BatteryStatsInternal;
 import android.os.Binder;
 import android.os.Bundle;
 import android.os.IBinder;
@@ -37,15 +38,17 @@
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.util.Slog;
+import android.util.StatsLog;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.net.NetworkStatsFactory;
+import com.android.internal.os.KernelWakelockReader;
+import com.android.internal.os.KernelWakelockStats;
+import com.android.server.LocalServices;
+import com.android.server.SystemService;
 
 import java.util.ArrayList;
 import java.util.List;
-
-import com.android.internal.annotations.GuardedBy;
-import com.android.internal.os.KernelWakelockReader;
-import com.android.internal.os.KernelWakelockStats;
-import com.android.server.SystemService;
-
 import java.util.Map;
 
 /**
@@ -290,35 +293,168 @@
     private final KernelWakelockReader mKernelWakelockReader = new KernelWakelockReader();
     private final KernelWakelockStats mTmpWakelockStats = new KernelWakelockStats();
 
+    private StatsLogEventWrapper[] addNetworkStats(int tag, NetworkStats stats, boolean withFGBG) {
+        List<StatsLogEventWrapper> ret = new ArrayList<>();
+        int size = stats.size();
+        NetworkStats.Entry entry = new NetworkStats.Entry(); // For recycling
+        for (int j = 0; j < size; j++) {
+            stats.getValues(j, entry);
+            StatsLogEventWrapper e = new StatsLogEventWrapper(tag, withFGBG ? 6 : 5);
+            e.writeInt(entry.uid);
+            if (withFGBG) {
+                e.writeInt(entry.set);
+            }
+            e.writeLong(entry.rxBytes);
+            e.writeLong(entry.rxPackets);
+            e.writeLong(entry.txBytes);
+            e.writeLong(entry.txPackets);
+            ret.add(e);
+        }
+        return ret.toArray(new StatsLogEventWrapper[ret.size()]);
+    }
+
+    /**
+     * Allows rollups per UID but keeping the set (foreground/background) slicing.
+     * Adapted from groupedByUid in frameworks/base/core/java/android/net/NetworkStats.java
+     */
+    private NetworkStats rollupNetworkStatsByFGBG(NetworkStats stats) {
+        final NetworkStats ret = new NetworkStats(stats.getElapsedRealtime(), 1);
+
+        final NetworkStats.Entry entry = new NetworkStats.Entry();
+        entry.iface = NetworkStats.IFACE_ALL;
+        entry.tag = NetworkStats.TAG_NONE;
+        entry.metered = NetworkStats.METERED_ALL;
+        entry.roaming = NetworkStats.ROAMING_ALL;
+
+        int size = stats.size();
+        NetworkStats.Entry recycle = new NetworkStats.Entry(); // Used for retrieving values
+        for (int i = 0; i < size; i++) {
+            stats.getValues(i, recycle);
+
+            // Skip specific tags, since already counted in TAG_NONE
+            if (recycle.tag != NetworkStats.TAG_NONE) continue;
+
+            entry.set = recycle.set; // Allows slicing by background/foreground
+            entry.uid = recycle.uid;
+            entry.rxBytes = recycle.rxBytes;
+            entry.rxPackets = recycle.rxPackets;
+            entry.txBytes = recycle.txBytes;
+            entry.txPackets = recycle.txPackets;
+            // Operations purposefully omitted since we don't use them for statsd.
+            ret.combineValues(entry);
+        }
+        return ret;
+    }
+
     @Override // Binder call
     public StatsLogEventWrapper[] pullData(int pullCode) {
         enforceCallingPermission();
-        if (DEBUG) {
+        if (DEBUG)
             Slog.d(TAG, "Pulling " + pullCode);
-        }
 
-        List<StatsLogEventWrapper> ret = new ArrayList<>();
         switch (pullCode) {
-            case PULL_CODE_KERNEL_WAKELOCKS: {
+            case StatsLog.WIFI_BYTES_TRANSFERRED: {
+                long token = Binder.clearCallingIdentity();
+                try {
+                    // TODO: Consider caching the following call to get BatteryStatsInternal.
+                    BatteryStatsInternal bs = LocalServices.getService(BatteryStatsInternal.class);
+                    String[] ifaces = bs.getWifiIfaces();
+                    if (ifaces.length == 0) {
+                        return null;
+                    }
+                    NetworkStatsFactory nsf = new NetworkStatsFactory();
+                    // Combine all the metrics per Uid into one record.
+                    NetworkStats stats = nsf.readNetworkStatsDetail(NetworkStats.UID_ALL, ifaces,
+                            NetworkStats.TAG_NONE, null).groupedByUid();
+                    return addNetworkStats(pullCode, stats, false);
+                } catch (java.io.IOException e) {
+                    Slog.e(TAG, "Pulling netstats for wifi bytes has error", e);
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
+                break;
+            }
+            case StatsLog.MOBILE_BYTES_TRANSFERRED: {
+                long token = Binder.clearCallingIdentity();
+                try {
+                    BatteryStatsInternal bs = LocalServices.getService(BatteryStatsInternal.class);
+                    String[] ifaces = bs.getMobileIfaces();
+                    if (ifaces.length == 0) {
+                        return null;
+                    }
+                    NetworkStatsFactory nsf = new NetworkStatsFactory();
+                    // Combine all the metrics per Uid into one record.
+                    NetworkStats stats = nsf.readNetworkStatsDetail(NetworkStats.UID_ALL, ifaces,
+                        NetworkStats.TAG_NONE, null).groupedByUid();
+                    return addNetworkStats(pullCode, stats, false);
+                } catch (java.io.IOException e) {
+                    Slog.e(TAG, "Pulling netstats for mobile bytes has error", e);
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
+                break;
+            }
+            case StatsLog.WIFI_BYTES_TRANSFERRED_BY_FG_BG: {
+                long token = Binder.clearCallingIdentity();
+                try {
+                    BatteryStatsInternal bs = LocalServices.getService(BatteryStatsInternal.class);
+                    String[] ifaces = bs.getWifiIfaces();
+                    if (ifaces.length == 0) {
+                        return null;
+                    }
+                    NetworkStatsFactory nsf = new NetworkStatsFactory();
+                    NetworkStats stats = rollupNetworkStatsByFGBG(
+                            nsf.readNetworkStatsDetail(NetworkStats.UID_ALL, ifaces,
+                            NetworkStats.TAG_NONE, null));
+                    return addNetworkStats(pullCode, stats, true);
+                } catch (java.io.IOException e) {
+                    Slog.e(TAG, "Pulling netstats for wifi bytes w/ fg/bg has error", e);
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
+                break;
+            }
+            case StatsLog.MOBILE_BYTES_TRANSFERRED_BY_FG_BG: {
+                long token = Binder.clearCallingIdentity();
+                try {
+                    BatteryStatsInternal bs = LocalServices.getService(BatteryStatsInternal.class);
+                    String[] ifaces = bs.getMobileIfaces();
+                    if (ifaces.length == 0) {
+                        return null;
+                    }
+                    NetworkStatsFactory nsf = new NetworkStatsFactory();
+                    NetworkStats stats = rollupNetworkStatsByFGBG(
+                            nsf.readNetworkStatsDetail(NetworkStats.UID_ALL, ifaces,
+                            NetworkStats.TAG_NONE, null));
+                    return addNetworkStats(pullCode, stats, true);
+                } catch (java.io.IOException e) {
+                    Slog.e(TAG, "Pulling netstats for mobile bytes w/ fg/bg has error", e);
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
+                break;
+            }
+            case StatsLog.KERNEL_WAKELOCKS_REPORTED: {
                 final KernelWakelockStats wakelockStats =
                         mKernelWakelockReader.readKernelWakelockStats(mTmpWakelockStats);
+                List<StatsLogEventWrapper> ret = new ArrayList();
                 for (Map.Entry<String, KernelWakelockStats.Entry> ent : wakelockStats.entrySet()) {
                     String name = ent.getKey();
                     KernelWakelockStats.Entry kws = ent.getValue();
-                    StatsLogEventWrapper e = new StatsLogEventWrapper(41, 4);
+                    StatsLogEventWrapper e = new StatsLogEventWrapper(pullCode, 4);
+                    e.writeString(name);
                     e.writeInt(kws.mCount);
                     e.writeInt(kws.mVersion);
                     e.writeLong(kws.mTotalTime);
-                    e.writeString(name);
                     ret.add(e);
                 }
-                break;
+                return ret.toArray(new StatsLogEventWrapper[ret.size()]);
             }
             default:
                 Slog.w(TAG, "No such pollable data as " + pullCode);
                 return null;
         }
-        return ret.toArray(new StatsLogEventWrapper[ret.size()]);
+        return null;
     }
 
     @Override // Binder call