Merge changes If51b6676,Ia6863a70,I802d2316 am: 76851bd64d am: 4740b85389

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

Change-Id: Ibb4cfb5be191efe6e6e821eb17d905fb96b02ebb

6 files changed, 276 insertions, 264 deletions

diff --git a/packages/ConnectivityT/framework-t/src/android/app/usage/NetworkStatsManager.java b/packages/ConnectivityT/framework-t/src/android/app/usage/NetworkStatsManager.java
index cc7b2a5a30e2..f74edb1a01e5 100644
--- a/packages/ConnectivityT/framework-t/src/android/app/usage/NetworkStatsManager.java
+++ b/packages/ConnectivityT/framework-t/src/android/app/usage/NetworkStatsManager.java
@@ -142,7 +142,15 @@ public class NetworkStatsManager {
         setAugmentWithSubscriptionPlan(true);
     }
 
-    /** @hide */
+    /**
+     * Set poll on open flag to indicate the poll is needed before service gets statistics
+     * result. This is default enabled. However, for any non-privileged caller, the poll might
+     * be omitted in case of rate limiting.
+     *
+     * @param pollOnOpen true if poll is needed.
+     * @hide
+     */
+    // @SystemApi(client = MODULE_LIBRARIES)
     public void setPollOnOpen(boolean pollOnOpen) {
         if (pollOnOpen) {
             mFlags |= FLAG_POLL_ON_OPEN;
@@ -863,4 +871,74 @@ public class NetworkStatsManager {
             return msg.getData().getParcelable(key);
         }
     }
+
+    /**
+     * Mark given UID as being in foreground for stats purposes.
+     *
+     * @hide
+     */
+    // @SystemApi
+    @RequiresPermission(anyOf = {
+            NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
+            android.Manifest.permission.NETWORK_STACK})
+    public void setUidForeground(int uid, boolean uidForeground) {
+        try {
+            mService.setUidForeground(uid, uidForeground);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Advise persistence threshold; may be overridden internally.
+     *
+     * @hide
+     */
+    // @SystemApi
+    @RequiresPermission(anyOf = {
+            NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
+            android.Manifest.permission.NETWORK_STACK})
+    public void advisePersistThreshold(long thresholdBytes) {
+        try {
+            mService.advisePersistThreshold(thresholdBytes);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Force update of statistics.
+     *
+     * @hide
+     */
+    // @SystemApi
+    @RequiresPermission(anyOf = {
+            NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
+            android.Manifest.permission.NETWORK_STACK})
+    public void forceUpdate() {
+        try {
+            mService.forceUpdate();
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Set the warning and limit to all registered custom network stats providers.
+     * Note that invocation of any interface will be sent to all providers.
+     *
+     * @hide
+     */
+    // @SystemApi
+    @RequiresPermission(anyOf = {
+            NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
+            android.Manifest.permission.NETWORK_STACK})
+    public void setStatsProviderWarningAndLimitAsync(@NonNull String iface, long warning,
+            long limit) {
+        try {
+            mService.setStatsProviderWarningAndLimitAsync(iface, warning, limit);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
 }
diff --git a/packages/ConnectivityT/framework-t/src/android/net/INetworkStatsService.aidl b/packages/ConnectivityT/framework-t/src/android/net/INetworkStatsService.aidl
index 12937b5cb2c7..a4babb543dbd 100644
--- a/packages/ConnectivityT/framework-t/src/android/net/INetworkStatsService.aidl
+++ b/packages/ConnectivityT/framework-t/src/android/net/INetworkStatsService.aidl
@@ -94,4 +94,16 @@ interface INetworkStatsService {
     /** Registers a network stats provider */
     INetworkStatsProviderCallback registerNetworkStatsProvider(String tag,
             in INetworkStatsProvider provider);
+
+    /** Mark given UID as being in foreground for stats purposes. */
+    void setUidForeground(int uid, boolean uidForeground);
+
+    /** Advise persistence threshold; may be overridden internally. */
+    void advisePersistThreshold(long thresholdBytes);
+
+    /**
+     * Set the warning and limit to all registered custom network stats providers.
+     * Note that invocation of any interface will be sent to all providers.
+     */
+     void setStatsProviderWarningAndLimitAsync(String iface, long warning, long limit);
 }
diff --git a/packages/ConnectivityT/service/src/com/android/server/net/NetworkStatsManagerInternal.java b/packages/ConnectivityT/service/src/com/android/server/net/NetworkStatsManagerInternal.java
deleted file mode 100644
index 0e9a9da6804b..000000000000
--- a/packages/ConnectivityT/service/src/com/android/server/net/NetworkStatsManagerInternal.java
+++ /dev/null
@@ -1,45 +0,0 @@
-/*
- * Copyright (C) 2018 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 com.android.server.net;
-
-import android.annotation.NonNull;
-import android.net.NetworkStats;
-import android.net.NetworkTemplate;
-
-public abstract class NetworkStatsManagerInternal {
-    /** Return network layer usage total for traffic that matches template. */
-    public abstract long getNetworkTotalBytes(NetworkTemplate template, long start, long end);
-
-    /** Return network layer usage per-UID for traffic that matches template. */
-    public abstract NetworkStats getNetworkUidBytes(NetworkTemplate template, long start, long end);
-
-    /** Mark given UID as being in foreground for stats purposes. */
-    public abstract void setUidForeground(int uid, boolean uidForeground);
-
-    /** Advise persistance threshold; may be overridden internally. */
-    public abstract void advisePersistThreshold(long thresholdBytes);
-
-    /** Force update of statistics. */
-    public abstract void forceUpdate();
-
-    /**
-     * 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 setStatsProviderWarningAndLimitAsync(@NonNull String iface, long warning,
-            long limit);
-}
diff --git a/packages/ConnectivityT/service/src/com/android/server/net/NetworkStatsService.java b/packages/ConnectivityT/service/src/com/android/server/net/NetworkStatsService.java
index efc8c555c3c7..e15acf3ef616 100644
--- a/packages/ConnectivityT/service/src/com/android/server/net/NetworkStatsService.java
+++ b/packages/ConnectivityT/service/src/com/android/server/net/NetworkStatsService.java
@@ -431,7 +431,6 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
                 new DefaultNetworkStatsSettings(context), new NetworkStatsFactory(netd),
                 new NetworkStatsObservers(), getDefaultSystemDir(), getDefaultBaseDir(),
                 new Dependencies());
-        service.registerLocalService();
 
         return service;
     }
@@ -512,11 +511,6 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         }
     }
 
-    private void registerLocalService() {
-        LocalServices.addService(NetworkStatsManagerInternal.class,
-                new NetworkStatsManagerInternalImpl());
-    }
-
     /**
      * Observer that watches for {@link INetdUnsolicitedEventListener} alerts.
      */
@@ -1007,7 +1001,8 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
     }
 
     @VisibleForTesting
-    void setUidForeground(int uid, boolean uidForeground) {
+    public void setUidForeground(int uid, boolean uidForeground) {
+        PermissionUtils.enforceNetworkStackPermission(mContext);
         synchronized (mStatsLock) {
             final int set = uidForeground ? SET_FOREGROUND : SET_DEFAULT;
             final int oldSet = mActiveUidCounterSet.get(uid, SET_DEFAULT);
@@ -1043,7 +1038,7 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
 
     @Override
     public void forceUpdate() {
-        mContext.enforceCallingOrSelfPermission(READ_NETWORK_USAGE_HISTORY, TAG);
+        PermissionUtils.enforceNetworkStackPermission(mContext);
 
         final long token = Binder.clearCallingIdentity();
         try {
@@ -1053,7 +1048,9 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         }
     }
 
-    private void advisePersistThreshold(long thresholdBytes) {
+    /** Advise persistence threshold; may be overridden internally. */
+    public void advisePersistThreshold(long thresholdBytes) {
+        PermissionUtils.enforceNetworkStackPermission(mContext);
         // clamp threshold into safe range
         mPersistThreshold = NetworkStatsUtils.constrain(thresholdBytes,
                 128 * KB_IN_BYTES, 2 * MB_IN_BYTES);
@@ -1690,52 +1687,19 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         removeUidsLocked(CollectionUtils.toIntArray(uids));
     }
 
-    private class NetworkStatsManagerInternalImpl extends NetworkStatsManagerInternal {
-        @Override
-        public long getNetworkTotalBytes(NetworkTemplate template, long start, long end) {
-            Trace.traceBegin(TRACE_TAG_NETWORK, "getNetworkTotalBytes");
-            try {
-                return NetworkStatsService.this.getNetworkTotalBytes(template, start, end);
-            } finally {
-                Trace.traceEnd(TRACE_TAG_NETWORK);
-            }
-        }
-
-        @Override
-        public NetworkStats getNetworkUidBytes(NetworkTemplate template, long start, long end) {
-            Trace.traceBegin(TRACE_TAG_NETWORK, "getNetworkUidBytes");
-            try {
-                return NetworkStatsService.this.getNetworkUidBytes(template, start, end);
-            } finally {
-                Trace.traceEnd(TRACE_TAG_NETWORK);
-            }
-        }
-
-        @Override
-        public void setUidForeground(int uid, boolean uidForeground) {
-            NetworkStatsService.this.setUidForeground(uid, uidForeground);
-        }
-
-        @Override
-        public void advisePersistThreshold(long thresholdBytes) {
-            NetworkStatsService.this.advisePersistThreshold(thresholdBytes);
-        }
-
-        @Override
-        public void forceUpdate() {
-            NetworkStatsService.this.forceUpdate();
-        }
-
-        @Override
-        public void setStatsProviderWarningAndLimitAsync(
-                @NonNull String iface, long warning, long limit) {
-            if (LOGV) {
-                Log.v(TAG, "setStatsProviderWarningAndLimitAsync("
-                        + iface + "," + warning + "," + limit + ")");
-            }
-            invokeForAllStatsProviderCallbacks((cb) -> cb.mProvider.onSetWarningAndLimit(iface,
-                    warning, limit));
+    /**
+     * 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 void setStatsProviderWarningAndLimitAsync(
+            @NonNull String iface, long warning, long limit) {
+        PermissionUtils.enforceNetworkStackPermission(mContext);
+        if (LOGV) {
+            Log.v(TAG, "setStatsProviderWarningAndLimitAsync("
+                    + iface + "," + warning + "," + limit + ")");
         }
+        invokeForAllStatsProviderCallbacks((cb) -> cb.mProvider.onSetWarningAndLimit(iface,
+                warning, limit));
     }
 
     @Override
diff --git a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
index d752b68d650b..e1cbdb7d22fe 100644
--- a/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
+++ b/services/core/java/com/android/server/net/NetworkPolicyManagerService.java
@@ -151,6 +151,8 @@ import android.app.IUidObserver;
 import android.app.Notification;
 import android.app.NotificationManager;
 import android.app.PendingIntent;
+import android.app.usage.NetworkStats;
+import android.app.usage.NetworkStatsManager;
 import android.app.usage.UsageStatsManagerInternal;
 import android.content.BroadcastReceiver;
 import android.content.ComponentName;
@@ -181,7 +183,6 @@ import android.net.NetworkRequest;
 import android.net.NetworkSpecifier;
 import android.net.NetworkStack;
 import android.net.NetworkStateSnapshot;
-import android.net.NetworkStats;
 import android.net.NetworkTemplate;
 import android.net.TelephonyNetworkSpecifier;
 import android.net.TrafficStats;
@@ -443,7 +444,7 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
 
     private final Context mContext;
     private final IActivityManager mActivityManager;
-    private NetworkStatsManagerInternal mNetworkStats;
+    private NetworkStatsManager mNetworkStats;
     private final INetworkManagementService mNetworkManager;
     private UsageStatsManagerInternal mUsageStats;
     private AppStandbyInternal mAppStandby;
@@ -455,6 +456,8 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
     private ConnectivityManager mConnManager;
     private PowerManagerInternal mPowerManagerInternal;
     private PowerWhitelistManager mPowerWhitelistManager;
+    @NonNull
+    private final Dependencies mDeps;
 
     /** Current cached value of the current Battery Saver mode's setting for restrict background. */
     @GuardedBy("mUidRulesFirstLock")
@@ -706,7 +709,7 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
     public NetworkPolicyManagerService(Context context, IActivityManager activityManager,
             INetworkManagementService networkManagement) {
         this(context, activityManager, networkManagement, AppGlobals.getPackageManager(),
-                getDefaultClock(), getDefaultSystemDir(), false);
+                getDefaultClock(), getDefaultSystemDir(), false, new Dependencies(context));
     }
 
     private static @NonNull File getDefaultSystemDir() {
@@ -718,9 +721,59 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
                 Clock.systemUTC());
     }
 
+    static class Dependencies {
+        final Context mContext;
+        final NetworkStatsManager mNetworkStatsManager;
+        Dependencies(Context context) {
+            mContext = context;
+            mNetworkStatsManager = mContext.getSystemService(NetworkStatsManager.class);
+            // Query stats from NetworkStatsService will trigger a poll by default.
+            // But since NPMS listens stats updated event, and will query stats
+            // after the event. A polling -> updated -> query -> polling loop will be introduced
+            // if polls on open. Hence, while NPMS manages it's poll requests explicitly, set
+            // flag to false to prevent a polling loop.
+            mNetworkStatsManager.setPollOnOpen(false);
+        }
+
+        long getNetworkTotalBytes(NetworkTemplate template, long start, long end) {
+            Trace.traceBegin(TRACE_TAG_NETWORK, "getNetworkTotalBytes");
+            try {
+                final NetworkStats.Bucket ret = mNetworkStatsManager
+                        .querySummaryForDevice(template, start, end);
+                return ret.getRxBytes() + ret.getTxBytes();
+            } catch (RuntimeException e) {
+                Slog.w(TAG, "Failed to read network stats: " + e);
+                return 0;
+            } finally {
+                Trace.traceEnd(TRACE_TAG_NETWORK);
+            }
+        }
+
+        @NonNull
+        List<NetworkStats.Bucket> getNetworkUidBytes(
+                @NonNull NetworkTemplate template, long start, long end) {
+            Trace.traceBegin(TRACE_TAG_NETWORK, "getNetworkUidBytes");
+            final List<NetworkStats.Bucket> buckets = new ArrayList<>();
+            try {
+                final NetworkStats stats = mNetworkStatsManager.querySummary(template, start, end);
+                while (stats.hasNextBucket()) {
+                    final NetworkStats.Bucket bucket = new NetworkStats.Bucket();
+                    stats.getNextBucket(bucket);
+                    buckets.add(bucket);
+                }
+            } catch (RuntimeException e) {
+                Slog.w(TAG, "Failed to read network stats: " + e);
+            } finally {
+                Trace.traceEnd(TRACE_TAG_NETWORK);
+            }
+            return buckets;
+        }
+    }
+
+    @VisibleForTesting
     public NetworkPolicyManagerService(Context context, IActivityManager activityManager,
             INetworkManagementService networkManagement, IPackageManager pm, Clock clock,
-            File systemDir, boolean suppressDefaultPolicy) {
+            File systemDir, boolean suppressDefaultPolicy, Dependencies deps) {
         mContext = Objects.requireNonNull(context, "missing context");
         mActivityManager = Objects.requireNonNull(activityManager, "missing activityManager");
         mNetworkManager = Objects.requireNonNull(networkManagement, "missing networkManagement");
@@ -741,10 +794,12 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
         mUidEventHandler = new Handler(mUidEventThread.getLooper(), mUidEventHandlerCallback);
 
         mSuppressDefaultPolicy = suppressDefaultPolicy;
+        mDeps = Objects.requireNonNull(deps, "missing Dependencies");
 
         mPolicyFile = new AtomicFile(new File(systemDir, "netpolicy.xml"), "net-policy");
 
         mAppOps = context.getSystemService(AppOpsManager.class);
+        mNetworkStats = context.getSystemService(NetworkStatsManager.class);
         mMultipathPolicyTracker = new MultipathPolicyTracker(mContext, mHandler);
         // Expose private service for system components to use.
         LocalServices.addService(NetworkPolicyManagerInternal.class,
@@ -844,7 +899,6 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
 
             mUsageStats = LocalServices.getService(UsageStatsManagerInternal.class);
             mAppStandby = LocalServices.getService(AppStandbyInternal.class);
-            mNetworkStats = LocalServices.getService(NetworkStatsManagerInternal.class);
 
             synchronized (mUidRulesFirstLock) {
                 synchronized (mNetworkPoliciesSecondLock) {
@@ -1167,21 +1221,34 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
     };
 
     /**
-     * Receiver that watches for {@link INetworkStatsService} updates, which we
+     * Receiver that watches for {@link NetworkStatsManager} updates, which we
      * use to check against {@link NetworkPolicy#warningBytes}.
      */
-    final private BroadcastReceiver mStatsReceiver = new BroadcastReceiver() {
+    private final NetworkStatsBroadcastReceiver mStatsReceiver =
+            new NetworkStatsBroadcastReceiver();
+    private class NetworkStatsBroadcastReceiver extends BroadcastReceiver {
+        private boolean mIsAnyIntentReceived = false;
         @Override
         public void onReceive(Context context, Intent intent) {
             // on background handler thread, and verified
             // READ_NETWORK_USAGE_HISTORY permission above.
 
+            mIsAnyIntentReceived = true;
+
             synchronized (mNetworkPoliciesSecondLock) {
                 updateNetworkRulesNL();
                 updateNetworkEnabledNL();
                 updateNotificationsNL();
             }
         }
+
+        /**
+         * Return whether any {@code ACTION_NETWORK_STATS_UPDATED} intent is received.
+         * Used to determine if NetworkStatsService is ready.
+         */
+        public boolean isAnyIntentReceived() {
+            return mIsAnyIntentReceived;
+        }
     };
 
     /**
@@ -1405,15 +1472,17 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
         long maxBytes = 0;
         int maxUid = 0;
 
-        final NetworkStats stats = getNetworkUidBytes(template, start, end);
-        NetworkStats.Entry entry = null;
-        for (int i = 0; i < stats.size(); i++) {
-            entry = stats.getValues(i, entry);
-            final long bytes = entry.rxBytes + entry.txBytes;
+        // Skip if not ready. NetworkStatsService will block public API calls until it is
+        // ready. To prevent NPMS be blocked on that, skip and fail fast instead.
+        if (!mStatsReceiver.isAnyIntentReceived()) return null;
+
+        final List<NetworkStats.Bucket> stats = mDeps.getNetworkUidBytes(template, start, end);
+        for (final NetworkStats.Bucket entry : stats) {
+            final long bytes = entry.getRxBytes() + entry.getTxBytes();
             totalBytes += bytes;
             if (bytes > maxBytes) {
                 maxBytes = bytes;
-                maxUid = entry.uid;
+                maxUid = entry.getUid();
             }
         }
 
@@ -5398,25 +5467,10 @@ public class NetworkPolicyManagerService extends INetworkPolicyManager.Stub {
 
     @Deprecated
     private long getTotalBytes(NetworkTemplate template, long start, long end) {
-        return getNetworkTotalBytes(template, start, end);
-    }
-
-    private long getNetworkTotalBytes(NetworkTemplate template, long start, long end) {
-        try {
-            return mNetworkStats.getNetworkTotalBytes(template, start, end);
-        } catch (RuntimeException e) {
-            Slog.w(TAG, "Failed to read network stats: " + e);
-            return 0;
-        }
-    }
-
-    private NetworkStats getNetworkUidBytes(NetworkTemplate template, long start, long end) {
-        try {
-            return mNetworkStats.getNetworkUidBytes(template, start, end);
-        } catch (RuntimeException e) {
-            Slog.w(TAG, "Failed to read network stats: " + e);
-            return new NetworkStats(SystemClock.elapsedRealtime(), 0);
-        }
+        // Skip if not ready. NetworkStatsService will block public API calls until it is
+        // ready. To prevent NPMS be blocked on that, skip and fail fast instead.
+        if (!mStatsReceiver.isAnyIntentReceived()) return 0;
+        return mDeps.getNetworkTotalBytes(template, start, end);
     }
 
     private boolean isBandwidthControlEnabled() {
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 66df0fecad60..9a6f61e7c6cf 100644
--- a/services/tests/servicestests/src/com/android/server/net/NetworkPolicyManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/net/NetworkPolicyManagerServiceTest.java
@@ -49,12 +49,8 @@ import static android.net.NetworkPolicyManager.blockedReasonsToString;
 import static android.net.NetworkPolicyManager.uidPoliciesToString;
 import static android.net.NetworkPolicyManager.uidRulesToString;
 import static android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK;
-import static android.net.NetworkStats.IFACE_ALL;
 import static android.net.NetworkStats.METERED_NO;
 import static android.net.NetworkStats.METERED_YES;
-import static android.net.NetworkStats.SET_ALL;
-import static android.net.NetworkStats.TAG_ALL;
-import static android.net.NetworkStats.TAG_NONE;
 import static android.net.NetworkTemplate.buildTemplateCarrierMetered;
 import static android.net.NetworkTemplate.buildTemplateWifi;
 import static android.net.TrafficStats.MB_IN_BYTES;
@@ -75,6 +71,7 @@ import static com.android.server.net.NetworkPolicyManagerService.TYPE_LIMIT_SNOO
 import static com.android.server.net.NetworkPolicyManagerService.TYPE_RAPID;
 import static com.android.server.net.NetworkPolicyManagerService.TYPE_WARNING;
 import static com.android.server.net.NetworkPolicyManagerService.UidBlockedState.getEffectiveBlockedReasons;
+import static com.android.server.net.NetworkStatsService.ACTION_NETWORK_STATS_UPDATED;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -108,6 +105,8 @@ import android.app.IActivityManager;
 import android.app.IUidObserver;
 import android.app.Notification;
 import android.app.NotificationManager;
+import android.app.usage.NetworkStats;
+import android.app.usage.NetworkStatsManager;
 import android.app.usage.UsageStatsManagerInternal;
 import android.content.Context;
 import android.content.Intent;
@@ -125,8 +124,6 @@ import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkPolicy;
 import android.net.NetworkStateSnapshot;
-import android.net.NetworkStats;
-import android.net.NetworkStatsHistory;
 import android.net.NetworkTemplate;
 import android.net.TelephonyNetworkSpecifier;
 import android.net.wifi.WifiInfo;
@@ -138,7 +135,6 @@ import android.os.PowerManagerInternal;
 import android.os.PowerSaveState;
 import android.os.RemoteException;
 import android.os.SimpleClock;
-import android.os.SystemClock;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.platform.test.annotations.Presubmit;
@@ -263,12 +259,13 @@ public class NetworkPolicyManagerServiceTest {
     private @Mock CarrierConfigManager mCarrierConfigManager;
     private @Mock TelephonyManager mTelephonyManager;
     private @Mock UserManager mUserManager;
+    private @Mock NetworkStatsManager mStatsManager;
+    private TestDependencies mDeps;
 
     private ArgumentCaptor<ConnectivityManager.NetworkCallback> mNetworkCallbackCaptor =
             ArgumentCaptor.forClass(ConnectivityManager.NetworkCallback.class);
 
     private ActivityManagerInternal mActivityManagerInternal;
-    private NetworkStatsManagerInternal mStatsService;
 
     private IUidObserver mUidObserver;
     private INetworkManagementEventObserver mNetworkObserver;
@@ -335,8 +332,47 @@ public class NetworkPolicyManagerServiceTest {
                 .setBatterySaverEnabled(false).build();
         final PowerManagerInternal pmInternal = addLocalServiceMock(PowerManagerInternal.class);
         when(pmInternal.getLowPowerState(anyInt())).thenReturn(state);
+    }
+
+    private class TestDependencies extends NetworkPolicyManagerService.Dependencies {
+        private final SparseArray<NetworkStats.Bucket> mMockedStats = new SparseArray<>();
 
-        mStatsService = addLocalServiceMock(NetworkStatsManagerInternal.class);
+        TestDependencies(Context context) {
+            super(context);
+        }
+
+        @Override
+        long getNetworkTotalBytes(NetworkTemplate template, long start, long end) {
+            int total = 0;
+            for (int i = 0; i < mMockedStats.size(); i++) {
+                NetworkStats.Bucket bucket = mMockedStats.valueAt(i);
+                total += bucket.getRxBytes() + bucket.getTxBytes();
+            }
+            return total;
+        }
+
+        @Override
+        List<NetworkStats.Bucket> getNetworkUidBytes(NetworkTemplate template, long start,
+                long end) {
+            final List<NetworkStats.Bucket> ret = new ArrayList<>();
+            for (int i = 0; i < mMockedStats.size(); i++) {
+                ret.add(mMockedStats.valueAt(i));
+            }
+            return ret;
+        }
+
+        private void setMockedTotalBytes(int uid, long rxBytes, long txBytes) {
+            final NetworkStats.Bucket bucket = mock(NetworkStats.Bucket.class);
+            when(bucket.getUid()).thenReturn(uid);
+            when(bucket.getRxBytes()).thenReturn(rxBytes);
+            when(bucket.getTxBytes()).thenReturn(txBytes);
+            mMockedStats.set(uid, bucket);
+        }
+
+        private void increaseMockedTotalBytes(int uid, long rxBytes, long txBytes) {
+            final NetworkStats.Bucket bucket = mMockedStats.get(uid);
+            setMockedTotalBytes(uid, bucket.getRxBytes() + rxBytes, bucket.getTxBytes() + txBytes);
+        }
     }
 
     @Before
@@ -376,6 +412,8 @@ public class NetworkPolicyManagerServiceTest {
                         return mConnManager;
                     case Context.USER_SERVICE:
                         return mUserManager;
+                    case Context.NETWORK_STATS_SERVICE:
+                        return mStatsManager;
                     default:
                         return super.getSystemService(name);
                 }
@@ -400,8 +438,9 @@ public class NetworkPolicyManagerServiceTest {
         }).when(mActivityManager).registerUidObserver(any(), anyInt(), anyInt(), any(String.class));
 
         mFutureIntent = newRestrictBackgroundChangedFuture();
+        mDeps = new TestDependencies(mServiceContext);
         mService = new NetworkPolicyManagerService(mServiceContext, mActivityManager,
-                mNetworkManager, mIpm, mClock, mPolicyDir, true);
+                mNetworkManager, mIpm, mClock, mPolicyDir, true, mDeps);
         mService.bindConnectivityManager();
         mPolicyListener = new NetworkPolicyListenerAnswer(mService);
 
@@ -456,6 +495,9 @@ public class NetworkPolicyManagerServiceTest {
         verify(mNetworkManager).registerObserver(networkObserver.capture());
         mNetworkObserver = networkObserver.getValue();
 
+        // Simulate NetworkStatsService broadcast stats updated to signal its readiness.
+        mServiceContext.sendBroadcast(new Intent(ACTION_NETWORK_STATS_UPDATED));
+
         NetworkPolicy defaultPolicy = mService.buildDefaultCarrierPolicy(0, "");
         mDefaultWarningBytes = defaultPolicy.warningBytes;
         mDefaultLimitBytes = defaultPolicy.limitBytes;
@@ -479,7 +521,6 @@ public class NetworkPolicyManagerServiceTest {
         LocalServices.removeServiceForTest(DeviceIdleInternal.class);
         LocalServices.removeServiceForTest(AppStandbyInternal.class);
         LocalServices.removeServiceForTest(UsageStatsManagerInternal.class);
-        LocalServices.removeServiceForTest(NetworkStatsManagerInternal.class);
     }
 
     @After
@@ -1108,10 +1149,7 @@ public class NetworkPolicyManagerServiceTest {
         when(mConnManager.getAllNetworkStateSnapshots()).thenReturn(snapshots);
 
         // pretend that 512 bytes total have happened
-        stats = new NetworkStats(getElapsedRealtime(), 1)
-                .insertEntry(TEST_IFACE, 256L, 2L, 256L, 2L);
-        when(mStatsService.getNetworkTotalBytes(sTemplateWifi, CYCLE_START, CYCLE_END))
-                .thenReturn(stats.getTotalBytes());
+        mDeps.setMockedTotalBytes(UID_A, 256L, 256L);
 
         mPolicyListener.expect().onMeteredIfacesChanged(any());
         setNetworkPolicies(new NetworkPolicy(
@@ -1124,26 +1162,6 @@ public class NetworkPolicyManagerServiceTest {
 
     @Test
     public void testNotificationWarningLimitSnooze() throws Exception {
-        // Create a place to store fake usage
-        final NetworkStatsHistory history = new NetworkStatsHistory(TimeUnit.HOURS.toMillis(1));
-        final NetworkStats stats = new NetworkStats(SystemClock.elapsedRealtime(), 0);
-        when(mStatsService.getNetworkTotalBytes(any(), anyLong(), anyLong()))
-                .thenAnswer(new Answer<Long>() {
-                    @Override
-                    public Long answer(InvocationOnMock invocation) throws Throwable {
-                        final NetworkStatsHistory.Entry entry = history.getValues(
-                                invocation.getArgument(1), invocation.getArgument(2), null);
-                        return entry.rxBytes + entry.txBytes;
-                    }
-                });
-        when(mStatsService.getNetworkUidBytes(any(), anyLong(), anyLong()))
-                .thenAnswer(new Answer<NetworkStats>() {
-                    @Override
-                    public NetworkStats answer(InvocationOnMock invocation) throws Throwable {
-                        return stats;
-                    }
-                });
-
         // Get active mobile network in place
         expectMobileDefaults();
         mService.updateNetworks();
@@ -1161,9 +1179,7 @@ public class NetworkPolicyManagerServiceTest {
 
         // Normal usage means no notification
         {
-            history.clear();
-            history.recordData(start, end,
-                    new NetworkStats.Entry(DataUnit.MEGABYTES.toBytes(360), 0L, 0L, 0L, 0));
+            mDeps.setMockedTotalBytes(UID_A, DataUnit.MEGABYTES.toBytes(360), 0);
 
             reset(mTelephonyManager, mNetworkManager, mNotifManager);
             TelephonyManager tmSub = expectMobileDefaults();
@@ -1178,9 +1194,7 @@ public class NetworkPolicyManagerServiceTest {
 
         // Push over warning
         {
-            history.clear();
-            history.recordData(start, end,
-                    new NetworkStats.Entry(DataUnit.MEGABYTES.toBytes(1799), 0L, 0L, 0L, 0));
+            mDeps.setMockedTotalBytes(UID_A, DataUnit.MEGABYTES.toBytes(1799), 0);
 
             reset(mTelephonyManager, mNetworkManager, mNotifManager);
             TelephonyManager tmSub = expectMobileDefaults();
@@ -1196,9 +1210,7 @@ public class NetworkPolicyManagerServiceTest {
 
         // Push over warning, but with a config that isn't from an identified carrier
         {
-            history.clear();
-            history.recordData(start, end,
-                    new NetworkStats.Entry(DataUnit.MEGABYTES.toBytes(1799), 0L, 0L, 0L, 0));
+            mDeps.setMockedTotalBytes(UID_A, DataUnit.MEGABYTES.toBytes(1799), 0);
 
             reset(mTelephonyManager, mNetworkManager, mNotifManager);
             TelephonyManager tmSub = expectMobileDefaults();
@@ -1215,9 +1227,7 @@ public class NetworkPolicyManagerServiceTest {
 
         // Push over limit
         {
-            history.clear();
-            history.recordData(start, end,
-                    new NetworkStats.Entry(DataUnit.MEGABYTES.toBytes(1810), 0L, 0L, 0L, 0));
+            mDeps.setMockedTotalBytes(UID_A, DataUnit.MEGABYTES.toBytes(1810), 0);
 
             reset(mTelephonyManager, mNetworkManager, mNotifManager);
             TelephonyManager tmSub = expectMobileDefaults();
@@ -1248,26 +1258,6 @@ public class NetworkPolicyManagerServiceTest {
 
     @Test
     public void testNotificationRapid() throws Exception {
-        // Create a place to store fake usage
-        final NetworkStatsHistory history = new NetworkStatsHistory(TimeUnit.HOURS.toMillis(1));
-        final NetworkStats stats = new NetworkStats(SystemClock.elapsedRealtime(), 0);
-        when(mStatsService.getNetworkTotalBytes(any(), anyLong(), anyLong()))
-                .thenAnswer(new Answer<Long>() {
-                    @Override
-                    public Long answer(InvocationOnMock invocation) throws Throwable {
-                        final NetworkStatsHistory.Entry entry = history.getValues(
-                                invocation.getArgument(1), invocation.getArgument(2), null);
-                        return entry.rxBytes + entry.txBytes;
-                    }
-                });
-        when(mStatsService.getNetworkUidBytes(any(), anyLong(), anyLong()))
-                .thenAnswer(new Answer<NetworkStats>() {
-                    @Override
-                    public NetworkStats answer(InvocationOnMock invocation) throws Throwable {
-                        return stats;
-                    }
-                });
-
         // Get active mobile network in place
         expectMobileDefaults();
         mService.updateNetworks();
@@ -1285,9 +1275,7 @@ public class NetworkPolicyManagerServiceTest {
 
         // Using 20% data in 20% time is normal
         {
-            history.clear();
-            history.recordData(start, end,
-                    new NetworkStats.Entry(DataUnit.MEGABYTES.toBytes(360), 0L, 0L, 0L, 0));
+            mDeps.setMockedTotalBytes(UID_A, DataUnit.MEGABYTES.toBytes(360), 0);
 
             reset(mNotifManager);
             mService.updateNetworks();
@@ -1297,16 +1285,9 @@ public class NetworkPolicyManagerServiceTest {
         // Using 80% data in 20% time is alarming; but spread equally among
         // three UIDs means we get generic alert
         {
-            history.clear();
-            history.recordData(start, end,
-                    new NetworkStats.Entry(DataUnit.MEGABYTES.toBytes(1440), 0L, 0L, 0L, 0));
-            stats.clear();
-            stats.insertEntry(IFACE_ALL, UID_A, SET_ALL, TAG_ALL,
-                    DataUnit.MEGABYTES.toBytes(480), 0, 0, 0, 0);
-            stats.insertEntry(IFACE_ALL, UID_B, SET_ALL, TAG_ALL,
-                    DataUnit.MEGABYTES.toBytes(480), 0, 0, 0, 0);
-            stats.insertEntry(IFACE_ALL, UID_C, SET_ALL, TAG_ALL,
-                    DataUnit.MEGABYTES.toBytes(480), 0, 0, 0, 0);
+            mDeps.setMockedTotalBytes(UID_A, DataUnit.MEGABYTES.toBytes(480), 0);
+            mDeps.setMockedTotalBytes(UID_B, DataUnit.MEGABYTES.toBytes(480), 0);
+            mDeps.setMockedTotalBytes(UID_C, DataUnit.MEGABYTES.toBytes(480), 0);
 
             reset(mNotifManager);
             mService.updateNetworks();
@@ -1325,14 +1306,9 @@ public class NetworkPolicyManagerServiceTest {
         // Using 80% data in 20% time is alarming; but mostly done by one UID
         // means we get specific alert
         {
-            history.clear();
-            history.recordData(start, end,
-                    new NetworkStats.Entry(DataUnit.MEGABYTES.toBytes(1440), 0L, 0L, 0L, 0));
-            stats.clear();
-            stats.insertEntry(IFACE_ALL, UID_A, SET_ALL, TAG_ALL,
-                    DataUnit.MEGABYTES.toBytes(960), 0, 0, 0, 0);
-            stats.insertEntry(IFACE_ALL, UID_B, SET_ALL, TAG_ALL,
-                    DataUnit.MEGABYTES.toBytes(480), 0, 0, 0, 0);
+            mDeps.setMockedTotalBytes(UID_A, DataUnit.MEGABYTES.toBytes(960), 0);
+            mDeps.setMockedTotalBytes(UID_B, DataUnit.MEGABYTES.toBytes(480), 0);
+            mDeps.setMockedTotalBytes(UID_C, 0, 0);
 
             reset(mNotifManager);
             mService.updateNetworks();
@@ -1362,13 +1338,10 @@ public class NetworkPolicyManagerServiceTest {
 
         // bring up wifi network with metered policy
         snapshots = List.of(buildWifi());
-        stats = new NetworkStats(getElapsedRealtime(), 1)
-                .insertEntry(TEST_IFACE, 0L, 0L, 0L, 0L);
+        mDeps.setMockedTotalBytes(UID_A, 0L, 0L);
 
         {
             when(mConnManager.getAllNetworkStateSnapshots()).thenReturn(snapshots);
-            when(mStatsService.getNetworkTotalBytes(sTemplateWifi, TIME_FEB_15,
-                    currentTimeMillis())).thenReturn(stats.getTotalBytes());
 
             mPolicyListener.expect().onMeteredIfacesChanged(any());
             setNetworkPolicies(new NetworkPolicy(
@@ -1647,18 +1620,6 @@ public class NetworkPolicyManagerServiceTest {
         final NetworkPolicyManagerInternal internal = LocalServices
                 .getService(NetworkPolicyManagerInternal.class);
 
-        // Create a place to store fake usage
-        final NetworkStatsHistory history = new NetworkStatsHistory(TimeUnit.HOURS.toMillis(1));
-        final NetworkStats stats = new NetworkStats(SystemClock.elapsedRealtime(), 0);
-        when(mStatsService.getNetworkTotalBytes(any(), anyLong(), anyLong()))
-                .thenAnswer(invocation -> {
-                    final NetworkStatsHistory.Entry entry = history.getValues(
-                            invocation.getArgument(1), invocation.getArgument(2), null);
-                    return entry.rxBytes + entry.txBytes;
-                });
-        when(mStatsService.getNetworkUidBytes(any(), anyLong(), anyLong()))
-                .thenReturn(stats);
-
         // Get active mobile network in place
         expectMobileDefaults();
         mService.updateNetworks();
@@ -1669,9 +1630,7 @@ public class NetworkPolicyManagerServiceTest {
         setCurrentTimeMillis(end);
 
         // Get some data usage in place
-        history.clear();
-        history.recordData(start, end,
-                new NetworkStats.Entry(DataUnit.MEGABYTES.toBytes(360), 0L, 0L, 0L, 0));
+        mDeps.setMockedTotalBytes(UID_A, DataUnit.MEGABYTES.toBytes(360), 0);
 
         // No data plan
         {
@@ -1786,20 +1745,11 @@ public class NetworkPolicyManagerServiceTest {
                 true);
     }
 
-    private void increaseMockedTotalBytes(NetworkStats stats, long rxBytes, long txBytes) {
-        stats.insertEntry(TEST_IFACE, UID_A, SET_ALL, TAG_NONE,
-                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 {
         mService.onStatsProviderWarningOrLimitReached();
         // Wait for processing of MSG_STATS_PROVIDER_WARNING_OR_LIMIT_REACHED.
         postMsgAndWaitForCompletion();
-        verify(mStatsService).forceUpdate();
+        verify(mStatsManager).forceUpdate();
         // Wait for processing of MSG_*_INTERFACE_QUOTAS.
         postMsgAndWaitForCompletion();
     }
@@ -1812,13 +1762,12 @@ public class NetworkPolicyManagerServiceTest {
     public void testStatsProviderWarningAndLimitReached() throws Exception {
         final int CYCLE_DAY = 15;
 
-        final NetworkStats stats = new NetworkStats(0L, 1);
-        increaseMockedTotalBytes(stats, 2999, 2000);
+        mDeps.setMockedTotalBytes(UID_A, 2999, 2000);
 
         // Get active mobile network in place
         expectMobileDefaults();
         mService.updateNetworks();
-        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, Long.MAX_VALUE,
+        verify(mStatsManager).setStatsProviderWarningAndLimitAsync(TEST_IFACE, Long.MAX_VALUE,
                 Long.MAX_VALUE);
 
         // Set warning to 7KB and limit to 10KB.
@@ -1828,32 +1777,32 @@ public class NetworkPolicyManagerServiceTest {
         postMsgAndWaitForCompletion();
 
         // Verifies that remaining quotas are set to providers.
-        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, 2001L, 5001L);
-        reset(mStatsService);
+        verify(mStatsManager).setStatsProviderWarningAndLimitAsync(TEST_IFACE, 2001L, 5001L);
+        reset(mStatsManager);
 
         // 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);
+        mDeps.increaseMockedTotalBytes(UID_A, 1000, 999);
         triggerOnStatsProviderWarningOrLimitReached();
-        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, 2L, 3002L);
-        reset(mStatsService);
+        verify(mStatsManager).setStatsProviderWarningAndLimitAsync(TEST_IFACE, 2L, 3002L);
+        reset(mStatsManager);
 
         // 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);
+        mDeps.increaseMockedTotalBytes(UID_A, 1000L, 1000);
         triggerOnStatsProviderWarningOrLimitReached();
-        verify(mStatsService).setStatsProviderWarningAndLimitAsync(
+        verify(mStatsManager).setStatsProviderWarningAndLimitAsync(
                 TEST_IFACE, Long.MAX_VALUE, 1002L);
-        reset(mStatsService);
+        reset(mStatsManager);
 
         // 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);
+        mDeps.increaseMockedTotalBytes(UID_A, 1000L, 1000);
         triggerOnStatsProviderWarningOrLimitReached();
-        verify(mStatsService).setStatsProviderWarningAndLimitAsync(TEST_IFACE, Long.MAX_VALUE, 1L);
-        reset(mStatsService);
+        verify(mStatsManager).setStatsProviderWarningAndLimitAsync(TEST_IFACE, Long.MAX_VALUE, 1L);
+        reset(mStatsManager);
     }
 
     private void enableRestrictedMode(boolean enable) throws Exception {
@@ -2143,7 +2092,7 @@ public class NetworkPolicyManagerServiceTest {
     }
 
     private void verifyAdvisePersistThreshold() throws Exception {
-        verify(mStatsService).advisePersistThreshold(anyLong());
+        verify(mStatsManager).advisePersistThreshold(anyLong());
     }
 
     private static class TestAbstractFuture<T> extends AbstractFuture<T> {