Change to get network history from NetworkStatsManager.

- in DataUsageController, change to use
NetworkStatsManager.querySummaryForDevice() to get network history.
- also change to use querySummaryForDevice() instead of querySummary()
in chart loader, since it returns the aggregated bucket, so we do not
need to do the extra work for summing up the usage data.
- move the mapping of network template to network type to
DataUsageController and add the mapping for wifi wildcard.

Bug: 111751694
Test: make RunSettingsLibRoboTests
Change-Id: I6523869563cd53f791fc955468e1842b4f6adba4

5 files changed, 186 insertions, 87 deletions

diff --git a/packages/SettingsLib/src/com/android/settingslib/net/DataUsageController.java b/packages/SettingsLib/src/com/android/settingslib/net/DataUsageController.java
index eeaa9872a452..183d4856ac6f 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.telephony.TelephonyManager.SIM_STATE_READY;
 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.os.ServiceManager;
 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 @@ import java.util.Locale;
 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 @@ public class DataUsageController {
     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 @@ public class DataUsageController {
         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 @@ public class DataUsageController {
     }
 
     @VisibleForTesting
+    @Deprecated
     INetworkStatsSession getSession() {
         if (mSession == null) {
             try {
@@ -128,71 +136,72 @@ public class DataUsageController {
     }
 
     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 @@ public class DataUsageController {
                 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 @@ public class DataUsageController {
         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 @@ public class DataUsageController {
                 .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 @@ public class DataUsageController {
         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 7ae3398d42ea..ec5a0b5cc4cd 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 @@ public class NetworkCycleChartDataLoader
     @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 @@ public class NetworkCycleChartDataLoader
         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 b1c2c3a2d2e5..d9578014e846 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 abstract class NetworkCycleDataLoader<D> extends AsyncTaskLoader<D> {
 
         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 1be856a59fb5..b6ac467c729f 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.anyInt;
 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 @@ import org.robolectric.RuntimeEnvironment;
 @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 @@ public class DataUsageControllerTest {
                 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 class DataUsageControllerTest {
     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 @@ public class DataUsageControllerTest {
         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 cad88b1f97a4..0a036317910e 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 @@ public class NetworkCycleChartDataLoaderTest {
     }
 
     @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 @@ public class NetworkCycleChartDataLoaderTest {
 
         mLoader.recordUsage(start, end);
 
-        verify(mNetworkStatsManager).querySummary(networkType, subId, start, end);
+        verify(mNetworkStatsManager).querySummaryForDevice(networkType, subId, start, end);
     }
 }