Merge "[SP29] Send interface warning bytes to NetworkStatsProvider" am: a819b9fcf5 am: efe96175f6

Original change: https://android-review.googlesource.com/c/platform/frameworks/base/+/1470723

Change-Id: I6d799819362fd7d89b6be691d15c49314c867734
diff --git a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
index e9d5fe6..c127858 100644
--- a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
+++ b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
@@ -423,8 +423,8 @@
     private static final int MSG_LIMIT_REACHED = 5;
     private static final int MSG_RESTRICT_BACKGROUND_CHANGED = 6;
     private static final int MSG_ADVISE_PERSIST_THRESHOLD = 7;
-    private static final int MSG_UPDATE_INTERFACE_QUOTA = 10;
-    private static final int MSG_REMOVE_INTERFACE_QUOTA = 11;
+    private static final int MSG_UPDATE_INTERFACE_QUOTAS = 10;
+    private static final int MSG_REMOVE_INTERFACE_QUOTAS = 11;
     private static final int MSG_POLICIES_CHANGED = 13;
     private static final int MSG_RESET_FIREWALL_RULES_BY_UID = 15;
     private static final int MSG_SUBSCRIPTION_OVERRIDE = 16;
@@ -2036,33 +2036,44 @@
             final boolean hasWarning = policy.warningBytes != LIMIT_DISABLED;
             final boolean hasLimit = policy.limitBytes != LIMIT_DISABLED;
             long limitBytes = Long.MAX_VALUE;
-            if (hasLimit && policy.hasCycle()) {
+            long warningBytes = Long.MAX_VALUE;
+            if ((hasLimit || hasWarning) && policy.hasCycle()) {
                 final Pair<ZonedDateTime, ZonedDateTime> cycle = NetworkPolicyManager
                         .cycleIterator(policy).next();
                 final long start = cycle.first.toInstant().toEpochMilli();
                 final long end = cycle.second.toInstant().toEpochMilli();
                 final long totalBytes = getTotalBytes(policy.template, start, end);
 
-                if (policy.lastLimitSnooze < start) {
+                // If the limit notification is not snoozed, the limit quota needs to be calculated.
+                if (hasLimit && policy.lastLimitSnooze < start) {
                     // remaining "quota" bytes are based on total usage in
                     // current cycle. kernel doesn't like 0-byte rules, so we
                     // set 1-byte quota and disable the radio later.
                     limitBytes = Math.max(1, policy.limitBytes - totalBytes);
                 }
+
+                // If the warning notification was snoozed by user, or the service already knows
+                // it is over warning bytes, doesn't need to calculate warning bytes.
+                if (hasWarning && policy.lastWarningSnooze < start
+                        && !policy.isOverWarning(totalBytes)) {
+                    warningBytes = Math.max(1, policy.warningBytes - totalBytes);
+                }
             }
 
-            if (hasLimit || policy.metered) {
+            if (hasWarning || hasLimit || policy.metered) {
                 if (matchingIfaces.size() > 1) {
                     // TODO: switch to shared quota once NMS supports
                     Slog.w(TAG, "shared quota unsupported; generating rule for each iface");
                 }
 
-                // Set the interface limit. For interfaces which has no cycle, or metered with
-                // no policy limit, or snoozed limit notification; we still need to put iptables
-                // rule hooks to restrict apps for data saver, so push really high quota.
+                // Set the interface warning and limit. For interfaces which has no cycle,
+                // or metered with no policy quotas, or snoozed notification; we still need to put
+                // iptables rule hooks to restrict apps for data saver, so push really high quota.
+                // TODO: Push NetworkStatsProvider.QUOTA_UNLIMITED instead of Long.MAX_VALUE to
+                //  providers.
                 for (int j = matchingIfaces.size() - 1; j >= 0; j--) {
                     final String iface = matchingIfaces.valueAt(j);
-                    setInterfaceQuotaAsync(iface, limitBytes);
+                    setInterfaceQuotasAsync(iface, warningBytes, limitBytes);
                     newMeteredIfaces.add(iface);
                 }
             }
@@ -2085,7 +2096,7 @@
                 for (int j = matchingIfaces.size() - 1; j >= 0; j--) {
                     final String iface = matchingIfaces.valueAt(j);
                     if (!newMeteredIfaces.contains(iface)) {
-                        setInterfaceQuotaAsync(iface, Long.MAX_VALUE);
+                        setInterfaceQuotasAsync(iface, Long.MAX_VALUE, Long.MAX_VALUE);
                         newMeteredIfaces.add(iface);
                     }
                 }
@@ -2097,7 +2108,7 @@
             for (int i = mMeteredIfaces.size() - 1; i >= 0; i--) {
                 final String iface = mMeteredIfaces.valueAt(i);
                 if (!newMeteredIfaces.contains(iface)) {
-                    removeInterfaceQuotaAsync(iface);
+                    removeInterfaceQuotasAsync(iface);
                 }
             }
             mMeteredIfaces = newMeteredIfaces;
@@ -5038,19 +5049,20 @@
                     mNetworkStats.advisePersistThreshold(persistThreshold);
                     return true;
                 }
-                case MSG_UPDATE_INTERFACE_QUOTA: {
-                    final String iface = (String) msg.obj;
-                    // int params need to be stitched back into a long
-                    final long quota = ((long) msg.arg1 << 32) | (msg.arg2 & 0xFFFFFFFFL);
-                    removeInterfaceQuota(iface);
-                    setInterfaceQuota(iface, quota);
-                    mNetworkStats.setStatsProviderLimitAsync(iface, quota);
+                case MSG_UPDATE_INTERFACE_QUOTAS: {
+                    final IfaceQuotas val = (IfaceQuotas) msg.obj;
+                    // TODO: Consider set a new limit before removing the original one.
+                    removeInterfaceLimit(val.iface);
+                    setInterfaceLimit(val.iface, val.limit);
+                    mNetworkStats.setStatsProviderWarningAndLimitAsync(val.iface, val.warning,
+                            val.limit);
                     return true;
                 }
-                case MSG_REMOVE_INTERFACE_QUOTA: {
+                case MSG_REMOVE_INTERFACE_QUOTAS: {
                     final String iface = (String) msg.obj;
-                    removeInterfaceQuota(iface);
-                    mNetworkStats.setStatsProviderLimitAsync(iface, QUOTA_UNLIMITED);
+                    removeInterfaceLimit(iface);
+                    mNetworkStats.setStatsProviderWarningAndLimitAsync(iface, QUOTA_UNLIMITED,
+                            QUOTA_UNLIMITED);
                     return true;
                 }
                 case MSG_RESET_FIREWALL_RULES_BY_UID: {
@@ -5198,15 +5210,32 @@
         }
     }
 
-    private void setInterfaceQuotaAsync(String iface, long quotaBytes) {
-        // long quotaBytes split up into two ints to fit in message
-        mHandler.obtainMessage(MSG_UPDATE_INTERFACE_QUOTA, (int) (quotaBytes >> 32),
-                (int) (quotaBytes & 0xFFFFFFFF), iface).sendToTarget();
+    private static final class IfaceQuotas {
+        @NonNull public final String iface;
+        // Warning and limit bytes of interface qutoas, could be QUOTA_UNLIMITED or Long.MAX_VALUE
+        // if not set. 0 is not acceptable since kernel doesn't like 0-byte rules.
+        public final long warning;
+        public final long limit;
+
+        private IfaceQuotas(@NonNull String iface, long warning, long limit) {
+            this.iface = iface;
+            this.warning = warning;
+            this.limit = limit;
+        }
     }
 
-    private void setInterfaceQuota(String iface, long quotaBytes) {
+    private void setInterfaceQuotasAsync(@NonNull String iface,
+            long warningBytes, long limitBytes) {
+        mHandler.obtainMessage(MSG_UPDATE_INTERFACE_QUOTAS,
+                new IfaceQuotas(iface, warningBytes, limitBytes)).sendToTarget();
+    }
+
+    private void setInterfaceLimit(String iface, long limitBytes) {
         try {
-            mNetworkManager.setInterfaceQuota(iface, quotaBytes);
+            // For legacy design the data warning is covered by global alert, where the
+            // kernel will notify upper layer for a small amount of change of traffic
+            // statistics. Thus, passing warning is not needed.
+            mNetworkManager.setInterfaceQuota(iface, limitBytes);
         } catch (IllegalStateException e) {
             Log.wtf(TAG, "problem setting interface quota", e);
         } catch (RemoteException e) {
@@ -5214,11 +5243,11 @@
         }
     }
 
-    private void removeInterfaceQuotaAsync(String iface) {
-        mHandler.obtainMessage(MSG_REMOVE_INTERFACE_QUOTA, iface).sendToTarget();
+    private void removeInterfaceQuotasAsync(String iface) {
+        mHandler.obtainMessage(MSG_REMOVE_INTERFACE_QUOTAS, iface).sendToTarget();
     }
 
-    private void removeInterfaceQuota(String iface) {
+    private void removeInterfaceLimit(String iface) {
         try {
             mNetworkManager.removeInterfaceQuota(iface);
         } catch (IllegalStateException e) {
diff --git a/services/core/java/com/android/server/net/NetworkStatsManagerInternal.java b/services/core/java/com/android/server/net/NetworkStatsManagerInternal.java
index 0cb0bc2c..0e9a9da 100644
--- a/services/core/java/com/android/server/net/NetworkStatsManagerInternal.java
+++ b/services/core/java/com/android/server/net/NetworkStatsManagerInternal.java
@@ -37,8 +37,9 @@
     public abstract void forceUpdate();
 
     /**
-     * Set the quota limit to all registered custom network stats providers.
+     * Set the warning and limit to all registered custom network stats providers.
      * Note that invocation of any interface will be sent to all providers.
      */
-    public abstract void setStatsProviderLimitAsync(@NonNull String iface, long quota);
+    public abstract void setStatsProviderWarningAndLimitAsync(@NonNull String iface, long warning,
+            long limit);
 }
diff --git a/services/core/java/com/android/server/net/NetworkStatsService.java b/services/core/java/com/android/server/net/NetworkStatsService.java
index 445a425..19f5e3c 100644
--- a/services/core/java/com/android/server/net/NetworkStatsService.java
+++ b/services/core/java/com/android/server/net/NetworkStatsService.java
@@ -1693,11 +1693,14 @@
         }
 
         @Override
-        public void setStatsProviderLimitAsync(@NonNull String iface, long quota) {
-            if (LOGV) Slog.v(TAG, "setStatsProviderLimitAsync(" + iface + "," + quota + ")");
-            // TODO: Set warning accordingly.
+        public void setStatsProviderWarningAndLimitAsync(
+                @NonNull String iface, long warning, long limit) {
+            if (LOGV) {
+                Slog.v(TAG, "setStatsProviderWarningAndLimitAsync("
+                        + iface + "," + warning + "," + limit + ")");
+            }
             invokeForAllStatsProviderCallbacks((cb) -> cb.mProvider.onSetWarningAndLimit(iface,
-                    NetworkStatsProvider.QUOTA_UNLIMITED, quota));
+                    warning, limit));
         }
     }
 
diff --git a/services/tests/servicestests/src/com/android/server/net/NetworkPolicyManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/net/NetworkPolicyManagerServiceTest.java
index d405113..100d3ea 100644
--- a/services/tests/servicestests/src/com/android/server/net/NetworkPolicyManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/net/NetworkPolicyManagerServiceTest.java
@@ -1773,57 +1773,75 @@
                 true);
     }
 
-    /**
-     * Test that when StatsProvider triggers limit reached, new limit will be calculated and
-     * re-armed.
-     */
-    @Test
-    public void testStatsProviderLimitReached() throws Exception {
-        final int CYCLE_DAY = 15;
-
-        final NetworkStats stats = new NetworkStats(0L, 1);
+    private void increaseMockedTotalBytes(NetworkStats stats, long rxBytes, long txBytes) {
         stats.insertEntry(TEST_IFACE, UID_A, SET_ALL, TAG_NONE,
-                2999, 1, 2000, 1, 0);
+                rxBytes, 1, txBytes, 1, 0);
         when(mStatsService.getNetworkTotalBytes(any(), anyLong(), anyLong()))
                 .thenReturn(stats.getTotalBytes());
         when(mStatsService.getNetworkUidBytes(any(), anyLong(), anyLong()))
                 .thenReturn(stats);
+    }
+
+    private void triggerOnStatsProviderWarningOrLimitReached() throws InterruptedException {
+        final NetworkPolicyManagerInternal npmi = LocalServices
+                .getService(NetworkPolicyManagerInternal.class);
+        npmi.onStatsProviderWarningOrLimitReached("TEST");
+        // Wait for processing of MSG_STATS_PROVIDER_WARNING_OR_LIMIT_REACHED.
+        postMsgAndWaitForCompletion();
+        verify(mStatsService).forceUpdate();
+        // Wait for processing of MSG_*_INTERFACE_QUOTAS.
+        postMsgAndWaitForCompletion();
+    }
+
+    /**
+     * Test that when StatsProvider triggers warning and limit reached, new quotas will be
+     * calculated and re-armed.
+     */
+    @Test
+    public void testStatsProviderWarningAndLimitReached() throws Exception {
+        final int CYCLE_DAY = 15;
+
+        final NetworkStats stats = new NetworkStats(0L, 1);
+        increaseMockedTotalBytes(stats, 2999, 2000);
 
         // Get active mobile network in place
         expectMobileDefaults();
         mService.updateNetworks();
-        verify(mStatsService).setStatsProviderLimitAsync(TEST_IFACE, Long.MAX_VALUE);
+        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, Long.MAX_VALUE,
+                Long.MAX_VALUE);
 
-        // Set limit to 10KB.
+        // Set warning to 7KB and limit to 10KB.
         setNetworkPolicies(new NetworkPolicy(
-                sTemplateMobileAll, CYCLE_DAY, TIMEZONE_UTC, WARNING_DISABLED, 10000L,
-                true));
+                sTemplateMobileAll, CYCLE_DAY, TIMEZONE_UTC, 7000L, 10000L, true));
         postMsgAndWaitForCompletion();
 
-        // Verifies that remaining quota is set to providers.
-        verify(mStatsService).setStatsProviderLimitAsync(TEST_IFACE, 10000L - 4999L);
-
+        // Verifies that remaining quotas are set to providers.
+        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, 2001L, 5001L);
         reset(mStatsService);
 
-        // Increase the usage.
-        stats.insertEntry(TEST_IFACE, UID_A, SET_ALL, TAG_NONE,
-                1000, 1, 999, 1, 0);
-        when(mStatsService.getNetworkTotalBytes(any(), anyLong(), anyLong()))
-                .thenReturn(stats.getTotalBytes());
-        when(mStatsService.getNetworkUidBytes(any(), anyLong(), anyLong()))
-                .thenReturn(stats);
+        // Increase the usage and simulates that limit reached fires earlier by provider,
+        // but actually the quota is not yet reached. Verifies that the limit reached leads to
+        // a force update and new quotas should be set.
+        increaseMockedTotalBytes(stats, 1000, 999);
+        triggerOnStatsProviderWarningOrLimitReached();
+        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, 2L, 3002L);
+        reset(mStatsService);
 
-        // Simulates that limit reached fires earlier by provider, but actually the quota is not
-        // yet reached.
-        final NetworkPolicyManagerInternal npmi = LocalServices
-                .getService(NetworkPolicyManagerInternal.class);
-        npmi.onStatsProviderWarningOrLimitReached("TEST");
+        // Increase the usage and simulate warning reached, the new warning should be unlimited
+        // since service will disable warning quota to stop lower layer from keep triggering
+        // warning reached event.
+        increaseMockedTotalBytes(stats, 1000L, 1000);
+        triggerOnStatsProviderWarningOrLimitReached();
+        verify(mStatsService).setStatsProviderWarningAndLimitAsync(
+                TEST_IFACE, Long.MAX_VALUE, 1002L);
+        reset(mStatsService);
 
-        // Verifies that the limit reached leads to a force update and new limit should be set.
-        postMsgAndWaitForCompletion();
-        verify(mStatsService).forceUpdate();
-        postMsgAndWaitForCompletion();
-        verify(mStatsService).setStatsProviderLimitAsync(TEST_IFACE, 10000L - 4999L - 1999L);
+        // Increase the usage that over the warning and limit, the new limit should set to 1 to
+        // block the network traffic.
+        increaseMockedTotalBytes(stats, 1000L, 1000);
+        triggerOnStatsProviderWarningOrLimitReached();
+        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, Long.MAX_VALUE, 1L);
+        reset(mStatsService);
     }
 
     /**