Merge "Change to get network history from NetworkStatsManager."
diff --git a/packages/SettingsLib/src/com/android/settingslib/net/DataUsageController.java b/packages/SettingsLib/src/com/android/settingslib/net/DataUsageController.java
index eeaa987..183d485 100644
--- a/packages/SettingsLib/src/com/android/settingslib/net/DataUsageController.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/DataUsageController.java
@@ -24,6 +24,8 @@
 import static android.text.format.DateUtils.FORMAT_ABBREV_MONTH;
 import static android.text.format.DateUtils.FORMAT_SHOW_DATE;
 
+import android.app.usage.NetworkStats.Bucket;
+import android.app.usage.NetworkStatsManager;
 import android.content.Context;
 import android.net.ConnectivityManager;
 import android.net.INetworkStatsService;
@@ -37,6 +39,7 @@
 import android.telephony.SubscriptionManager;
 import android.telephony.TelephonyManager;
 import android.text.format.DateUtils;
+import android.util.FeatureFlagUtils;
 import android.util.Log;
 import android.util.Range;
 
@@ -51,6 +54,8 @@
 public class DataUsageController {
 
     private static final String TAG = "DataUsageController";
+    @VisibleForTesting
+    static final String DATA_USAGE_V2 = "settings_data_usage_v2";
     private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
     private static final int FIELDS = FIELD_RX_BYTES | FIELD_TX_BYTES;
     private static final StringBuilder PERIOD_BUILDER = new StringBuilder(50);
@@ -62,6 +67,7 @@
     private final ConnectivityManager mConnectivityManager;
     private final INetworkStatsService mStatsService;
     private final NetworkPolicyManager mPolicyManager;
+    private final NetworkStatsManager mNetworkStatsManager;
 
     private INetworkStatsSession mSession;
     private Callback mCallback;
@@ -74,6 +80,7 @@
         mStatsService = INetworkStatsService.Stub.asInterface(
                 ServiceManager.getService(Context.NETWORK_STATS_SERVICE));
         mPolicyManager = NetworkPolicyManager.from(mContext);
+        mNetworkStatsManager = context.getSystemService(NetworkStatsManager.class);
     }
 
     public void setNetworkController(NetworkNameProvider networkController) {
@@ -89,6 +96,7 @@
     }
 
     @VisibleForTesting
+    @Deprecated
     INetworkStatsSession getSession() {
         if (mSession == null) {
             try {
@@ -128,71 +136,72 @@
     }
 
     public DataUsageInfo getDataUsageInfo(NetworkTemplate template) {
-        final INetworkStatsSession session = getSession();
-        if (session == null) {
-            return warn("no stats session");
-        }
         final NetworkPolicy policy = findNetworkPolicy(template);
-        try {
-            final NetworkStatsHistory history = session.getHistoryForNetwork(template, FIELDS);
-            final long now = System.currentTimeMillis();
-            final long start, end;
-            final Iterator<Range<ZonedDateTime>> it =
-                    (policy != null) ? policy.cycleIterator() : null;
-            if (it != null && it.hasNext()) {
-                final Range<ZonedDateTime> cycle = it.next();
-                start = cycle.getLower().toInstant().toEpochMilli();
-                end = cycle.getUpper().toInstant().toEpochMilli();
-            } else {
-                // period = last 4 wks
-                end = now;
-                start = now - DateUtils.WEEK_IN_MILLIS * 4;
-            }
-            final long callStart = System.currentTimeMillis();
-            final NetworkStatsHistory.Entry entry = history.getValues(start, end, now, null);
-            final long callEnd = System.currentTimeMillis();
-            if (DEBUG) Log.d(TAG, String.format("history call from %s to %s now=%s took %sms: %s",
-                    new Date(start), new Date(end), new Date(now), callEnd - callStart,
-                    historyEntryToString(entry)));
-            if (entry == null) {
-                return warn("no entry data");
-            }
-            final long totalBytes = entry.rxBytes + entry.txBytes;
-            final DataUsageInfo usage = new DataUsageInfo();
-            usage.startDate = start;
-            usage.usageLevel = totalBytes;
-            usage.period = formatDateRange(start, end);
-            usage.cycleStart = start;
-            usage.cycleEnd = end;
-
-            if (policy != null) {
-                usage.limitLevel = policy.limitBytes > 0 ? policy.limitBytes : 0;
-                usage.warningLevel = policy.warningBytes > 0 ? policy.warningBytes : 0;
-            } else {
-                usage.warningLevel = getDefaultWarningLevel();
-            }
-            if (usage != null && mNetworkController != null) {
-                usage.carrier = mNetworkController.getMobileDataNetworkName();
-            }
-            return usage;
-        } catch (RemoteException e) {
-            return warn("remote call failed");
+        final long now = System.currentTimeMillis();
+        final long start, end;
+        final Iterator<Range<ZonedDateTime>> it = (policy != null) ? policy.cycleIterator() : null;
+        if (it != null && it.hasNext()) {
+            final Range<ZonedDateTime> cycle = it.next();
+            start = cycle.getLower().toInstant().toEpochMilli();
+            end = cycle.getUpper().toInstant().toEpochMilli();
+        } else {
+            // period = last 4 wks
+            end = now;
+            start = now - DateUtils.WEEK_IN_MILLIS * 4;
         }
+        final long totalBytes;
+        final long callStart = System.currentTimeMillis();
+        if (FeatureFlagUtils.isEnabled(mContext, DATA_USAGE_V2)) {
+            totalBytes = getUsageLevel(template, start, end);
+        } else {
+            totalBytes = getUsageLevel(template, start, end, now);
+        }
+        if (totalBytes < 0L) {
+            return warn("no entry data");
+        }
+        final DataUsageInfo usage = new DataUsageInfo();
+        usage.startDate = start;
+        usage.usageLevel = totalBytes;
+        usage.period = formatDateRange(start, end);
+        usage.cycleStart = start;
+        usage.cycleEnd = end;
+
+        if (policy != null) {
+            usage.limitLevel = policy.limitBytes > 0 ? policy.limitBytes : 0;
+            usage.warningLevel = policy.warningBytes > 0 ? policy.warningBytes : 0;
+        } else {
+            usage.warningLevel = getDefaultWarningLevel();
+        }
+        if (usage != null && mNetworkController != null) {
+            usage.carrier = mNetworkController.getMobileDataNetworkName();
+        }
+        return usage;
     }
 
     /**
      * Get the total usage level recorded in the network history
      * @param template the network template to retrieve the network history
-     * @return the total usage level recorded in the network history
+     * @return the total usage level recorded in the network history or -1L if there is error
+     * retrieving the data.
      */
-    public long getHistoriclUsageLevel(NetworkTemplate template) {
+    public long getHistoricalUsageLevel(NetworkTemplate template) {
+        if (FeatureFlagUtils.isEnabled(mContext, DATA_USAGE_V2)) {
+            return getUsageLevel(template, 0L /* start */, System.currentTimeMillis() /* end */);
+        } else {
+            final long now = System.currentTimeMillis();
+            return getUsageLevel(template, 0L /* start */, now /* end */, now);
+        }
+    }
+
+    @Deprecated
+    private long getUsageLevel(NetworkTemplate template, long start, long end, long now) {
         final INetworkStatsSession session = getSession();
         if (session != null) {
             try {
-                final NetworkStatsHistory history = session.getHistoryForNetwork(template, FIELDS);
-                final long now = System.currentTimeMillis();
-                final NetworkStatsHistory.Entry entry =
-                        history.getValues(0L /* start */, now /* end */, now, null /* recycle */);
+                final NetworkStatsHistory history =
+                    session.getHistoryForNetwork(template, FIELDS);
+                final NetworkStatsHistory.Entry entry = history.getValues(
+                        start, end, System.currentTimeMillis() /* now */, null /* recycle */);
                 if (entry != null) {
                     return entry.rxBytes + entry.txBytes;
                 }
@@ -201,7 +210,21 @@
                 Log.w(TAG, "Failed to get data usage, remote call failed");
             }
         }
-        return 0L;
+        return -1L;
+    }
+
+    private long getUsageLevel(NetworkTemplate template, long start, long end) {
+        try {
+            final Bucket bucket = mNetworkStatsManager.querySummaryForDevice(
+                getNetworkType(template), getActiveSubscriberId(mContext), start, end);
+            if (bucket != null) {
+                return bucket.getRxBytes() + bucket.getTxBytes();
+            }
+            Log.w(TAG, "Failed to get data usage, no entry data");
+        } catch (RemoteException e) {
+            Log.w(TAG, "Failed to get data usage, remote call failed");
+        }
+        return -1L;
     }
 
     private NetworkPolicy findNetworkPolicy(NetworkTemplate template) {
@@ -218,6 +241,7 @@
         return null;
     }
 
+    @Deprecated
     private static String historyEntryToString(NetworkStatsHistory.Entry entry) {
         return entry == null ? null : new StringBuilder("Entry[")
                 .append("bucketDuration=").append(entry.bucketDuration)
@@ -231,6 +255,17 @@
                 .append(']').toString();
     }
 
+    private static String statsBucketToString(Bucket bucket) {
+        return bucket == null ? null : new StringBuilder("Entry[")
+            .append("bucketDuration=").append(bucket.getEndTimeStamp() - bucket.getStartTimeStamp())
+            .append(",bucketStart=").append(bucket.getStartTimeStamp())
+            .append(",rxBytes=").append(bucket.getRxBytes())
+            .append(",rxPackets=").append(bucket.getRxPackets())
+            .append(",txBytes=").append(bucket.getTxBytes())
+            .append(",txPackets=").append(bucket.getTxPackets())
+            .append(']').toString();
+    }
+
     public void setMobileDataEnabled(boolean enabled) {
         Log.d(TAG, "setMobileDataEnabled: enabled=" + enabled);
         mTelephonyManager.setDataEnabled(enabled);
@@ -249,6 +284,25 @@
         return mTelephonyManager.getDataEnabled();
     }
 
+    static int getNetworkType(NetworkTemplate networkTemplate) {
+        if (networkTemplate == null) {
+            return ConnectivityManager.TYPE_NONE;
+        }
+        final int matchRule = networkTemplate.getMatchRule();
+        switch (matchRule) {
+            case NetworkTemplate.MATCH_MOBILE:
+            case NetworkTemplate.MATCH_MOBILE_WILDCARD:
+                return ConnectivityManager.TYPE_MOBILE;
+            case NetworkTemplate.MATCH_WIFI:
+            case NetworkTemplate.MATCH_WIFI_WILDCARD:
+                return  ConnectivityManager.TYPE_WIFI;
+            case NetworkTemplate.MATCH_ETHERNET:
+                return  ConnectivityManager.TYPE_ETHERNET;
+            default:
+                return ConnectivityManager.TYPE_MOBILE;
+        }
+    }
+
     private static String getActiveSubscriberId(Context context) {
         final TelephonyManager tele = TelephonyManager.from(context);
         final String actualSubscriberId = tele.getSubscriberId(
diff --git a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleChartDataLoader.java b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleChartDataLoader.java
index 7ae3398..ec5a0b5 100644
--- a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleChartDataLoader.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleChartDataLoader.java
@@ -43,9 +43,9 @@
     @Override
     void recordUsage(long start, long end) {
         try {
-            final NetworkStats stats = mNetworkStatsManager.querySummary(
+            final NetworkStats.Bucket bucket = mNetworkStatsManager.querySummaryForDevice(
                 mNetworkType, mSubId, start, end);
-            final long total = getTotalUsage(stats);
+            final long total = bucket == null ? 0L : bucket.getRxBytes() + bucket.getTxBytes();
             if (total > 0L) {
                 final NetworkCycleChartData.Builder builder = new NetworkCycleChartData.Builder();
                 builder.setUsageBuckets(getUsageBuckets(start, end))
@@ -80,9 +80,11 @@
         while (bucketEnd <= end) {
             long usage = 0L;
             try {
-                final NetworkStats stats = mNetworkStatsManager.querySummary(
+                final NetworkStats.Bucket bucket = mNetworkStatsManager.querySummaryForDevice(
                     mNetworkType, mSubId, bucketStart, bucketEnd);
-                usage = getTotalUsage(stats);
+                if (bucket != null) {
+                    usage = bucket.getRxBytes() + bucket.getTxBytes();
+                }
             } catch (RemoteException e) {
                 Log.e(TAG, "Exception querying network detail.", e);
             }
diff --git a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataLoader.java b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataLoader.java
index b1c2c3a..d957801 100644
--- a/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataLoader.java
+++ b/packages/SettingsLib/src/com/android/settingslib/net/NetworkCycleDataLoader.java
@@ -176,31 +176,11 @@
 
         public Builder<T> setNetworkTemplate(NetworkTemplate template) {
             mNetworkTemplate = template;
-            setNetworkType();
+            mNetworkType = DataUsageController.getNetworkType(template);
             return this;
         }
 
         public abstract T build();
-
-        private void setNetworkType() {
-            if (mNetworkTemplate != null) {
-                final int matchRule = mNetworkTemplate.getMatchRule();
-                switch (matchRule) {
-                    case NetworkTemplate.MATCH_MOBILE:
-                    case NetworkTemplate.MATCH_MOBILE_WILDCARD:
-                        mNetworkType = ConnectivityManager.TYPE_MOBILE;
-                        break;
-                    case NetworkTemplate.MATCH_WIFI:
-                        mNetworkType = ConnectivityManager.TYPE_WIFI;
-                        break;
-                    case NetworkTemplate.MATCH_ETHERNET:
-                        mNetworkType = ConnectivityManager.TYPE_ETHERNET;
-                        break;
-                    default:
-                        mNetworkType = ConnectivityManager.TYPE_MOBILE;
-                }
-            }
-        }
     }
 
 }
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/DataUsageControllerTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/DataUsageControllerTest.java
index 1be856a..b6ac467 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/DataUsageControllerTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/DataUsageControllerTest.java
@@ -24,16 +24,23 @@
 import static org.mockito.Mockito.anyLong;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.eq;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.usage.NetworkStats;
+import android.app.usage.NetworkStatsManager;
 import android.content.Context;
+import android.net.ConnectivityManager;
 import android.net.INetworkStatsSession;
 import android.net.NetworkStatsHistory;
 import android.net.NetworkStatsHistory.Entry;
 import android.net.NetworkTemplate;
 import android.os.RemoteException;
+import android.telephony.TelephonyManager;
 import android.text.format.DateUtils;
+import android.util.FeatureFlagUtils;
 
 import com.android.settingslib.SettingsLibRobolectricTestRunner;
 
@@ -47,8 +54,14 @@
 @RunWith(SettingsLibRobolectricTestRunner.class)
 public class DataUsageControllerTest {
 
+    private static final String SUB_ID = "Test Subscriber";
+
     @Mock
     private INetworkStatsSession mSession;
+    @Mock
+    private TelephonyManager mTelephonyManager;
+    @Mock
+    private NetworkStatsManager mNetworkStatsManager;
 
     private Context mContext;
     private DataUsageController mController;
@@ -63,13 +76,14 @@
                 new NetworkStatsHistory(DateUtils.DAY_IN_MILLIS /* bucketDuration */));
         doReturn(mNetworkStatsHistory)
                 .when(mSession).getHistoryForNetwork(any(NetworkTemplate.class), anyInt());
+        doReturn(SUB_ID).when(mTelephonyManager).getSubscriberId(anyInt());
     }
 
     @Test
-    public void getHistoriclUsageLevel_noNetworkSession_shouldReturn0() {
+    public void getHistoricalUsageLevel_noNetworkSession_shouldReturnNegative1() {
         doReturn(null).when(mController).getSession();
 
-        assertThat(mController.getHistoriclUsageLevel(null /* template */)).isEqualTo(0L);
+        assertThat(mController.getHistoricalUsageLevel(null /* template */)).isEqualTo(-1L);
 
     }
 
@@ -77,13 +91,13 @@
     public void getHistoriclUsageLevel_noUsageData_shouldReturn0() {
         doReturn(mSession).when(mController).getSession();
 
-        assertThat(mController.getHistoriclUsageLevel(NetworkTemplate.buildTemplateWifiWildcard()))
+        assertThat(mController.getHistoricalUsageLevel(NetworkTemplate.buildTemplateWifiWildcard()))
                 .isEqualTo(0L);
 
     }
 
     @Test
-    public void getHistoriclUsageLevel_hasUsageData_shouldReturnTotalUsage() {
+    public void getHistoricalUsageLevel_hasUsageData_shouldReturnTotalUsage() {
         doReturn(mSession).when(mController).getSession();
         final long receivedBytes = 743823454L;
         final long transmittedBytes = 16574289L;
@@ -94,8 +108,57 @@
         when(mNetworkStatsHistory.getValues(eq(0L), anyLong(), anyLong(), nullable(Entry.class)))
                 .thenReturn(entry);
 
-        assertThat(mController.getHistoriclUsageLevel(NetworkTemplate.buildTemplateWifiWildcard()))
+        assertThat(mController.getHistoricalUsageLevel(NetworkTemplate.buildTemplateWifiWildcard()))
                 .isEqualTo(receivedBytes + transmittedBytes);
 
     }
+
+    @Test
+    public void getHistoricalUsageLevel_v2_shouldQuerySummaryForDevice() throws Exception {
+        final Context context = mock(Context.class);
+        FeatureFlagUtils.setEnabled(context, DataUsageController.DATA_USAGE_V2, true);
+        when(context.getSystemService(Context.TELEPHONY_SERVICE)).thenReturn(mTelephonyManager);
+        when(context.getSystemService(NetworkStatsManager.class)).thenReturn(mNetworkStatsManager);
+        final DataUsageController controller = new DataUsageController(context);
+
+        controller.getHistoricalUsageLevel(NetworkTemplate.buildTemplateWifiWildcard());
+
+        verify(mNetworkStatsManager).querySummaryForDevice(eq(ConnectivityManager.TYPE_WIFI),
+                eq(SUB_ID), eq(0L) /* startTime */, anyLong() /* endTime */);
+    }
+
+    @Test
+    public void getHistoricalUsageLevel_v2NoUsageData_shouldReturn0() throws Exception {
+        final Context context = mock(Context.class);
+        FeatureFlagUtils.setEnabled(context, DataUsageController.DATA_USAGE_V2, true);
+        when(context.getSystemService(Context.TELEPHONY_SERVICE)).thenReturn(mTelephonyManager);
+        when(context.getSystemService(NetworkStatsManager.class)).thenReturn(mNetworkStatsManager);
+        when(mNetworkStatsManager.querySummaryForDevice(eq(ConnectivityManager.TYPE_WIFI),
+                eq(SUB_ID), eq(0L) /* startTime */, anyLong() /* endTime */))
+                .thenReturn(mock(NetworkStats.Bucket.class));
+        final DataUsageController controller = new DataUsageController(context);
+
+        assertThat(controller.getHistoricalUsageLevel(NetworkTemplate.buildTemplateWifiWildcard()))
+            .isEqualTo(0L);
+    }
+
+    @Test
+    public void getHistoricalUsageLevel_v2HasUsageData_shouldReturnTotalUsage()
+            throws Exception {
+        final Context context = mock(Context.class);
+        FeatureFlagUtils.setEnabled(context, DataUsageController.DATA_USAGE_V2, true);
+        when(context.getSystemService(Context.TELEPHONY_SERVICE)).thenReturn(mTelephonyManager);
+        when(context.getSystemService(NetworkStatsManager.class)).thenReturn(mNetworkStatsManager);
+        final long receivedBytes = 743823454L;
+        final long transmittedBytes = 16574289L;
+        final NetworkStats.Bucket bucket = mock(NetworkStats.Bucket.class);
+        when(bucket.getRxBytes()).thenReturn(receivedBytes);
+        when(bucket.getTxBytes()).thenReturn(transmittedBytes);
+        when(mNetworkStatsManager.querySummaryForDevice(eq(ConnectivityManager.TYPE_WIFI),
+                eq(SUB_ID), eq(0L) /* startTime */, anyLong() /* endTime */)).thenReturn(bucket);
+        final DataUsageController controller = new DataUsageController(context);
+
+        assertThat(controller.getHistoricalUsageLevel(NetworkTemplate.buildTemplateWifiWildcard()))
+                .isEqualTo(receivedBytes + transmittedBytes);
+    }
 }
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleChartDataLoaderTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleChartDataLoaderTest.java
index cad88b1..0a03631 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleChartDataLoaderTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/net/NetworkCycleChartDataLoaderTest.java
@@ -58,7 +58,7 @@
     }
 
     @Test
-    public void recordUsage_shouldQueryNetworkSummary() throws RemoteException {
+    public void recordUsage_shouldQueryNetworkSummaryForDevice() throws RemoteException {
         final long end = System.currentTimeMillis();
         final long start = end - (DateUtils.WEEK_IN_MILLIS * 4);
         final int networkType = ConnectivityManager.TYPE_MOBILE;
@@ -68,6 +68,6 @@
 
         mLoader.recordUsage(start, end);
 
-        verify(mNetworkStatsManager).querySummary(networkType, subId, start, end);
+        verify(mNetworkStatsManager).querySummaryForDevice(networkType, subId, start, end);
     }
 }