From 5a0df01ad7638b3339381b47c495d618f556b289 Mon Sep 17 00:00:00 2001
From: Jeff Sharkey <jsharkey@android.com>
Date: Fri, 20 Apr 2018 10:59:09 -0600
Subject: [PATCH] Give CTS a way to force-poll network stats.

Collecting network statistics is pretty heavy, which is why we're
throttling callers.  However, to keep CTS running fast, we provide a
way for tests to force a poll event, instead of making them wait for
the throttle timeout.

Bug: 77908520
Test: atest cts/tests/tests/app.usage/src/android/app/usage/cts/NetworkUsageStatsTest.java
Change-Id: Ia792f0cd495023366ff8c4839df54e7da2ae8331
---
 .../android/app/usage/NetworkStatsManager.java    | 15 ++++++++++++++-
 .../android/server/net/NetworkStatsService.java   |  3 ++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/core/java/android/app/usage/NetworkStatsManager.java b/core/java/android/app/usage/NetworkStatsManager.java
index 0b21196f5d..9f46f207d6 100644
--- a/core/java/android/app/usage/NetworkStatsManager.java
+++ b/core/java/android/app/usage/NetworkStatsManager.java
@@ -20,6 +20,7 @@ import static com.android.internal.util.Preconditions.checkNotNull;
 
 import android.annotation.Nullable;
 import android.annotation.SystemService;
+import android.annotation.TestApi;
 import android.app.usage.NetworkStats.Bucket;
 import android.content.Context;
 import android.net.ConnectivityManager;
@@ -111,7 +112,9 @@ public class NetworkStatsManager {
     /** @hide */
     public static final int FLAG_POLL_ON_OPEN = 1 << 0;
     /** @hide */
-    public static final int FLAG_AUGMENT_WITH_SUBSCRIPTION_PLAN = 1 << 1;
+    public static final int FLAG_POLL_FORCE = 1 << 1;
+    /** @hide */
+    public static final int FLAG_AUGMENT_WITH_SUBSCRIPTION_PLAN = 1 << 2;
 
     private int mFlags;
 
@@ -140,6 +143,16 @@ public class NetworkStatsManager {
         }
     }
 
+    /** @hide */
+    @TestApi
+    public void setPollForce(boolean pollForce) {
+        if (pollForce) {
+            mFlags |= FLAG_POLL_FORCE;
+        } else {
+            mFlags &= ~FLAG_POLL_FORCE;
+        }
+    }
+
     /** @hide */
     public void setAugmentWithSubscriptionPlan(boolean augmentWithSubscriptionPlan) {
         if (augmentWithSubscriptionPlan) {
diff --git a/services/core/java/com/android/server/net/NetworkStatsService.java b/services/core/java/com/android/server/net/NetworkStatsService.java
index ae7058d42e..9ef6c66b07 100644
--- a/services/core/java/com/android/server/net/NetworkStatsService.java
+++ b/services/core/java/com/android/server/net/NetworkStatsService.java
@@ -544,7 +544,8 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         final int usedFlags = isRateLimitedForPoll(callingUid)
                 ? flags & (~NetworkStatsManager.FLAG_POLL_ON_OPEN)
                 : flags;
-        if ((usedFlags & NetworkStatsManager.FLAG_POLL_ON_OPEN) != 0) {
+        if ((usedFlags & (NetworkStatsManager.FLAG_POLL_ON_OPEN
+                | NetworkStatsManager.FLAG_POLL_FORCE)) != 0) {
             final long ident = Binder.clearCallingIdentity();
             try {
                 performPoll(FLAG_PERSIST_ALL);