Merge "[SP21] Address comments for API council review about aosp/1172143"

api/system-current.txt

@@ -1240,7 +1240,8 @@ package android.app.usage {
   }
 
   public class NetworkStatsManager {
-    method @NonNull @RequiresPermission(anyOf={android.Manifest.permission.NETWORK_STATS_PROVIDER, android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK}) public android.net.netstats.provider.NetworkStatsProviderCallback registerNetworkStatsProvider(@NonNull String, @NonNull android.net.netstats.provider.AbstractNetworkStatsProvider);
+    method @NonNull @RequiresPermission(anyOf={android.Manifest.permission.NETWORK_STATS_PROVIDER, android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK}) public void registerNetworkStatsProvider(@NonNull String, @NonNull android.net.netstats.provider.NetworkStatsProvider);
+    method @NonNull @RequiresPermission(anyOf={android.Manifest.permission.NETWORK_STATS_PROVIDER, android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK}) public void unregisterNetworkStatsProvider(@NonNull android.net.netstats.provider.NetworkStatsProvider);
   }
 
   public static final class UsageEvents.Event {
@@ -5181,21 +5182,17 @@ package android.net.metrics {
 
 package android.net.netstats.provider {
 
-  public abstract class AbstractNetworkStatsProvider {
+  public abstract class NetworkStatsProvider {
-    ctor public AbstractNetworkStatsProvider();
+    ctor public NetworkStatsProvider();
-    method public abstract void requestStatsUpdate(int);
+    method public void notifyAlertReached();
-    method public abstract void setAlert(long);
+    method public void notifyLimitReached();
-    method public abstract void setLimit(@NonNull String, long);
+    method public void notifyStatsUpdated(int, @NonNull android.net.NetworkStats, @NonNull android.net.NetworkStats);
+    method public abstract void onRequestStatsUpdate(int);
+    method public abstract void onSetAlert(long);
+    method public abstract void onSetLimit(@NonNull String, long);
     field public static final int QUOTA_UNLIMITED = -1; // 0xffffffff
   }
 
-  public class NetworkStatsProviderCallback {
-    method public void onAlertReached();
-    method public void onLimitReached();
-    method public void onStatsUpdated(int, @NonNull android.net.NetworkStats, @NonNull android.net.NetworkStats);
-    method public void unregister();
-  }
-
 }
 
 package android.net.sip {

core/java/android/app/usage/NetworkStatsManager.java

@@ -31,9 +31,8 @@ import android.net.INetworkStatsService;
 import android.net.NetworkIdentity;
 import android.net.NetworkStack;
 import android.net.NetworkTemplate;
-import android.net.netstats.provider.AbstractNetworkStatsProvider;
+import android.net.netstats.provider.INetworkStatsProviderCallback;
-import android.net.netstats.provider.NetworkStatsProviderCallback;
+import android.net.netstats.provider.NetworkStatsProvider;
-import android.net.netstats.provider.NetworkStatsProviderWrapper;
 import android.os.Binder;
 import android.os.Handler;
 import android.os.Looper;
@@ -528,34 +527,53 @@ public class NetworkStatsManager {
 
     /**
      * Registers a custom provider of {@link android.net.NetworkStats} to provide network statistics
-     * to the system. To unregister, invoke {@link NetworkStatsProviderCallback#unregister()}.
+     * to the system. To unregister, invoke {@link #unregisterNetworkStatsProvider}.
      * Note that no de-duplication of statistics between providers is performed, so each provider
-     * must only report network traffic that is not being reported by any other provider.
+     * must only report network traffic that is not being reported by any other provider. Also note
+     * that the provider cannot be re-registered after unregistering.
      *
      * @param tag a human readable identifier of the custom network stats provider. This is only
      *            used for debugging.
-     * @param provider the subclass of {@link AbstractNetworkStatsProvider} that needs to be
+     * @param provider the subclass of {@link NetworkStatsProvider} that needs to be
      *                 registered to the system.
-     * @return a {@link NetworkStatsProviderCallback}, which can be used to report events to the
-     *         system or unregister the provider.
      * @hide
      */
     @SystemApi
     @RequiresPermission(anyOf = {
             android.Manifest.permission.NETWORK_STATS_PROVIDER,
             NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK})
-    @NonNull public NetworkStatsProviderCallback registerNetworkStatsProvider(
+    @NonNull public void registerNetworkStatsProvider(
             @NonNull String tag,
-            @NonNull AbstractNetworkStatsProvider provider) {
+            @NonNull NetworkStatsProvider provider) {
         try {
-            final NetworkStatsProviderWrapper wrapper = new NetworkStatsProviderWrapper(provider);
+            if (provider.getProviderCallbackBinder() != null) {
-            return new NetworkStatsProviderCallback(
+                throw new IllegalArgumentException("provider is already registered");
-                    mService.registerNetworkStatsProvider(tag, wrapper));
+            }
+            final INetworkStatsProviderCallback cbBinder =
+                    mService.registerNetworkStatsProvider(tag, provider.getProviderBinder());
+            provider.setProviderCallbackBinder(cbBinder);
+        } catch (RemoteException e) {
+            e.rethrowAsRuntimeException();
+        }
+    }
+
+    /**
+     * Unregisters an instance of {@link NetworkStatsProvider}.
+     *
+     * @param provider the subclass of {@link NetworkStatsProvider} that needs to be
+     *                 unregistered to the system.
+     * @hide
+     */
+    @SystemApi
+    @RequiresPermission(anyOf = {
+            android.Manifest.permission.NETWORK_STATS_PROVIDER,
+            NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK})
+    @NonNull public void unregisterNetworkStatsProvider(@NonNull NetworkStatsProvider provider) {
+        try {
+            provider.getProviderCallbackBinderOrThrow().unregister();
         } catch (RemoteException e) {
             e.rethrowAsRuntimeException();
         }
-        // Unreachable code, but compiler doesn't know about it.
-        return null;
     }
 
     private static NetworkTemplate createTemplate(int networkType, String subscriberId) {

core/java/android/net/netstats/provider/AbstractNetworkStatsProvider.java

@@ -1,70 +0,0 @@
-/*
- * Copyright (C) 2020 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 android.net.netstats.provider;
-
-import android.annotation.NonNull;
-import android.annotation.SystemApi;
-import android.net.NetworkStats;
-
-/**
- * A base class that allows external modules to implement a custom network statistics provider.
- * @hide
- */
-@SystemApi
-public abstract class AbstractNetworkStatsProvider {
-    /**
-     * A value used by {@link #setLimit} and {@link #setAlert} indicates there is no limit.
-     */
-    public static final int QUOTA_UNLIMITED = -1;
-
-    /**
-     * Called by {@code NetworkStatsService} when global polling is needed. Custom
-     * implementation of providers MUST respond to it by calling
-     * {@link NetworkStatsProviderCallback#onStatsUpdated} within one minute. Responding
-     * later than this may cause the stats to be dropped.
-     *
-     * @param token a positive number identifying the new state of the system under which
-     *              {@link NetworkStats} have to be gathered from now on. When this is called,
-     *              custom implementations of providers MUST report the latest stats with the
-     *              previous token, under which stats were being gathered so far.
-     */
-    public abstract void requestStatsUpdate(int token);
-
-    /**
-     * Called by {@code NetworkStatsService} when setting the interface quota for the specified
-     * upstream interface. When this is called, the custom implementation should block all egress
-     * packets on the {@code iface} associated with the provider when {@code quotaBytes} bytes have
-     * been reached, and MUST respond to it by calling
-     * {@link NetworkStatsProviderCallback#onLimitReached()}.
-     *
-     * @param iface the interface requiring the operation.
-     * @param quotaBytes the quota defined as the number of bytes, starting from zero and counting
-     *                   from now. A value of {@link #QUOTA_UNLIMITED} indicates there is no limit.
-     */
-    public abstract void setLimit(@NonNull String iface, long quotaBytes);
-
-    /**
-     * Called by {@code NetworkStatsService} when setting the alert bytes. Custom implementations
-     * MUST call {@link NetworkStatsProviderCallback#onAlertReached()} when {@code quotaBytes} bytes
-     * have been reached. Unlike {@link #setLimit(String, long)}, the custom implementation should
-     * not block all egress packets.
-     *
-     * @param quotaBytes the quota defined as the number of bytes, starting from zero and counting
-     *                   from now. A value of {@link #QUOTA_UNLIMITED} indicates there is no alert.
-     */
-    public abstract void setAlert(long quotaBytes);
-}

core/java/android/net/netstats/provider/INetworkStatsProvider.aidl

@@ -22,7 +22,7 @@ package android.net.netstats.provider;
  * @hide
  */
 oneway interface INetworkStatsProvider {
-    void requestStatsUpdate(int token);
+    void onRequestStatsUpdate(int token);
-    void setLimit(String iface, long quotaBytes);
+    void onSetLimit(String iface, long quotaBytes);
-    void setAlert(long quotaBytes);
+    void onSetAlert(long quotaBytes);
 }

core/java/android/net/netstats/provider/INetworkStatsProviderCallback.aidl

@@ -24,8 +24,8 @@ import android.net.NetworkStats;
  * @hide
  */
 oneway interface INetworkStatsProviderCallback {
-    void onStatsUpdated(int token, in NetworkStats ifaceStats, in NetworkStats uidStats);
+    void notifyStatsUpdated(int token, in NetworkStats ifaceStats, in NetworkStats uidStats);
-    void onAlertReached();
+    void notifyAlertReached();
-    void onLimitReached();
+    void notifyLimitReached();
     void unregister();
 }

core/java/android/net/netstats/provider/NetworkStatsProvider.java

@@ -0,0 +1,195 @@
+/*
+ * Copyright (C) 2020 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 android.net.netstats.provider;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.SystemApi;
+import android.net.NetworkStats;
+import android.os.RemoteException;
+
+/**
+ * A base class that allows external modules to implement a custom network statistics provider.
+ * @hide
+ */
+@SystemApi
+public abstract class NetworkStatsProvider {
+    /**
+     * A value used by {@link #onSetLimit} and {@link #onSetAlert} indicates there is no limit.
+     */
+    public static final int QUOTA_UNLIMITED = -1;
+
+    @NonNull private final INetworkStatsProvider mProviderBinder =
+            new INetworkStatsProvider.Stub() {
+
+        @Override
+        public void onRequestStatsUpdate(int token) {
+            NetworkStatsProvider.this.onRequestStatsUpdate(token);
+        }
+
+        @Override
+        public void onSetLimit(String iface, long quotaBytes) {
+            NetworkStatsProvider.this.onSetLimit(iface, quotaBytes);
+        }
+
+        @Override
+        public void onSetAlert(long quotaBytes) {
+            NetworkStatsProvider.this.onSetAlert(quotaBytes);
+        }
+    };
+
+    // The binder given by the service when successfully registering. Only null before registering,
+    // never null once non-null.
+    @Nullable
+    private INetworkStatsProviderCallback mProviderCbBinder;
+
+    /**
+     * Return the binder invoked by the service and redirect function calls to the overridden
+     * methods.
+     * @hide
+     */
+    @NonNull
+    public INetworkStatsProvider getProviderBinder() {
+        return mProviderBinder;
+    }
+
+    /**
+     * Store the binder that was returned by the service when successfully registering. Note that
+     * the provider cannot be re-registered. Hence this method can only be called once per provider.
+     *
+     * @hide
+     */
+    public void setProviderCallbackBinder(@NonNull INetworkStatsProviderCallback binder) {
+        if (mProviderCbBinder != null) {
+            throw new IllegalArgumentException("provider is already registered");
+        }
+        mProviderCbBinder = binder;
+    }
+
+    /**
+     * Get the binder that was returned by the service when successfully registering. Or null if the
+     * provider was never registered.
+     *
+     * @hide
+     */
+    @Nullable
+    public INetworkStatsProviderCallback getProviderCallbackBinder() {
+        return mProviderCbBinder;
+    }
+
+    /**
+     * Get the binder that was returned by the service when successfully registering. Throw an
+     * {@link IllegalStateException} if the provider is not registered.
+     *
+     * @hide
+     */
+    @NonNull
+    public INetworkStatsProviderCallback getProviderCallbackBinderOrThrow() {
+        if (mProviderCbBinder == null) {
+            throw new IllegalStateException("the provider is not registered");
+        }
+        return mProviderCbBinder;
+    }
+
+    /**
+     * Notify the system of new network statistics.
+     *
+     * Send the network statistics recorded since the last call to {@link #notifyStatsUpdated}. Must
+     * be called as soon as possible after {@link NetworkStatsProvider#onRequestStatsUpdate(int)}
+     * being called. Responding later increases the probability stats will be dropped. The
+     * provider can also call this whenever it wants to reports new stats for any reason.
+     * Note that the system will not necessarily immediately propagate the statistics to
+     * reflect the update.
+     *
+     * @param token the token under which these stats were gathered. Providers can call this method
+     *              with the current token as often as they want, until the token changes.
+     *              {@see NetworkStatsProvider#onRequestStatsUpdate()}
+     * @param ifaceStats the {@link NetworkStats} per interface to be reported.
+     *                   The provider should not include any traffic that is already counted by
+     *                   kernel interface counters.
+     * @param uidStats the same stats as above, but counts {@link NetworkStats}
+     *                 per uid.
+     */
+    public void notifyStatsUpdated(int token, @NonNull NetworkStats ifaceStats,
+            @NonNull NetworkStats uidStats) {
+        try {
+            getProviderCallbackBinderOrThrow().notifyStatsUpdated(token, ifaceStats, uidStats);
+        } catch (RemoteException e) {
+            e.rethrowAsRuntimeException();
+        }
+    }
+
+    /**
+     * Notify system that the quota set by {@code onSetAlert} has been reached.
+     */
+    public void notifyAlertReached() {
+        try {
+            getProviderCallbackBinderOrThrow().notifyAlertReached();
+        } catch (RemoteException e) {
+            e.rethrowAsRuntimeException();
+        }
+    }
+
+    /**
+     * Notify system that the quota set by {@code onSetLimit} has been reached.
+     */
+    public void notifyLimitReached() {
+        try {
+            getProviderCallbackBinderOrThrow().notifyLimitReached();
+        } catch (RemoteException e) {
+            e.rethrowAsRuntimeException();
+        }
+    }
+
+    /**
+     * Called by {@code NetworkStatsService} when it requires to know updated stats.
+     * The provider MUST respond by calling {@link #notifyStatsUpdated} as soon as possible.
+     * Responding later increases the probability stats will be dropped. Memory allowing, the
+     * system will try to take stats into account up to one minute after calling
+     * {@link #onRequestStatsUpdate}.
+     *
+     * @param token a positive number identifying the new state of the system under which
+     *              {@link NetworkStats} have to be gathered from now on. When this is called,
+     *              custom implementations of providers MUST tally and report the latest stats with
+     *              the previous token, under which stats were being gathered so far.
+     */
+    public abstract void onRequestStatsUpdate(int token);
+
+    /**
+     * Called by {@code NetworkStatsService} when setting the interface quota for the specified
+     * upstream interface. When this is called, the custom implementation should block all egress
+     * packets on the {@code iface} associated with the provider when {@code quotaBytes} bytes have
+     * been reached, and MUST respond to it by calling
+     * {@link NetworkStatsProvider#notifyLimitReached()}.
+     *
+     * @param iface the interface requiring the operation.
+     * @param quotaBytes the quota defined as the number of bytes, starting from zero and counting
+     *                   from now. A value of {@link #QUOTA_UNLIMITED} indicates there is no limit.
+     */
+    public abstract void onSetLimit(@NonNull String iface, long quotaBytes);
+
+    /**
+     * Called by {@code NetworkStatsService} when setting the alert bytes. Custom implementations
+     * MUST call {@link NetworkStatsProvider#notifyAlertReached()} when {@code quotaBytes} bytes
+     * have been reached. Unlike {@link #onSetLimit(String, long)}, the custom implementation should
+     * not block all egress packets.
+     *
+     * @param quotaBytes the quota defined as the number of bytes, starting from zero and counting
+     *                   from now. A value of {@link #QUOTA_UNLIMITED} indicates there is no alert.
+     */
+    public abstract void onSetAlert(long quotaBytes);
+}

core/java/android/net/netstats/provider/NetworkStatsProviderCallback.java

@@ -1,98 +0,0 @@
-/*
- * Copyright (C) 2020 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 android.net.netstats.provider;
-
-import android.annotation.NonNull;
-import android.annotation.SuppressLint;
-import android.annotation.SystemApi;
-import android.net.NetworkStats;
-import android.os.RemoteException;
-
-/**
- * A callback class that allows callers to report events to the system.
- * @hide
- */
-@SystemApi
-@SuppressLint("CallbackMethodName")
-public class NetworkStatsProviderCallback {
-    @NonNull private final INetworkStatsProviderCallback mBinder;
-
-    /** @hide */
-    public NetworkStatsProviderCallback(@NonNull INetworkStatsProviderCallback binder) {
-        mBinder = binder;
-    }
-
-    /**
-     * Notify the system of new network statistics.
-     *
-     * Send the network statistics recorded since the last call to {@link #onStatsUpdated}. Must be
-     * called within one minute of {@link AbstractNetworkStatsProvider#requestStatsUpdate(int)}
-     * being called. The provider can also call this whenever it wants to reports new stats for any
-     * reason. Note that the system will not necessarily immediately propagate the statistics to
-     * reflect the update.
-     *
-     * @param token the token under which these stats were gathered. Providers can call this method
-     *              with the current token as often as they want, until the token changes.
-     *              {@see AbstractNetworkStatsProvider#requestStatsUpdate()}
-     * @param ifaceStats the {@link NetworkStats} per interface to be reported.
-     *                   The provider should not include any traffic that is already counted by
-     *                   kernel interface counters.
-     * @param uidStats the same stats as above, but counts {@link NetworkStats}
-     *                 per uid.
-     */
-    public void onStatsUpdated(int token, @NonNull NetworkStats ifaceStats,
-            @NonNull NetworkStats uidStats) {
-        try {
-            mBinder.onStatsUpdated(token, ifaceStats, uidStats);
-        } catch (RemoteException e) {
-            e.rethrowAsRuntimeException();
-        }
-    }
-
-    /**
-     * Notify system that the quota set by {@code setAlert} has been reached.
-     */
-    public void onAlertReached() {
-        try {
-            mBinder.onAlertReached();
-        } catch (RemoteException e) {
-            e.rethrowAsRuntimeException();
-        }
-    }
-
-    /**
-     * Notify system that the quota set by {@code setLimit} has been reached.
-     */
-    public void onLimitReached() {
-        try {
-            mBinder.onLimitReached();
-        } catch (RemoteException e) {
-            e.rethrowAsRuntimeException();
-        }
-    }
-
-    /**
-     * Unregister the provider and the referencing callback.
-     */
-    public void unregister() {
-        try {
-            mBinder.unregister();
-        } catch (RemoteException e) {
-            // Ignore error.
-        }
-    }
-}

core/java/android/net/netstats/provider/NetworkStatsProviderWrapper.java

@@ -1,48 +0,0 @@
-/*
- * Copyright (C) 2020 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 android.net.netstats.provider;
-
-import android.annotation.NonNull;
-
-/**
- * A wrapper class of {@link INetworkStatsProvider} that hides the binder interface from exposing
- * to outer world.
- *
- * @hide
- */
-public class NetworkStatsProviderWrapper extends INetworkStatsProvider.Stub {
-    @NonNull final AbstractNetworkStatsProvider mProvider;
-
-    public NetworkStatsProviderWrapper(AbstractNetworkStatsProvider provider) {
-        mProvider = provider;
-    }
-
-    @Override
-    public void requestStatsUpdate(int token) {
-        mProvider.requestStatsUpdate(token);
-    }
-
-    @Override
-    public void setLimit(@NonNull String iface, long quotaBytes) {
-        mProvider.setLimit(iface, quotaBytes);
-    }
-
-    @Override
-    public void setAlert(long quotaBytes) {
-        mProvider.setAlert(quotaBytes);
-    }
-}

packages/Tethering/src/com/android/server/connectivity/tethering/OffloadController.java

@@ -39,8 +39,7 @@ import android.net.RouteInfo;
 import android.net.netlink.ConntrackMessage;
 import android.net.netlink.NetlinkConstants;
 import android.net.netlink.NetlinkSocket;
-import android.net.netstats.provider.AbstractNetworkStatsProvider;
+import android.net.netstats.provider.NetworkStatsProvider;
-import android.net.netstats.provider.NetworkStatsProviderCallback;
 import android.net.util.SharedLog;
 import android.os.Handler;
 import android.provider.Settings;
@@ -89,8 +88,8 @@ public class OffloadController {
     private final Handler mHandler;
     private final OffloadHardwareInterface mHwInterface;
     private final ContentResolver mContentResolver;
-    private final @NonNull OffloadTetheringStatsProvider mStatsProvider;
+    @Nullable
-    private final @Nullable NetworkStatsProviderCallback mStatsProviderCb;
+    private final OffloadTetheringStatsProvider mStatsProvider;
     private final SharedLog mLog;
     private final HashMap<String, LinkProperties> mDownstreams;
     private boolean mConfigInitialized;
@@ -124,19 +123,18 @@ public class OffloadController {
         mHandler = h;
         mHwInterface = hwi;
         mContentResolver = contentResolver;
-        mStatsProvider = new OffloadTetheringStatsProvider();
         mLog = log.forSubComponent(TAG);
         mDownstreams = new HashMap<>();
         mExemptPrefixes = new HashSet<>();
         mLastLocalPrefixStrs = new HashSet<>();
-        NetworkStatsProviderCallback providerCallback = null;
+        OffloadTetheringStatsProvider provider = new OffloadTetheringStatsProvider();
         try {
-            providerCallback = nsm.registerNetworkStatsProvider(
+            nsm.registerNetworkStatsProvider(getClass().getSimpleName(), provider);
-                    getClass().getSimpleName(), mStatsProvider);
         } catch (RuntimeException e) {
             Log.wtf(TAG, "Cannot register offload stats provider: " + e);
+            provider = null;
         }
-        mStatsProviderCb = providerCallback;
+        mStatsProvider = provider;
     }
 
     /** Start hardware offload. */
@@ -185,7 +183,7 @@ public class OffloadController {
                         // and we need to synchronize stats and limits between
                         // software and hardware forwarding.
                         updateStatsForAllUpstreams();
-                        mStatsProvider.pushTetherStats();
+                        if (mStatsProvider != null) mStatsProvider.pushTetherStats();
                     }
 
                     @Override
@@ -198,7 +196,7 @@ public class OffloadController {
                         // limits set take into account any software tethering
                         // traffic that has been happening in the meantime.
                         updateStatsForAllUpstreams();
-                        mStatsProvider.pushTetherStats();
+                        if (mStatsProvider != null) mStatsProvider.pushTetherStats();
                         // [2] (Re)Push all state.
                         computeAndPushLocalPrefixes(UpdateType.FORCE);
                         pushAllDownstreamState();
@@ -217,10 +215,12 @@ public class OffloadController {
                         // TODO: rev the HAL so that it provides an interface name.
 
                         updateStatsForCurrentUpstream();
+                        if (mStatsProvider != null) {
                             mStatsProvider.pushTetherStats();
-                        // Push stats to service does not cause the service react to it immediately.
+                            // Push stats to service does not cause the service react to it
-                        // Inform the service about limit reached.
+                            // immediately. Inform the service about limit reached.
-                        if (mStatsProviderCb != null) mStatsProviderCb.onLimitReached();
+                            mStatsProvider.notifyLimitReached();
+                        }
                     }
 
                     @Override
@@ -263,13 +263,17 @@ public class OffloadController {
     }
 
     @VisibleForTesting
-    class OffloadTetheringStatsProvider extends AbstractNetworkStatsProvider {
+    class OffloadTetheringStatsProvider extends NetworkStatsProvider {
         // These stats must only ever be touched on the handler thread.
         @NonNull
         private NetworkStats mIfaceStats = new NetworkStats(0L, 0);
         @NonNull
         private NetworkStats mUidStats = new NetworkStats(0L, 0);
 
+        /**
+         * A helper function that collect tether stats from local hashmap. Note that this does not
+         * invoke binder call.
+         */
         @VisibleForTesting
         @NonNull
         NetworkStats getTetherStats(@NonNull StatsType how) {
@@ -287,7 +291,7 @@ public class OffloadController {
         }
 
         @Override
-        public void setLimit(String iface, long quotaBytes) {
+        public void onSetLimit(String iface, long quotaBytes) {
             // Listen for all iface is necessary since upstream might be changed after limit
             // is set.
             mHandler.post(() -> {
@@ -315,13 +319,12 @@ public class OffloadController {
          */
         public void pushTetherStats() {
             // TODO: remove the accumulated stats and report the diff from HAL directly.
-            if (null == mStatsProviderCb) return;
             final NetworkStats ifaceDiff =
                     getTetherStats(StatsType.STATS_PER_IFACE).subtract(mIfaceStats);
             final NetworkStats uidDiff =
                     getTetherStats(StatsType.STATS_PER_UID).subtract(mUidStats);
             try {
-                mStatsProviderCb.onStatsUpdated(0 /* token */, ifaceDiff, uidDiff);
+                notifyStatsUpdated(0 /* token */, ifaceDiff, uidDiff);
                 mIfaceStats = mIfaceStats.add(ifaceDiff);
                 mUidStats = mUidStats.add(uidDiff);
             } catch (RuntimeException e) {
@@ -330,7 +333,7 @@ public class OffloadController {
         }
 
         @Override
-        public void requestStatsUpdate(int token) {
+        public void onRequestStatsUpdate(int token) {
             // Do not attempt to update stats by querying the offload HAL
             // synchronously from a different thread than the Handler thread. http://b/64771555.
             mHandler.post(() -> {
@@ -340,7 +343,7 @@ public class OffloadController {
         }
 
         @Override
-        public void setAlert(long quotaBytes) {
+        public void onSetAlert(long quotaBytes) {
             // TODO: Ask offload HAL to notify alert without stopping traffic.
         }
     }

packages/Tethering/tests/unit/src/com/android/server/connectivity/tethering/OffloadControllerTest.java

@@ -33,6 +33,8 @@ import static com.android.testutils.MiscAssertsKt.assertContainsAll;
 import static com.android.testutils.MiscAssertsKt.assertThrows;
 import static com.android.testutils.NetworkStatsUtilsKt.orderInsensitiveEquals;
 
+import static junit.framework.Assert.assertNotNull;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.anyInt;
@@ -61,8 +63,7 @@ import android.net.LinkProperties;
 import android.net.NetworkStats;
 import android.net.NetworkStats.Entry;
 import android.net.RouteInfo;
-import android.net.netstats.provider.AbstractNetworkStatsProvider;
+import android.net.netstats.provider.INetworkStatsProviderCallback;
-import android.net.netstats.provider.NetworkStatsProviderCallback;
 import android.net.util.SharedLog;
 import android.os.Handler;
 import android.os.Looper;
@@ -108,12 +109,10 @@ public class OffloadControllerTest {
     @Mock private ApplicationInfo mApplicationInfo;
     @Mock private Context mContext;
     @Mock private NetworkStatsManager mStatsManager;
-    @Mock private NetworkStatsProviderCallback mTetherStatsProviderCb;
+    @Mock private INetworkStatsProviderCallback mTetherStatsProviderCb;
+    private OffloadController.OffloadTetheringStatsProvider mTetherStatsProvider;
     private final ArgumentCaptor<ArrayList> mStringArrayCaptor =
             ArgumentCaptor.forClass(ArrayList.class);
-    private final ArgumentCaptor<OffloadController.OffloadTetheringStatsProvider>
-            mTetherStatsProviderCaptor =
-            ArgumentCaptor.forClass(OffloadController.OffloadTetheringStatsProvider.class);
     private final ArgumentCaptor<OffloadHardwareInterface.ControlCallback> mControlCallbackCaptor =
             ArgumentCaptor.forClass(OffloadHardwareInterface.ControlCallback.class);
     private MockContentResolver mContentResolver;
@@ -126,8 +125,6 @@ public class OffloadControllerTest {
         mContentResolver.addProvider(Settings.AUTHORITY, new FakeSettingsProvider());
         when(mContext.getContentResolver()).thenReturn(mContentResolver);
         FakeSettingsProvider.clearSettingsProvider();
-        when(mStatsManager.registerNetworkStatsProvider(anyString(), any()))
-                .thenReturn(mTetherStatsProviderCb);
     }
 
     @After public void tearDown() throws Exception {
@@ -154,8 +151,14 @@ public class OffloadControllerTest {
     private OffloadController makeOffloadController() throws Exception {
         OffloadController offload = new OffloadController(new Handler(Looper.getMainLooper()),
                 mHardware, mContentResolver, mStatsManager, new SharedLog("test"));
+        final ArgumentCaptor<OffloadController.OffloadTetheringStatsProvider>
+                tetherStatsProviderCaptor =
+                ArgumentCaptor.forClass(OffloadController.OffloadTetheringStatsProvider.class);
         verify(mStatsManager).registerNetworkStatsProvider(anyString(),
-                mTetherStatsProviderCaptor.capture());
+                tetherStatsProviderCaptor.capture());
+        mTetherStatsProvider = tetherStatsProviderCaptor.getValue();
+        assertNotNull(mTetherStatsProvider);
+        mTetherStatsProvider.setProviderCallbackBinder(mTetherStatsProviderCb);
         return offload;
     }
 
@@ -413,9 +416,6 @@ public class OffloadControllerTest {
         final OffloadController offload = makeOffloadController();
         offload.start();
 
-        final OffloadController.OffloadTetheringStatsProvider provider =
-                mTetherStatsProviderCaptor.getValue();
-
         final String ethernetIface = "eth1";
         final String mobileIface = "rmnet_data0";
 
@@ -443,8 +443,8 @@ public class OffloadControllerTest {
         inOrder.verify(mHardware, times(1)).getForwardedStats(eq(mobileIface));
 
         // Verify that the fetched stats are stored.
-        final NetworkStats ifaceStats = provider.getTetherStats(STATS_PER_IFACE);
+        final NetworkStats ifaceStats = mTetherStatsProvider.getTetherStats(STATS_PER_IFACE);
-        final NetworkStats uidStats = provider.getTetherStats(STATS_PER_UID);
+        final NetworkStats uidStats = mTetherStatsProvider.getTetherStats(STATS_PER_UID);
         final NetworkStats expectedIfaceStats = new NetworkStats(0L, 2)
                 .addValues(buildTestEntry(STATS_PER_IFACE, mobileIface, 999, 99999))
                 .addValues(buildTestEntry(STATS_PER_IFACE, ethernetIface, 12345, 54321));
@@ -462,13 +462,12 @@ public class OffloadControllerTest {
                 NetworkStats.class);
 
         // Force pushing stats update to verify the stats reported.
-        provider.pushTetherStats();
+        mTetherStatsProvider.pushTetherStats();
-        verify(mTetherStatsProviderCb, times(1)).onStatsUpdated(anyInt(),
+        verify(mTetherStatsProviderCb, times(1))
-                ifaceStatsCaptor.capture(), uidStatsCaptor.capture());
+                .notifyStatsUpdated(anyInt(), ifaceStatsCaptor.capture(), uidStatsCaptor.capture());
         assertTrue(orderInsensitiveEquals(expectedIfaceStats, ifaceStatsCaptor.getValue()));
         assertTrue(orderInsensitiveEquals(expectedUidStats, uidStatsCaptor.getValue()));
 
-
         when(mHardware.getForwardedStats(eq(ethernetIface))).thenReturn(
                 new ForwardedStats(100000, 100000));
         offload.setUpstreamLinkProperties(null);
@@ -483,8 +482,8 @@ public class OffloadControllerTest {
         inOrder.verifyNoMoreInteractions();
 
         // Verify that the stored stats is accumulated.
-        final NetworkStats ifaceStatsAccu = provider.getTetherStats(STATS_PER_IFACE);
+        final NetworkStats ifaceStatsAccu = mTetherStatsProvider.getTetherStats(STATS_PER_IFACE);
-        final NetworkStats uidStatsAccu = provider.getTetherStats(STATS_PER_UID);
+        final NetworkStats uidStatsAccu = mTetherStatsProvider.getTetherStats(STATS_PER_UID);
         final NetworkStats expectedIfaceStatsAccu = new NetworkStats(0L, 2)
                 .addValues(buildTestEntry(STATS_PER_IFACE, mobileIface, 999, 99999))
                 .addValues(buildTestEntry(STATS_PER_IFACE, ethernetIface, 112345, 154321));
@@ -498,7 +497,7 @@ public class OffloadControllerTest {
 
         // Verify that only diff of stats is reported.
         reset(mTetherStatsProviderCb);
-        provider.pushTetherStats();
+        mTetherStatsProvider.pushTetherStats();
         final NetworkStats expectedIfaceStatsDiff = new NetworkStats(0L, 2)
                 .addValues(buildTestEntry(STATS_PER_IFACE, mobileIface, 0, 0))
                 .addValues(buildTestEntry(STATS_PER_IFACE, ethernetIface, 100000, 100000));
@@ -506,8 +505,8 @@ public class OffloadControllerTest {
         final NetworkStats expectedUidStatsDiff = new NetworkStats(0L, 2)
                 .addValues(buildTestEntry(STATS_PER_UID, mobileIface, 0, 0))
                 .addValues(buildTestEntry(STATS_PER_UID, ethernetIface, 100000, 100000));
-        verify(mTetherStatsProviderCb, times(1)).onStatsUpdated(anyInt(),
+        verify(mTetherStatsProviderCb, times(1))
-                ifaceStatsCaptor.capture(), uidStatsCaptor.capture());
+                .notifyStatsUpdated(anyInt(), ifaceStatsCaptor.capture(), uidStatsCaptor.capture());
         assertTrue(orderInsensitiveEquals(expectedIfaceStatsDiff, ifaceStatsCaptor.getValue()));
         assertTrue(orderInsensitiveEquals(expectedUidStatsDiff, uidStatsCaptor.getValue()));
     }
@@ -529,19 +528,18 @@ public class OffloadControllerTest {
         lp.setInterfaceName(ethernetIface);
         offload.setUpstreamLinkProperties(lp);
 
-        AbstractNetworkStatsProvider provider = mTetherStatsProviderCaptor.getValue();
         final InOrder inOrder = inOrder(mHardware);
         when(mHardware.setUpstreamParameters(any(), any(), any(), any())).thenReturn(true);
         when(mHardware.setDataLimit(anyString(), anyLong())).thenReturn(true);
 
         // Applying an interface quota to the current upstream immediately sends it to the hardware.
-        provider.setLimit(ethernetIface, ethernetLimit);
+        mTetherStatsProvider.onSetLimit(ethernetIface, ethernetLimit);
         waitForIdle();
         inOrder.verify(mHardware).setDataLimit(ethernetIface, ethernetLimit);
         inOrder.verifyNoMoreInteractions();
 
         // Applying an interface quota to another upstream does not take any immediate action.
-        provider.setLimit(mobileIface, mobileLimit);
+        mTetherStatsProvider.onSetLimit(mobileIface, mobileLimit);
         waitForIdle();
         inOrder.verify(mHardware, never()).setDataLimit(anyString(), anyLong());
 
@@ -554,7 +552,7 @@ public class OffloadControllerTest {
 
         // Setting a limit of ITetheringStatsProvider.QUOTA_UNLIMITED causes the limit to be set
         // to Long.MAX_VALUE.
-        provider.setLimit(mobileIface, ITetheringStatsProvider.QUOTA_UNLIMITED);
+        mTetherStatsProvider.onSetLimit(mobileIface, ITetheringStatsProvider.QUOTA_UNLIMITED);
         waitForIdle();
         inOrder.verify(mHardware).setDataLimit(mobileIface, Long.MAX_VALUE);
 
@@ -562,7 +560,7 @@ public class OffloadControllerTest {
         when(mHardware.setUpstreamParameters(any(), any(), any(), any())).thenReturn(false);
         lp.setInterfaceName(ethernetIface);
         offload.setUpstreamLinkProperties(lp);
-        provider.setLimit(mobileIface, mobileLimit);
+        mTetherStatsProvider.onSetLimit(mobileIface, mobileLimit);
         waitForIdle();
         inOrder.verify(mHardware, never()).setDataLimit(anyString(), anyLong());
 
@@ -571,7 +569,7 @@ public class OffloadControllerTest {
         when(mHardware.setDataLimit(anyString(), anyLong())).thenReturn(false);
         lp.setInterfaceName(mobileIface);
         offload.setUpstreamLinkProperties(lp);
-        provider.setLimit(mobileIface, mobileLimit);
+        mTetherStatsProvider.onSetLimit(mobileIface, mobileLimit);
         waitForIdle();
         inOrder.verify(mHardware).getForwardedStats(ethernetIface);
         inOrder.verify(mHardware).stopOffloadControl();
@@ -587,7 +585,7 @@ public class OffloadControllerTest {
 
         OffloadHardwareInterface.ControlCallback callback = mControlCallbackCaptor.getValue();
         callback.onStoppedLimitReached();
-        verify(mTetherStatsProviderCb, times(1)).onStatsUpdated(anyInt(), any(), any());
+        verify(mTetherStatsProviderCb, times(1)).notifyStatsUpdated(anyInt(), any(), any());
     }
 
     @Test
@@ -691,7 +689,7 @@ public class OffloadControllerTest {
         verify(mHardware, times(1)).getForwardedStats(eq(RMNET0));
         verify(mHardware, times(1)).getForwardedStats(eq(WLAN0));
         // TODO: verify the exact stats reported.
-        verify(mTetherStatsProviderCb, times(1)).onStatsUpdated(anyInt(), any(), any());
+        verify(mTetherStatsProviderCb, times(1)).notifyStatsUpdated(anyInt(), any(), any());
         verifyNoMoreInteractions(mTetherStatsProviderCb);
         verifyNoMoreInteractions(mHardware);
     }
@@ -756,7 +754,7 @@ public class OffloadControllerTest {
         // Verify forwarded stats behaviour.
         verify(mHardware, times(1)).getForwardedStats(eq(RMNET0));
         verify(mHardware, times(1)).getForwardedStats(eq(WLAN0));
-        verify(mTetherStatsProviderCb, times(1)).onStatsUpdated(anyInt(), any(), any());
+        verify(mTetherStatsProviderCb, times(1)).notifyStatsUpdated(anyInt(), any(), any());
         verifyNoMoreInteractions(mTetherStatsProviderCb);
 
         // TODO: verify local prefixes and downstreams are also pushed to the HAL.

services/core/java/com/android/server/net/NetworkPolicyManagerInternal.java

@@ -21,7 +21,7 @@ import static com.android.server.net.NetworkPolicyManagerService.isUidNetworking
 import android.annotation.NonNull;
 import android.net.Network;
 import android.net.NetworkTemplate;
-import android.net.netstats.provider.AbstractNetworkStatsProvider;
+import android.net.netstats.provider.NetworkStatsProvider;
 import android.telephony.SubscriptionPlan;
 
 import java.util.Set;
@@ -130,8 +130,8 @@ public abstract class NetworkPolicyManagerInternal {
             Set<String> packageNames, int userId);
 
     /**
-     *  Notifies that the specified {@link AbstractNetworkStatsProvider} has reached its quota
+     *  Notifies that the specified {@link NetworkStatsProvider} has reached its quota
-     *  which was set through {@link AbstractNetworkStatsProvider#setLimit(String, long)}.
+     *  which was set through {@link NetworkStatsProvider#onSetLimit(String, long)}.
      *
      * @param tag the human readable identifier of the custom network stats provider.
      */

services/core/java/com/android/server/net/NetworkPolicyManagerService.java

@@ -75,7 +75,7 @@ import static android.net.NetworkTemplate.MATCH_MOBILE;
 import static android.net.NetworkTemplate.MATCH_WIFI;
 import static android.net.NetworkTemplate.buildTemplateMobileAll;
 import static android.net.TrafficStats.MB_IN_BYTES;
-import static android.net.netstats.provider.AbstractNetworkStatsProvider.QUOTA_UNLIMITED;
+import static android.net.netstats.provider.NetworkStatsProvider.QUOTA_UNLIMITED;
 import static android.os.Trace.TRACE_TAG_NETWORK;
 import static android.provider.Settings.Global.NETPOLICY_OVERRIDE_ENABLED;
 import static android.provider.Settings.Global.NETPOLICY_QUOTA_ENABLED;

services/core/java/com/android/server/net/NetworkStatsService.java

@@ -103,7 +103,7 @@ import android.net.NetworkTemplate;
 import android.net.TrafficStats;
 import android.net.netstats.provider.INetworkStatsProvider;
 import android.net.netstats.provider.INetworkStatsProviderCallback;
-import android.net.netstats.provider.NetworkStatsProviderCallback;
+import android.net.netstats.provider.NetworkStatsProvider;
 import android.os.BestClock;
 import android.os.Binder;
 import android.os.DropBoxManager;
@@ -558,7 +558,7 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         } catch (RemoteException e) {
             // ignored; service lives in system_server
         }
-        invokeForAllStatsProviderCallbacks((cb) -> cb.mProvider.setAlert(mGlobalAlertBytes));
+        invokeForAllStatsProviderCallbacks((cb) -> cb.mProvider.onSetAlert(mGlobalAlertBytes));
     }
 
     @Override
@@ -1376,7 +1376,8 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         Trace.traceBegin(TRACE_TAG_NETWORK, "provider.requestStatsUpdate");
         final int registeredCallbackCount = mStatsProviderCbList.getRegisteredCallbackCount();
         mStatsProviderSem.drainPermits();
-        invokeForAllStatsProviderCallbacks((cb) -> cb.mProvider.requestStatsUpdate(0 /* unused */));
+        invokeForAllStatsProviderCallbacks(
+                (cb) -> cb.mProvider.onRequestStatsUpdate(0 /* unused */));
         try {
             mStatsProviderSem.tryAcquire(registeredCallbackCount,
                     MAX_STATS_PROVIDER_POLL_WAIT_TIME_MS, TimeUnit.MILLISECONDS);
@@ -1551,7 +1552,7 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         @Override
         public void setStatsProviderLimitAsync(@NonNull String iface, long quota) {
             Slog.v(TAG, "setStatsProviderLimitAsync(" + iface + "," + quota + ")");
-            invokeForAllStatsProviderCallbacks((cb) -> cb.mProvider.setLimit(iface, quota));
+            invokeForAllStatsProviderCallbacks((cb) -> cb.mProvider.onSetLimit(iface, quota));
         }
     }
 
@@ -1820,12 +1821,10 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
      *
      * @param tag a human readable identifier of the custom network stats provider.
      * @param provider the {@link INetworkStatsProvider} binder corresponding to the
-     *                 {@link android.net.netstats.provider.AbstractNetworkStatsProvider} to be
+     *                 {@link NetworkStatsProvider} to be registered.
-     *                 registered.
      *
-     * @return a binder interface of
+     * @return a {@link INetworkStatsProviderCallback} binder
-     *         {@link android.net.netstats.provider.NetworkStatsProviderCallback}, which can be
+     *         interface, which can be used to report events to the system.
-     *         used to report events to the system.
      */
     public @NonNull INetworkStatsProviderCallback registerNetworkStatsProvider(
             @NonNull String tag, @NonNull INetworkStatsProvider provider) {
@@ -1931,7 +1930,7 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         }
 
         @Override
-        public void onStatsUpdated(int token, @Nullable NetworkStats ifaceStats,
+        public void notifyStatsUpdated(int token, @Nullable NetworkStats ifaceStats,
                 @Nullable NetworkStats uidStats) {
             // TODO: 1. Use token to map ifaces to correct NetworkIdentity.
             //       2. Store the difference and store it directly to the recorder.
@@ -1943,12 +1942,12 @@ public class NetworkStatsService extends INetworkStatsService.Stub {
         }
 
         @Override
-        public void onAlertReached() throws RemoteException {
+        public void notifyAlertReached() throws RemoteException {
             mAlertObserver.limitReached(LIMIT_GLOBAL_ALERT, null /* unused */);
         }
 
         @Override
-        public void onLimitReached() {
+        public void notifyLimitReached() {
             Log.d(TAG, mTag + ": onLimitReached");
             LocalServices.getService(NetworkPolicyManagerInternal.class)
                     .onStatsProviderLimitReached(mTag);

tests/net/java/com/android/server/net/NetworkStatsServiceTest.java

@@ -1026,7 +1026,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         mService.forceUpdateIfaces(NETWORKS_WIFI, states, getActiveIface(states), new VpnInfo[0]);
 
         // Verifies that one requestStatsUpdate will be called during iface update.
-        provider.expectStatsUpdate(0 /* unused */);
+        provider.expectOnRequestStatsUpdate(0 /* unused */);
 
         // Create some initial traffic and report to the service.
         incrementCurrentTime(HOUR_IN_MILLIS);
@@ -1037,7 +1037,7 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
                 .addValues(new NetworkStats.Entry(TEST_IFACE, UID_RED, SET_DEFAULT,
                         0xF00D, METERED_YES, ROAMING_NO, DEFAULT_NETWORK_YES,
                         64L, 1L, 64L, 1L, 1L));
-        cb.onStatsUpdated(0 /* unused */, expectedStats, expectedStats);
+        cb.notifyStatsUpdated(0 /* unused */, expectedStats, expectedStats);
 
         // Make another empty mutable stats object. This is necessary since the new NetworkStats
         // object will be used to compare with the old one in NetworkStatsRecoder, two of them
@@ -1047,8 +1047,8 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         forcePollAndWaitForIdle();
 
         // Verifies that one requestStatsUpdate and setAlert will be called during polling.
-        provider.expectStatsUpdate(0 /* unused */);
+        provider.expectOnRequestStatsUpdate(0 /* unused */);
-        provider.expectSetAlert(MB_IN_BYTES);
+        provider.expectOnSetAlert(MB_IN_BYTES);
 
         // Verifies that service recorded history, does not verify uid tag part.
         assertUidTotal(sTemplateWifi, UID_RED, 128L, 2L, 128L, 2L, 1);
@@ -1082,13 +1082,13 @@ public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
         assertNotNull(cb);
 
         // Simulates alert quota of the provider has been reached.
-        cb.onAlertReached();
+        cb.notifyAlertReached();
         HandlerUtilsKt.waitForIdle(mHandlerThread, WAIT_TIMEOUT);
 
         // Verifies that polling is triggered by alert reached.
-        provider.expectStatsUpdate(0 /* unused */);
+        provider.expectOnRequestStatsUpdate(0 /* unused */);
         // Verifies that global alert will be re-armed.
-        provider.expectSetAlert(MB_IN_BYTES);
+        provider.expectOnSetAlert(MB_IN_BYTES);
     }
 
     private static File getBaseDir(File statsDir) {