diff --git a/OWNERS_core_networking b/OWNERS_core_networking
index 83f798a..6d8ed4a 100644
--- a/OWNERS_core_networking
+++ b/OWNERS_core_networking
@@ -1,10 +1,6 @@
-chiachangwang@google.com
-cken@google.com
 jchalard@google.com
 junyulai@google.com
-lifr@google.com
 lorenzo@google.com
-markchien@google.com
 martinwu@google.com
 maze@google.com
 motomuman@google.com
@@ -12,7 +8,5 @@
 prohr@google.com
 reminv@google.com
 satk@google.com
-waynema@google.com
 xiaom@google.com
-yumike@google.com
 yuyanghuang@google.com
diff --git a/framework-t/src/android/net/INetworkStatsService.aidl b/framework-t/src/android/net/INetworkStatsService.aidl
index c86f7fd..7f0c1fe 100644
--- a/framework-t/src/android/net/INetworkStatsService.aidl
+++ b/framework-t/src/android/net/INetworkStatsService.aidl
@@ -101,4 +101,7 @@
      * Note that invocation of any interface will be sent to all providers.
      */
      void setStatsProviderWarningAndLimitAsync(String iface, long warning, long limit);
+
+     /** Clear TrafficStats rate-limit caches. */
+     void clearTrafficStatsRateLimitCaches();
 }
diff --git a/framework-t/src/android/net/TrafficStats.java b/framework-t/src/android/net/TrafficStats.java
index a69b38d..77c8001 100644
--- a/framework-t/src/android/net/TrafficStats.java
+++ b/framework-t/src/android/net/TrafficStats.java
@@ -19,6 +19,7 @@
 import static android.annotation.SystemApi.Client.MODULE_LIBRARIES;
 
 import android.annotation.NonNull;
+import android.annotation.RequiresPermission;
 import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
@@ -692,6 +693,27 @@
         return UNSUPPORTED;
     }
 
+    /** Clear TrafficStats rate-limit caches.
+     *
+     * This is mainly for {@link com.android.server.net.NetworkStatsService} to
+     * clear rate-limit cache to avoid caching for TrafficStats API results.
+     * Tests might get stale values after generating network traffic, which
+     * generally need to wait for cache expiry to get updated values.
+     *
+     * @hide
+     */
+    @RequiresPermission(anyOf = {
+            NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
+            android.Manifest.permission.NETWORK_STACK,
+            android.Manifest.permission.NETWORK_SETTINGS})
+    public static void clearRateLimitCaches() {
+        try {
+            getStatsService().clearTrafficStatsRateLimitCaches();
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
     /**
      * Return the number of packets transmitted on the specified interface since the interface
      * was created. Statistics are measured at the network layer, so both TCP and
diff --git a/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java b/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
index 42a922d..385adc6 100644
--- a/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
+++ b/service-t/src/com/android/metrics/NetworkNsdReportedMetrics.java
@@ -247,12 +247,15 @@
      * @param isLegacy Whether this call is using legacy backend.
      * @param transactionId The transaction id of service resolution.
      * @param durationMs The duration before stop resolving the service.
+     * @param sentQueryCount The count of sent queries during resolving.
      */
-    public void reportServiceResolutionStop(boolean isLegacy, int transactionId, long durationMs) {
+    public void reportServiceResolutionStop(boolean isLegacy, int transactionId, long durationMs,
+            int sentQueryCount) {
         final Builder builder = makeReportedBuilder(isLegacy, transactionId);
         builder.setType(NsdEventType.NET_RESOLVE);
         builder.setQueryResult(MdnsQueryResult.MQR_SERVICE_RESOLUTION_STOP);
         builder.setEventDurationMillisec(durationMs);
+        builder.setSentQueryCount(sentQueryCount);
         mDependencies.statsWrite(builder.build());
     }
 
diff --git a/service-t/src/com/android/server/NsdService.java b/service-t/src/com/android/server/NsdService.java
index 8552eec..2e258ab 100644
--- a/service-t/src/com/android/server/NsdService.java
+++ b/service-t/src/com/android/server/NsdService.java
@@ -2851,7 +2851,8 @@
                                 request.getSentQueryCount());
                     } else if (listener instanceof ResolutionListener) {
                         mMetrics.reportServiceResolutionStop(false /* isLegacy */, transactionId,
-                                request.calculateRequestDurationMs(mClock.elapsedRealtime()));
+                                request.calculateRequestDurationMs(mClock.elapsedRealtime()),
+                                request.getSentQueryCount());
                     } else if (listener instanceof ServiceInfoListener) {
                         mMetrics.reportServiceInfoCallbackUnregistered(transactionId,
                                 request.calculateRequestDurationMs(mClock.elapsedRealtime()),
@@ -2892,7 +2893,8 @@
                     case NsdManager.RESOLVE_SERVICE:
                         stopResolveService(transactionId);
                         mMetrics.reportServiceResolutionStop(true /* isLegacy */, transactionId,
-                                request.calculateRequestDurationMs(mClock.elapsedRealtime()));
+                                request.calculateRequestDurationMs(mClock.elapsedRealtime()),
+                                NO_SENT_QUERY_COUNT);
                         break;
                     case NsdManager.REGISTER_SERVICE:
                         unregisterService(transactionId);
@@ -3108,7 +3110,8 @@
             mMetrics.reportServiceResolutionStop(
                     isLegacyClientRequest(request),
                     request.mTransactionId,
-                    request.calculateRequestDurationMs(mClock.elapsedRealtime()));
+                    request.calculateRequestDurationMs(mClock.elapsedRealtime()),
+                    request.getSentQueryCount());
             try {
                 mCb.onStopResolutionSucceeded(listenerKey);
             } catch (RemoteException e) {
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index 5e98ee1..d43c3da 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -16,6 +16,7 @@
 
 package com.android.server.net;
 
+import static android.Manifest.permission.NETWORK_SETTINGS;
 import static android.Manifest.permission.NETWORK_STATS_PROVIDER;
 import static android.Manifest.permission.READ_NETWORK_USAGE_HISTORY;
 import static android.Manifest.permission.UPDATE_DEVICE_STATS;
@@ -50,12 +51,17 @@
 import static android.net.NetworkTemplate.MATCH_WIFI;
 import static android.net.TrafficStats.KB_IN_BYTES;
 import static android.net.TrafficStats.MB_IN_BYTES;
+import static android.net.TrafficStats.TYPE_RX_BYTES;
+import static android.net.TrafficStats.TYPE_RX_PACKETS;
+import static android.net.TrafficStats.TYPE_TX_BYTES;
+import static android.net.TrafficStats.TYPE_TX_PACKETS;
 import static android.net.TrafficStats.UID_TETHERING;
 import static android.net.TrafficStats.UNSUPPORTED;
 import static android.net.netstats.NetworkStatsDataMigrationUtils.PREFIX_UID;
 import static android.net.netstats.NetworkStatsDataMigrationUtils.PREFIX_UID_TAG;
 import static android.net.netstats.NetworkStatsDataMigrationUtils.PREFIX_XT;
 import static android.os.Trace.TRACE_TAG_NETWORK;
+import static android.provider.DeviceConfig.NAMESPACE_TETHERING;
 import static android.system.OsConstants.ENOENT;
 import static android.telephony.SubscriptionManager.INVALID_SUBSCRIPTION_ID;
 import static android.text.format.DateUtils.DAY_IN_MILLIS;
@@ -64,6 +70,7 @@
 import static android.text.format.DateUtils.SECOND_IN_MILLIS;
 
 import static com.android.internal.annotations.VisibleForTesting.Visibility.PRIVATE;
+import static com.android.net.module.util.DeviceConfigUtils.getDeviceConfigPropertyInt;
 import static com.android.net.module.util.NetworkCapabilitiesUtils.getDisplayTransport;
 import static com.android.net.module.util.NetworkStatsUtils.LIMIT_GLOBAL_ALERT;
 import static com.android.server.net.NetworkStatsEventLogger.POLL_REASON_PERIODIC;
@@ -299,6 +306,12 @@
     static final String NETSTATS_FASTDATAINPUT_SUCCESSES_COUNTER_NAME = "fastdatainput.successes";
     static final String NETSTATS_FASTDATAINPUT_FALLBACKS_COUNTER_NAME = "fastdatainput.fallbacks";
 
+    static final String TRAFFIC_STATS_CACHE_EXPIRY_DURATION_NAME =
+            "trafficstats_cache_expiry_duration_ms";
+    static final String TRAFFIC_STATS_CACHE_MAX_ENTRIES_NAME = "trafficstats_cache_max_entries";
+    static final int DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS = 1000;
+    static final int DEFAULT_TRAFFIC_STATS_CACHE_MAX_ENTRIES = 400;
+
     private final Context mContext;
     private final NetworkStatsFactory mStatsFactory;
     private final AlarmManager mAlarmManager;
@@ -454,6 +467,13 @@
 
     private long mLastStatsSessionPoll;
 
+    private final TrafficStatsRateLimitCache mTrafficStatsTotalCache;
+    private final TrafficStatsRateLimitCache mTrafficStatsIfaceCache;
+    private final TrafficStatsRateLimitCache mTrafficStatsUidCache;
+    static final String TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG =
+            "trafficstats_rate_limit_cache_enabled_flag";
+    private final boolean mSupportTrafficStatsRateLimitCache;
+
     private final Object mOpenSessionCallsLock = new Object();
 
     /**
@@ -643,6 +663,16 @@
             mEventLogger = null;
         }
 
+        final long cacheExpiryDurationMs = mDeps.getTrafficStatsRateLimitCacheExpiryDuration();
+        final int cacheMaxEntries = mDeps.getTrafficStatsRateLimitCacheMaxEntries();
+        mSupportTrafficStatsRateLimitCache = mDeps.supportTrafficStatsRateLimitCache(mContext);
+        mTrafficStatsTotalCache = new TrafficStatsRateLimitCache(mClock,
+                cacheExpiryDurationMs, cacheMaxEntries);
+        mTrafficStatsIfaceCache = new TrafficStatsRateLimitCache(mClock,
+                cacheExpiryDurationMs, cacheMaxEntries);
+        mTrafficStatsUidCache = new TrafficStatsRateLimitCache(mClock,
+                cacheExpiryDurationMs, cacheMaxEntries);
+
         // TODO: Remove bpfNetMaps creation and always start SkDestroyListener
         // Following code is for the experiment to verify the SkDestroyListener refactoring. Based
         // on the experiment flag, BpfNetMaps starts C SkDestroyListener (existing code) or
@@ -696,7 +726,7 @@
          * Get the count of import legacy target attempts.
          */
         public int getImportLegacyTargetAttempts() {
-            return DeviceConfigUtils.getDeviceConfigPropertyInt(
+            return getDeviceConfigPropertyInt(
                     DeviceConfig.NAMESPACE_TETHERING,
                     NETSTATS_IMPORT_LEGACY_TARGET_ATTEMPTS,
                     DEFAULT_NETSTATS_IMPORT_LEGACY_TARGET_ATTEMPTS);
@@ -706,7 +736,7 @@
          * Get the count of using FastDataInput target attempts.
          */
         public int getUseFastDataInputTargetAttempts() {
-            return DeviceConfigUtils.getDeviceConfigPropertyInt(
+            return getDeviceConfigPropertyInt(
                     DeviceConfig.NAMESPACE_TETHERING,
                     NETSTATS_FASTDATAINPUT_TARGET_ATTEMPTS, 0);
         }
@@ -888,6 +918,76 @@
             return DeviceConfigUtils.isTetheringFeatureNotChickenedOut(
                     ctx, CONFIG_ENABLE_NETWORK_STATS_EVENT_LOGGER);
         }
+
+        /**
+         * Get whether TrafficStats rate-limit cache is supported.
+         *
+         * This method should only be called once in the constructor,
+         * to ensure that the code does not need to deal with flag values changing at runtime.
+         */
+        public boolean supportTrafficStatsRateLimitCache(@NonNull Context ctx) {
+            return SdkLevel.isAtLeastV() && DeviceConfigUtils.isTetheringFeatureNotChickenedOut(
+                    ctx, TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG);
+        }
+
+        /**
+         * Get TrafficStats rate-limit cache expiry.
+         *
+         * This method should only be called once in the constructor,
+         * to ensure that the code does not need to deal with flag values changing at runtime.
+         */
+        public int getTrafficStatsRateLimitCacheExpiryDuration() {
+            return getDeviceConfigPropertyInt(
+                    NAMESPACE_TETHERING, TRAFFIC_STATS_CACHE_EXPIRY_DURATION_NAME,
+                    DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS);
+        }
+
+        /**
+         * Get TrafficStats rate-limit cache max entries.
+         *
+         * This method should only be called once in the constructor,
+         * to ensure that the code does not need to deal with flag values changing at runtime.
+         */
+        public int getTrafficStatsRateLimitCacheMaxEntries() {
+            return getDeviceConfigPropertyInt(
+                    NAMESPACE_TETHERING, TRAFFIC_STATS_CACHE_MAX_ENTRIES_NAME,
+                    DEFAULT_TRAFFIC_STATS_CACHE_MAX_ENTRIES);
+        }
+
+        /**
+         * Retrieves native network total statistics.
+         *
+         * @return A NetworkStats.Entry containing the native statistics, or
+         *         null if an error occurs.
+         */
+        @Nullable
+        public NetworkStats.Entry nativeGetTotalStat() {
+            return NetworkStatsService.nativeGetTotalStat();
+        }
+
+        /**
+         * Retrieves native network interface statistics for the specified interface.
+         *
+         * @param iface The name of the network interface to query.
+         * @return A NetworkStats.Entry containing the native statistics for the interface, or
+         *         null if an error occurs.
+         */
+        @Nullable
+        public NetworkStats.Entry nativeGetIfaceStat(String iface) {
+            return NetworkStatsService.nativeGetIfaceStat(iface);
+        }
+
+        /**
+         * Retrieves native network uid statistics for the specified uid.
+         *
+         * @param uid The uid of the application to query.
+         * @return A NetworkStats.Entry containing the native statistics for the uid, or
+         *         null if an error occurs.
+         */
+        @Nullable
+        public NetworkStats.Entry nativeGetUidStat(int uid) {
+            return NetworkStatsService.nativeGetUidStat(uid);
+        }
     }
 
     /**
@@ -1983,53 +2083,106 @@
         if (callingUid != android.os.Process.SYSTEM_UID && callingUid != uid) {
             return UNSUPPORTED;
         }
-        return getEntryValueForType(nativeGetUidStat(uid), type);
+        if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
+
+        if (!mSupportTrafficStatsRateLimitCache) {
+            return getEntryValueForType(mDeps.nativeGetUidStat(uid), type);
+        }
+
+        final NetworkStats.Entry entry = mTrafficStatsUidCache.getOrCompute(IFACE_ALL, uid,
+                () -> mDeps.nativeGetUidStat(uid));
+
+        return getEntryValueForType(entry, type);
+    }
+
+    @Nullable
+    private NetworkStats.Entry getIfaceStatsInternal(@NonNull String iface) {
+        final NetworkStats.Entry entry = mDeps.nativeGetIfaceStat(iface);
+        if (entry == null) {
+            return null;
+        }
+        // When tethering offload is in use, nativeIfaceStats does not contain usage from
+        // offload, add it back here. Note that the included statistics might be stale
+        // since polling newest stats from hardware might impact system health and not
+        // suitable for TrafficStats API use cases.
+        entry.add(getProviderIfaceStats(iface));
+        return entry;
     }
 
     @Override
     public long getIfaceStats(@NonNull String iface, int type) {
         Objects.requireNonNull(iface);
-        final NetworkStats.Entry entry = nativeGetIfaceStat(iface);
-        final long value = getEntryValueForType(entry, type);
-        if (value == UNSUPPORTED) {
-            return UNSUPPORTED;
-        } else {
-            // When tethering offload is in use, nativeIfaceStats does not contain usage from
-            // offload, add it back here. Note that the included statistics might be stale
-            // since polling newest stats from hardware might impact system health and not
-            // suitable for TrafficStats API use cases.
-            entry.add(getProviderIfaceStats(iface));
-            return getEntryValueForType(entry, type);
+        if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
+
+        if (!mSupportTrafficStatsRateLimitCache) {
+            return getEntryValueForType(getIfaceStatsInternal(iface), type);
         }
+
+        final NetworkStats.Entry entry = mTrafficStatsIfaceCache.getOrCompute(iface, UID_ALL,
+                () -> getIfaceStatsInternal(iface));
+
+        return getEntryValueForType(entry, type);
     }
 
     private long getEntryValueForType(@Nullable NetworkStats.Entry entry, int type) {
         if (entry == null) return UNSUPPORTED;
+        if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
         switch (type) {
-            case TrafficStats.TYPE_RX_BYTES:
+            case TYPE_RX_BYTES:
                 return entry.rxBytes;
-            case TrafficStats.TYPE_TX_BYTES:
-                return entry.txBytes;
-            case TrafficStats.TYPE_RX_PACKETS:
+            case TYPE_RX_PACKETS:
                 return entry.rxPackets;
-            case TrafficStats.TYPE_TX_PACKETS:
+            case TYPE_TX_BYTES:
+                return entry.txBytes;
+            case TYPE_TX_PACKETS:
                 return entry.txPackets;
             default:
-                return UNSUPPORTED;
+                throw new IllegalStateException("Bug: Invalid type: "
+                        + type + " should not reach here.");
         }
     }
 
+    private boolean isEntryValueTypeValid(int type) {
+        switch (type) {
+            case TYPE_RX_BYTES:
+            case TYPE_RX_PACKETS:
+            case TYPE_TX_BYTES:
+            case TYPE_TX_PACKETS:
+                return true;
+            default :
+                return false;
+        }
+    }
+
+    @Nullable
+    private NetworkStats.Entry getTotalStatsInternal() {
+        final NetworkStats.Entry entry = mDeps.nativeGetTotalStat();
+        if (entry == null) {
+            return null;
+        }
+        entry.add(getProviderIfaceStats(IFACE_ALL));
+        return entry;
+    }
+
     @Override
     public long getTotalStats(int type) {
-        final NetworkStats.Entry entry = nativeGetTotalStat();
-        final long value = getEntryValueForType(entry, type);
-        if (value == UNSUPPORTED) {
-            return UNSUPPORTED;
-        } else {
-            // Refer to comment in getIfaceStats
-            entry.add(getProviderIfaceStats(IFACE_ALL));
-            return getEntryValueForType(entry, type);
+        if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
+        if (!mSupportTrafficStatsRateLimitCache) {
+            return getEntryValueForType(getTotalStatsInternal(), type);
         }
+
+        final NetworkStats.Entry entry = mTrafficStatsTotalCache.getOrCompute(IFACE_ALL, UID_ALL,
+                () -> getTotalStatsInternal());
+
+        return getEntryValueForType(entry, type);
+    }
+
+    @Override
+    public void clearTrafficStatsRateLimitCaches() {
+        PermissionUtils.enforceNetworkStackPermissionOr(mContext, NETWORK_SETTINGS);
+        mTrafficStatsUidCache.clear();
+        mTrafficStatsIfaceCache.clear();
+        mTrafficStatsTotalCache.clear();
     }
 
     private NetworkStats.Entry getProviderIfaceStats(@Nullable String iface) {
@@ -2785,6 +2938,14 @@
             } catch (IOException e) {
                 pw.println("(failed to dump FastDataInput counters)");
             }
+            pw.print("trafficstats.cache.supported", mSupportTrafficStatsRateLimitCache);
+            pw.println();
+            pw.print(TRAFFIC_STATS_CACHE_EXPIRY_DURATION_NAME,
+                    mDeps.getTrafficStatsRateLimitCacheExpiryDuration());
+            pw.println();
+            pw.print(TRAFFIC_STATS_CACHE_MAX_ENTRIES_NAME,
+                    mDeps.getTrafficStatsRateLimitCacheMaxEntries());
+            pw.println();
 
             pw.decreaseIndent();
 
diff --git a/service/src/com/android/server/BpfNetMaps.java b/service/src/com/android/server/BpfNetMaps.java
index 42c1628..fc6d8c4 100644
--- a/service/src/com/android/server/BpfNetMaps.java
+++ b/service/src/com/android/server/BpfNetMaps.java
@@ -918,25 +918,6 @@
         }
     }
 
-    /**
-     * Return whether the network is blocked by firewall chains for the given uid.
-     *
-     * Note that {@link #getDataSaverEnabled()} has a latency before V.
-     *
-     * @param uid The target uid.
-     * @param isNetworkMetered Whether the target network is metered.
-     *
-     * @return True if the network is blocked. Otherwise, false.
-     * @throws ServiceSpecificException if the read fails.
-     *
-     * @hide
-     */
-    @RequiresApi(Build.VERSION_CODES.TIRAMISU)
-    public boolean isUidNetworkingBlocked(final int uid, boolean isNetworkMetered) {
-        return BpfNetMapsUtils.isUidNetworkingBlocked(uid, isNetworkMetered,
-                sConfigurationMap, sUidOwnerMap, sDataSaverEnabledMap);
-    }
-
     /** Register callback for statsd to pull atom. */
     @RequiresApi(Build.VERSION_CODES.TIRAMISU)
     public void setPullAtomCallback(final Context context) {
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index b1ae019..a15a2bf 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -2238,11 +2238,7 @@
         final long ident = Binder.clearCallingIdentity();
         try {
             final boolean metered = nc == null ? true : nc.isMetered();
-            if (mDeps.isAtLeastV()) {
-                return mBpfNetMaps.isUidNetworkingBlocked(uid, metered);
-            } else {
-                return mPolicyManager.isUidNetworkingBlocked(uid, metered);
-            }
+            return mPolicyManager.isUidNetworkingBlocked(uid, metered);
         } finally {
             Binder.restoreCallingIdentity(ident);
         }
diff --git a/staticlibs/device/com/android/net/module/util/netlink/NetlinkConstants.java b/staticlibs/device/com/android/net/module/util/netlink/NetlinkConstants.java
index ad7a4d7..27e1a32 100644
--- a/staticlibs/device/com/android/net/module/util/netlink/NetlinkConstants.java
+++ b/staticlibs/device/com/android/net/module/util/netlink/NetlinkConstants.java
@@ -123,6 +123,7 @@
     public static final short RTM_NEWRULE                   = 32;
     public static final short RTM_DELRULE                   = 33;
     public static final short RTM_GETRULE                   = 34;
+    public static final short RTM_NEWPREFIX                 = 52;
     public static final short RTM_NEWNDUSEROPT              = 68;
 
     // Netfilter netlink message types are presented by two bytes: high byte subsystem and
@@ -148,6 +149,8 @@
     public static final int RTMGRP_IPV4_IFADDR = 0x10;
     public static final int RTMGRP_IPV6_IFADDR = 0x100;
     public static final int RTMGRP_IPV6_ROUTE  = 0x400;
+    public static final int RTNLGRP_IPV6_PREFIX = 18;
+    public static final int RTMGRP_IPV6_PREFIX = 1 << (RTNLGRP_IPV6_PREFIX - 1);
     public static final int RTNLGRP_ND_USEROPT = 20;
     public static final int RTMGRP_ND_USEROPT = 1 << (RTNLGRP_ND_USEROPT - 1);
 
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/AutoReleaseNetworkCallbackRule.kt b/staticlibs/testutils/devicetests/com/android/testutils/AutoReleaseNetworkCallbackRule.kt
index 28ae609..93422ad 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/AutoReleaseNetworkCallbackRule.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/AutoReleaseNetworkCallbackRule.kt
@@ -21,6 +21,7 @@
 import android.net.Network
 import android.net.NetworkCapabilities
 import android.net.NetworkRequest
+import android.os.Handler
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.testutils.RecorderCallback.CallbackEntry
 import java.util.Collections
@@ -97,6 +98,15 @@
         cellRequestCb = null
     }
 
+    private fun addCallback(
+        cb: TestableNetworkCallback,
+        registrar: (TestableNetworkCallback) -> Unit
+    ): TestableNetworkCallback {
+        registrar(cb)
+        cbToCleanup.add(cb)
+        return cb
+    }
+
     /**
      * File a request for a Network.
      *
@@ -109,14 +119,27 @@
     @JvmOverloads
     fun requestNetwork(
         request: NetworkRequest,
-        cb: TestableNetworkCallback = TestableNetworkCallback()
-    ): TestableNetworkCallback {
-        cm.requestNetwork(request, cb)
-        cbToCleanup.add(cb)
-        return cb
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler? = null
+    ) = addCallback(cb) {
+        if (handler == null) {
+            cm.requestNetwork(request, it)
+        } else {
+            cm.requestNetwork(request, it, handler)
+        }
     }
 
     /**
+     * Overload of [requestNetwork] that allows specifying a timeout.
+     */
+    @JvmOverloads
+    fun requestNetwork(
+        request: NetworkRequest,
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        timeoutMs: Int,
+    ) = addCallback(cb) { cm.requestNetwork(request, it, timeoutMs) }
+
+    /**
      * File a callback for a NetworkRequest.
      *
      * This will fail tests (throw) if the cell network cannot be obtained, or if it was already
@@ -129,13 +152,63 @@
     fun registerNetworkCallback(
         request: NetworkRequest,
         cb: TestableNetworkCallback = TestableNetworkCallback()
-    ): TestableNetworkCallback {
-        cm.registerNetworkCallback(request, cb)
-        cbToCleanup.add(cb)
-        return cb
+    ) = addCallback(cb) { cm.registerNetworkCallback(request, it) }
+
+    /**
+     * @see ConnectivityManager.registerDefaultNetworkCallback
+     */
+    @JvmOverloads
+    fun registerDefaultNetworkCallback(
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler? = null
+    ) = addCallback(cb) {
+        if (handler == null) {
+            cm.registerDefaultNetworkCallback(it)
+        } else {
+            cm.registerDefaultNetworkCallback(it, handler)
+        }
     }
 
     /**
+     * @see ConnectivityManager.registerSystemDefaultNetworkCallback
+     */
+    @JvmOverloads
+    fun registerSystemDefaultNetworkCallback(
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler
+    ) = addCallback(cb) { cm.registerSystemDefaultNetworkCallback(it, handler) }
+
+    /**
+     * @see ConnectivityManager.registerDefaultNetworkCallbackForUid
+     */
+    @JvmOverloads
+    fun registerDefaultNetworkCallbackForUid(
+        uid: Int,
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler
+    ) = addCallback(cb) { cm.registerDefaultNetworkCallbackForUid(uid, it, handler) }
+
+    /**
+     * @see ConnectivityManager.registerBestMatchingNetworkCallback
+     */
+    @JvmOverloads
+    fun registerBestMatchingNetworkCallback(
+        request: NetworkRequest,
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler
+    ) = addCallback(cb) { cm.registerBestMatchingNetworkCallback(request, it, handler) }
+
+    /**
+     * @see ConnectivityManager.requestBackgroundNetwork
+     */
+    @JvmOverloads
+    fun requestBackgroundNetwork(
+        request: NetworkRequest,
+        cb: TestableNetworkCallback = TestableNetworkCallback(),
+        handler: Handler
+    ) = addCallback(cb) { cm.requestBackgroundNetwork(request, it, handler) }
+
+    /**
      * Unregister a callback filed using registration methods in this class.
      */
     fun unregisterNetworkCallback(cb: NetworkCallback) {
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt b/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt
index 8090d5b..3857810 100644
--- a/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt
+++ b/staticlibs/testutils/devicetests/com/android/testutils/ConnectUtil.kt
@@ -86,6 +86,7 @@
         val callback = TestableNetworkCallback(timeoutMs = WIFI_CONNECT_TIMEOUT_MS)
         cm.registerNetworkCallback(NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_WIFI)
+                .addCapability(NET_CAPABILITY_INTERNET)
                 .build(), callback)
 
         return tryTest {
diff --git a/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt b/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt
new file mode 100644
index 0000000..4185b05
--- /dev/null
+++ b/staticlibs/testutils/devicetests/com/android/testutils/SetFeatureFlagsRule.kt
@@ -0,0 +1,82 @@
+/*
+ * Copyright (C) 2024 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.testutils.com.android.testutils
+
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+ * A JUnit Rule that sets feature flags based on `@FeatureFlag` annotations.
+ *
+ * This rule enables dynamic control of feature flag states during testing.
+ *
+ * **Usage:**
+ * ```kotlin
+ * class MyTestClass {
+ *   @get:Rule
+ *   val setFeatureFlagsRule = SetFeatureFlagsRule(setFlagsMethod = (name, enabled) -> {
+ *     // Custom handling code.
+ *   })
+ *
+ *   // ... test methods with @FeatureFlag annotations
+ *   @FeatureFlag("FooBar1", true)
+ *   @FeatureFlag("FooBar2", false)
+ *   @Test
+ *   fun testFooBar() {}
+ * }
+ * ```
+ */
+class SetFeatureFlagsRule(val setFlagsMethod: (name: String, enabled: Boolean) -> Unit) : TestRule {
+    /**
+     * This annotation marks a test method as requiring a specific feature flag to be configured.
+     *
+     * Use this on test methods to dynamically control feature flag states during testing.
+     *
+     * @param name The name of the feature flag.
+     * @param enabled The desired state (true for enabled, false for disabled) of the feature flag.
+     */
+    @Target(AnnotationTarget.FUNCTION)
+    @Retention(AnnotationRetention.RUNTIME)
+    annotation class FeatureFlag(val name: String, val enabled: Boolean = true)
+
+    /**
+     * This method is the core of the rule, executed by the JUnit framework before each test method.
+     *
+     * It retrieves the test method's metadata.
+     * If any `@FeatureFlag` annotation is found, it passes every feature flag's name
+     * and enabled state into the user-specified lambda to apply custom actions.
+     */
+    override fun apply(base: Statement, description: Description): Statement {
+        return object : Statement() {
+            override fun evaluate() {
+                val testMethod = description.testClass.getMethod(description.methodName)
+                val featureFlagAnnotations = testMethod.getAnnotationsByType(
+                    FeatureFlag::class.java
+                )
+
+                for (featureFlagAnnotation in featureFlagAnnotations) {
+                    setFlagsMethod(featureFlagAnnotation.name, featureFlagAnnotation.enabled)
+                }
+
+                // Execute the test method, which includes methods annotated with
+                // @Before, @Test and @After.
+                base.evaluate()
+            }
+        }
+    }
+}
diff --git a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
index 4d465ba..c0f1080 100644
--- a/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/ConnectivityManagerTest.java
@@ -71,7 +71,9 @@
 import static android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_METERED;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VPN;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_PARTIAL_CONNECTIVITY;
+import static android.net.NetworkCapabilities.NET_CAPABILITY_TRUSTED;
 import static android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED;
 import static android.net.NetworkCapabilities.TRANSPORT_TEST;
 import static android.net.NetworkCapabilities.TRANSPORT_VPN;
@@ -81,7 +83,6 @@
 import static android.net.cts.util.CtsNetUtils.HTTP_PORT;
 import static android.net.cts.util.CtsNetUtils.NETWORK_CALLBACK_ACTION;
 import static android.net.cts.util.CtsNetUtils.TEST_HOST;
-import static android.net.cts.util.CtsNetUtils.TestNetworkCallback;
 import static android.net.cts.util.CtsTetheringUtils.TestTetheringEventCallback;
 import static android.os.MessageQueue.OnFileDescriptorEventListener.EVENT_INPUT;
 import static android.os.Process.INVALID_UID;
@@ -333,8 +334,6 @@
     private final ArraySet<Integer> mNetworkTypes = new ArraySet<>();
     private UiAutomation mUiAutomation;
     private CtsNetUtils mCtsNetUtils;
-    // The registered callbacks.
-    private List<NetworkCallback> mRegisteredCallbacks = new ArrayList<>();
     // Used for cleanup purposes.
     private final List<Range<Integer>> mVpnRequiredUidRanges = new ArrayList<>();
 
@@ -425,15 +424,12 @@
         // All tests in this class require a working Internet connection as they start. Make
         // sure there is still one as they end that's ready to use for the next test to use.
         mTestValidationConfigRule.runAfterNextCleanup(() -> {
-            final TestNetworkCallback callback = new TestNetworkCallback();
-            registerDefaultNetworkCallback(callback);
-            try {
-                assertNotNull("Couldn't restore Internet connectivity",
-                        callback.waitForAvailable());
-            } finally {
-                // Unregister all registered callbacks.
-                unregisterRegisteredCallbacks();
-            }
+            // mTestValidationConfigRule has higher order than networkCallbackRule, so
+            // networkCallbackRule is the outer rule and will be cleaned up after this method.
+            final TestableNetworkCallback callback =
+                    networkCallbackRule.registerDefaultNetworkCallback();
+            assertNotNull("Couldn't restore Internet connectivity",
+                    callback.eventuallyExpect(CallbackEntry.AVAILABLE));
         });
     }
 
@@ -993,10 +989,10 @@
         // default network.
         return new NetworkRequest.Builder()
                 .clearCapabilities()
-                .addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
-                .addCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
-                .addCapability(NetworkCapabilities.NET_CAPABILITY_NOT_VPN)
-                .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
+                .addCapability(NET_CAPABILITY_NOT_RESTRICTED)
+                .addCapability(NET_CAPABILITY_TRUSTED)
+                .addCapability(NET_CAPABILITY_NOT_VPN)
+                .addCapability(NET_CAPABILITY_INTERNET)
                 .build();
     }
 
@@ -1027,10 +1023,10 @@
         final String invalidPrivateDnsServer = "invalidhostname.example.com";
         final String goodPrivateDnsServer = "dns.google";
         mCtsNetUtils.storePrivateDnsSetting();
-        final TestableNetworkCallback cb = new TestableNetworkCallback();
         final NetworkRequest networkRequest = new NetworkRequest.Builder()
                 .addCapability(NET_CAPABILITY_INTERNET).build();
-        registerNetworkCallback(networkRequest, cb);
+        final TestableNetworkCallback cb =
+                networkCallbackRule.registerNetworkCallback(networkRequest);
         final Network networkForPrivateDns = mCm.getActiveNetwork();
         try {
             // Verifying the good private DNS sever
@@ -1068,24 +1064,27 @@
         assumeTrue(mPackageManager.hasSystemFeature(FEATURE_WIFI));
 
         // We will register for a WIFI network being available or lost.
-        final TestNetworkCallback callback = new TestNetworkCallback();
-        registerNetworkCallback(makeWifiNetworkRequest(), callback);
+        final TestableNetworkCallback callback = networkCallbackRule.registerNetworkCallback(
+                makeWifiNetworkRequest());
 
-        final TestNetworkCallback defaultTrackingCallback = new TestNetworkCallback();
-        registerDefaultNetworkCallback(defaultTrackingCallback);
+        final TestableNetworkCallback defaultTrackingCallback =
+                networkCallbackRule.registerDefaultNetworkCallback();
 
-        final TestNetworkCallback systemDefaultCallback = new TestNetworkCallback();
-        final TestNetworkCallback perUidCallback = new TestNetworkCallback();
-        final TestNetworkCallback bestMatchingCallback = new TestNetworkCallback();
+        final TestableNetworkCallback systemDefaultCallback = new TestableNetworkCallback();
+        final TestableNetworkCallback perUidCallback = new TestableNetworkCallback();
+        final TestableNetworkCallback bestMatchingCallback = new TestableNetworkCallback();
         final Handler h = new Handler(Looper.getMainLooper());
         if (TestUtils.shouldTestSApis()) {
             assertThrows(SecurityException.class, () ->
-                    registerSystemDefaultNetworkCallback(systemDefaultCallback, h));
+                    networkCallbackRule.registerSystemDefaultNetworkCallback(
+                            systemDefaultCallback, h));
             runWithShellPermissionIdentity(() -> {
-                registerSystemDefaultNetworkCallback(systemDefaultCallback, h);
-                registerDefaultNetworkCallbackForUid(Process.myUid(), perUidCallback, h);
+                networkCallbackRule.registerSystemDefaultNetworkCallback(systemDefaultCallback, h);
+                networkCallbackRule.registerDefaultNetworkCallbackForUid(Process.myUid(),
+                        perUidCallback, h);
             }, NETWORK_SETTINGS);
-            registerBestMatchingNetworkCallback(makeDefaultRequest(), bestMatchingCallback, h);
+            networkCallbackRule.registerBestMatchingNetworkCallback(
+                    makeDefaultRequest(), bestMatchingCallback, h);
         }
 
         Network wifiNetwork = null;
@@ -1094,24 +1093,22 @@
         // Now we should expect to get a network callback about availability of the wifi
         // network even if it was already connected as a state-based action when the callback
         // is registered.
-        wifiNetwork = callback.waitForAvailable();
+        wifiNetwork = callback.eventuallyExpect(CallbackEntry.AVAILABLE).getNetwork();
         assertNotNull("Did not receive onAvailable for TRANSPORT_WIFI request",
                 wifiNetwork);
 
-        final Network defaultNetwork = defaultTrackingCallback.waitForAvailable();
+        final Network defaultNetwork = defaultTrackingCallback.eventuallyExpect(
+                CallbackEntry.AVAILABLE).getNetwork();
         assertNotNull("Did not receive onAvailable on default network callback",
                 defaultNetwork);
 
         if (TestUtils.shouldTestSApis()) {
-            assertNotNull("Did not receive onAvailable on system default network callback",
-                    systemDefaultCallback.waitForAvailable());
-            final Network perUidNetwork = perUidCallback.waitForAvailable();
-            assertNotNull("Did not receive onAvailable on per-UID default network callback",
-                    perUidNetwork);
+            systemDefaultCallback.eventuallyExpect(CallbackEntry.AVAILABLE);
+            final Network perUidNetwork = perUidCallback.eventuallyExpect(CallbackEntry.AVAILABLE)
+                    .getNetwork();
             assertEquals(defaultNetwork, perUidNetwork);
-            final Network bestMatchingNetwork = bestMatchingCallback.waitForAvailable();
-            assertNotNull("Did not receive onAvailable on best matching network callback",
-                    bestMatchingNetwork);
+            final Network bestMatchingNetwork = bestMatchingCallback.eventuallyExpect(
+                    CallbackEntry.AVAILABLE).getNetwork();
             assertEquals(defaultNetwork, bestMatchingNetwork);
         }
     }
@@ -1123,8 +1120,8 @@
         final Handler h = new Handler(Looper.getMainLooper());
         // Verify registerSystemDefaultNetworkCallback can be accessed via
         // CONNECTIVITY_USE_RESTRICTED_NETWORKS permission.
-        runWithShellPermissionIdentity(() ->
-                        registerSystemDefaultNetworkCallback(new TestNetworkCallback(), h),
+        runWithShellPermissionIdentity(
+                () -> networkCallbackRule.registerSystemDefaultNetworkCallback(h),
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS);
     }
 
@@ -1294,15 +1291,14 @@
      */
     @AppModeFull(reason = "CHANGE_NETWORK_STATE permission can't be granted to instant apps")
     @Test
-    public void testRequestNetworkCallback() throws Exception {
-        final TestNetworkCallback callback = new TestNetworkCallback();
-        requestNetwork(new NetworkRequest.Builder()
-                .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
-                .build(), callback);
+    public void testRequestNetworkCallback() {
+        final TestableNetworkCallback callback = networkCallbackRule.requestNetwork(
+                new NetworkRequest.Builder().addCapability(
+                                NET_CAPABILITY_INTERNET)
+                        .build());
 
         // Wait to get callback for availability of internet
-        Network internetNetwork = callback.waitForAvailable();
-        assertNotNull("Did not receive NetworkCallback#onAvailable for INTERNET", internetNetwork);
+        callback.eventuallyExpect(CallbackEntry.AVAILABLE).getNetwork();
     }
 
     /**
@@ -1320,16 +1316,13 @@
             }
         }
 
-        final TestNetworkCallback callback = new TestNetworkCallback();
-        requestNetwork(new NetworkRequest.Builder().addTransportType(TRANSPORT_WIFI).build(),
-                callback, 100);
 
+        final TestableNetworkCallback callback = networkCallbackRule.requestNetwork(
+                new NetworkRequest.Builder().addTransportType(TRANSPORT_WIFI).build(),
+                100 /* timeoutMs */);
         try {
             // Wait to get callback for unavailability of requested network
-            assertTrue("Did not receive NetworkCallback#onUnavailable",
-                    callback.waitForUnavailable());
-        } catch (InterruptedException e) {
-            fail("NetworkCallback wait was interrupted.");
+            callback.eventuallyExpect(CallbackEntry.UNAVAILABLE, 2_000 /* timeoutMs */);
         } finally {
             if (previousWifiEnabledState) {
                 mCtsNetUtils.connectToWifi();
@@ -1416,40 +1409,48 @@
             final boolean useSystemDefault)
             throws Exception {
         final CompletableFuture<Network> networkFuture = new CompletableFuture<>();
-        final NetworkCallback networkCallback = new NetworkCallback() {
-            @Override
-            public void onCapabilitiesChanged(Network network, NetworkCapabilities nc) {
-                if (!nc.hasTransport(targetTransportType)) return;
 
-                final boolean metered = !nc.hasCapability(NET_CAPABILITY_NOT_METERED);
-                final boolean validated = nc.hasCapability(NET_CAPABILITY_VALIDATED);
-                if (metered == requestedMeteredness && (!waitForValidation || validated)) {
-                    networkFuture.complete(network);
+        // Registering a callback here guarantees onCapabilitiesChanged is called immediately
+        // with the current setting. Therefore, if the setting has already been changed,
+        // this method will return right away, and if not, it'll wait for the setting to change.
+        final TestableNetworkCallback networkCallback;
+        if (useSystemDefault) {
+            networkCallback = runWithShellPermissionIdentity(() -> {
+                if (isAtLeastS()) {
+                    return networkCallbackRule.registerSystemDefaultNetworkCallback(
+                            new Handler(Looper.getMainLooper()));
+                } else {
+                    // registerSystemDefaultNetworkCallback is only supported on S+.
+                    return networkCallbackRule.requestNetwork(
+                            new NetworkRequest.Builder()
+                                    .clearCapabilities()
+                                    .addCapability(NET_CAPABILITY_NOT_RESTRICTED)
+                                    .addCapability(NET_CAPABILITY_TRUSTED)
+                                    .addCapability(NET_CAPABILITY_NOT_VPN)
+                                    .addCapability(NET_CAPABILITY_INTERNET)
+                                    .build(),
+                            new TestableNetworkCallback(),
+                            new Handler(Looper.getMainLooper()));
                 }
-            }
-        };
-
-        try {
-            // Registering a callback here guarantees onCapabilitiesChanged is called immediately
-            // with the current setting. Therefore, if the setting has already been changed,
-            // this method will return right away, and if not, it'll wait for the setting to change.
-            if (useSystemDefault) {
-                runWithShellPermissionIdentity(() ->
-                                registerSystemDefaultNetworkCallback(networkCallback,
-                                        new Handler(Looper.getMainLooper())),
-                        NETWORK_SETTINGS);
-            } else {
-                registerDefaultNetworkCallback(networkCallback);
-            }
-
-            // Changing meteredness on wifi involves reconnecting, which can take several seconds
-            // (involves re-associating, DHCP...).
-            return networkFuture.get(NETWORK_CALLBACK_TIMEOUT_MS, TimeUnit.MILLISECONDS);
-        } catch (TimeoutException e) {
-            throw new AssertionError("Timed out waiting for active network metered status to "
-                    + "change to " + requestedMeteredness + " ; network = "
-                    + mCm.getActiveNetwork(), e);
+            },
+            NETWORK_SETTINGS);
+        } else {
+            networkCallback = networkCallbackRule.registerDefaultNetworkCallback();
         }
+
+        return networkCallback.eventuallyExpect(
+                CallbackEntry.NETWORK_CAPS_UPDATED,
+                // Changing meteredness on wifi involves reconnecting, which can take several
+                // seconds (involves re-associating, DHCP...).
+                NETWORK_CALLBACK_TIMEOUT_MS,
+                cb -> {
+                    final NetworkCapabilities nc = cb.getCaps();
+                    if (!nc.hasTransport(targetTransportType)) return false;
+
+                    final boolean metered = !nc.hasCapability(NET_CAPABILITY_NOT_METERED);
+                    final boolean validated = nc.hasCapability(NET_CAPABILITY_VALIDATED);
+                    return metered == requestedMeteredness && (!waitForValidation || validated);
+                }).getNetwork();
     }
 
     private Network setWifiMeteredStatusAndWait(String ssid, boolean isMetered,
@@ -2091,15 +2092,15 @@
     }
 
     private void verifyBindSocketToRestrictedNetworkDisallowed() throws Exception {
-        final TestableNetworkCallback testNetworkCb = new TestableNetworkCallback();
         final NetworkRequest testRequest = new NetworkRequest.Builder()
                 .addTransportType(NetworkCapabilities.TRANSPORT_TEST)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
+                .removeCapability(NET_CAPABILITY_TRUSTED)
+                .removeCapability(NET_CAPABILITY_NOT_RESTRICTED)
                 .setNetworkSpecifier(CompatUtil.makeTestNetworkSpecifier(
                         TEST_RESTRICTED_NW_IFACE_NAME))
                 .build();
-        runWithShellPermissionIdentity(() -> requestNetwork(testRequest, testNetworkCb),
+        final TestableNetworkCallback testNetworkCb = runWithShellPermissionIdentity(
+                () -> networkCallbackRule.requestNetwork(testRequest),
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS,
                 // CONNECTIVITY_INTERNAL is for requesting restricted network because shell does not
                 // have CONNECTIVITY_USE_RESTRICTED_NETWORKS on R.
@@ -2115,7 +2116,7 @@
                     NETWORK_CALLBACK_TIMEOUT_MS,
                     entry -> network.equals(entry.getNetwork())
                             && (!((CallbackEntry.CapabilitiesChanged) entry).getCaps()
-                            .hasCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)));
+                            .hasCapability(NET_CAPABILITY_NOT_RESTRICTED)));
             // CtsNetTestCases package doesn't hold CONNECTIVITY_USE_RESTRICTED_NETWORKS, so it
             // does not allow to bind socket to restricted network.
             assertThrows(IOException.class, () -> network.bindSocket(socket));
@@ -2238,7 +2239,7 @@
 
     private void registerCallbackAndWaitForAvailable(@NonNull final NetworkRequest request,
             @NonNull final TestableNetworkCallback cb) {
-        registerNetworkCallback(request, cb);
+        networkCallbackRule.registerNetworkCallback(request, cb);
         waitForAvailable(cb);
     }
 
@@ -2346,21 +2347,13 @@
 
     private void verifySsidFromCallbackNetworkCapabilities(@NonNull String ssid, boolean hasSsid)
             throws Exception {
-        final CompletableFuture<NetworkCapabilities> foundNc = new CompletableFuture();
-        final NetworkCallback callback = new NetworkCallback() {
-            @Override
-            public void onCapabilitiesChanged(Network network, NetworkCapabilities nc) {
-                foundNc.complete(nc);
-            }
-        };
-
-        registerNetworkCallback(makeWifiNetworkRequest(), callback);
+        final TestableNetworkCallback callback =
+                networkCallbackRule.registerNetworkCallback(makeWifiNetworkRequest());
         // Registering a callback here guarantees onCapabilitiesChanged is called immediately
         // because WiFi network should be connected.
-        final NetworkCapabilities nc =
-                foundNc.get(NETWORK_CALLBACK_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+        final NetworkCapabilities nc = callback.eventuallyExpect(
+                CallbackEntry.NETWORK_CAPS_UPDATED, NETWORK_CALLBACK_TIMEOUT_MS).getCaps();
         // Verify if ssid is contained in the NetworkCapabilities received from callback.
-        assertNotNull("NetworkCapabilities of the network is null", nc);
         assertEquals(hasSsid, Pattern.compile(ssid).matcher(nc.toString()).find());
     }
 
@@ -2389,8 +2382,8 @@
         final NetworkRequest testRequest = new NetworkRequest.Builder()
                 .addTransportType(TRANSPORT_TEST)
                 // Test networks do not have NOT_VPN or TRUSTED capabilities by default
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_VPN)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
+                .removeCapability(NET_CAPABILITY_NOT_VPN)
+                .removeCapability(NET_CAPABILITY_TRUSTED)
                 .setNetworkSpecifier(CompatUtil.makeTestNetworkSpecifier(
                         testNetworkInterface.getInterfaceName()))
                 .build();
@@ -2399,14 +2392,15 @@
         final TestableNetworkCallback callback = new TestableNetworkCallback();
         final Handler handler = new Handler(Looper.getMainLooper());
         assertThrows(SecurityException.class,
-                () -> requestBackgroundNetwork(testRequest, callback, handler));
+                () -> networkCallbackRule.requestBackgroundNetwork(testRequest, callback, handler));
 
         Network testNetwork = null;
         try {
             // Request background test network via Shell identity which has NETWORK_SETTINGS
             // permission granted.
             runWithShellPermissionIdentity(
-                    () -> requestBackgroundNetwork(testRequest, callback, handler),
+                    () -> networkCallbackRule.requestBackgroundNetwork(
+                            testRequest, callback, handler),
                     new String[] { android.Manifest.permission.NETWORK_SETTINGS });
 
             // Register the test network agent which has no foreground request associated to it.
@@ -2499,9 +2493,10 @@
         final int otherUid = UserHandle.getUid(5, Process.FIRST_APPLICATION_UID);
         final Handler handler = new Handler(Looper.getMainLooper());
 
-        registerDefaultNetworkCallback(myUidCallback, handler);
-        runWithShellPermissionIdentity(() -> registerDefaultNetworkCallbackForUid(
-                otherUid, otherUidCallback, handler), NETWORK_SETTINGS);
+        networkCallbackRule.registerDefaultNetworkCallback(myUidCallback, handler);
+        runWithShellPermissionIdentity(
+                () -> networkCallbackRule.registerDefaultNetworkCallbackForUid(
+                        otherUid, otherUidCallback, handler), NETWORK_SETTINGS);
 
         final Network defaultNetwork = myUidCallback.expect(CallbackEntry.AVAILABLE).getNetwork();
         final List<DetailedBlockedStatusCallback> allCallbacks =
@@ -2557,14 +2552,14 @@
         assertNotNull(info);
         assertEquals(DetailedState.CONNECTED, info.getDetailedState());
 
-        final TestableNetworkCallback callback = new TestableNetworkCallback();
+        final TestableNetworkCallback callback;
         try {
             mCmShim.setLegacyLockdownVpnEnabled(true);
 
             // setLegacyLockdownVpnEnabled is asynchronous and only takes effect when the
             // ConnectivityService handler thread processes it. Ensure it has taken effect by doing
             // something that blocks until the handler thread is idle.
-            registerDefaultNetworkCallback(callback);
+            callback = networkCallbackRule.registerDefaultNetworkCallback();
             waitForAvailable(callback);
 
             // Test one of the effects of setLegacyLockdownVpnEnabled: the fact that any NetworkInfo
@@ -2831,9 +2826,9 @@
     private void registerTestOemNetworkPreferenceCallbacks(
             @NonNull final TestableNetworkCallback defaultCallback,
             @NonNull final TestableNetworkCallback systemDefaultCallback) {
-        registerDefaultNetworkCallback(defaultCallback);
+        networkCallbackRule.registerDefaultNetworkCallback(defaultCallback);
         runWithShellPermissionIdentity(() ->
-                registerSystemDefaultNetworkCallback(systemDefaultCallback,
+                networkCallbackRule.registerSystemDefaultNetworkCallback(systemDefaultCallback,
                         new Handler(Looper.getMainLooper())), NETWORK_SETTINGS);
     }
 
@@ -2949,18 +2944,18 @@
                         + " unless device supports WiFi",
                 mPackageManager.hasSystemFeature(FEATURE_WIFI));
 
-        final TestNetworkCallback cb = new TestNetworkCallback();
         try {
             // Wait for partial connectivity to be detected on the network
             final Network network = preparePartialConnectivity();
 
-            requestNetwork(makeWifiNetworkRequest(), cb);
+            final TestableNetworkCallback cb = networkCallbackRule.requestNetwork(
+                    makeWifiNetworkRequest());
             runAsShell(NETWORK_SETTINGS, () -> {
                 // The always bit is verified in NetworkAgentTest
                 mCm.setAcceptPartialConnectivity(network, false /* accept */, false /* always */);
             });
             // Reject partial connectivity network should cause the network being torn down
-            assertEquals(network, cb.waitForLost());
+            assertEquals(network, cb.eventuallyExpect(CallbackEntry.LOST).getNetwork());
         } finally {
             mHttpServer.stop();
             // Wifi will not automatically reconnect to the network. ensureWifiDisconnected cannot
@@ -2988,7 +2983,6 @@
         assumeTrue("testAcceptPartialConnectivity_validatedNetwork cannot execute"
                         + " unless device supports WiFi and telephony", canRunTest);
 
-        final TestableNetworkCallback wifiCb = new TestableNetworkCallback();
         try {
             // Ensure at least one default network candidate connected.
             networkCallbackRule.requestCell();
@@ -2998,7 +2992,8 @@
             // guarantee that it won't become the default in the future.
             assertNotEquals(wifiNetwork, mCm.getActiveNetwork());
 
-            registerNetworkCallback(makeWifiNetworkRequest(), wifiCb);
+            final TestableNetworkCallback wifiCb = networkCallbackRule.registerNetworkCallback(
+                    makeWifiNetworkRequest());
             runAsShell(NETWORK_SETTINGS, () -> {
                 mCm.setAcceptUnvalidated(wifiNetwork, false /* accept */, false /* always */);
             });
@@ -3025,8 +3020,6 @@
         assumeTrue("testSetAvoidUnvalidated cannot execute"
                 + " unless device supports WiFi and telephony", canRunTest);
 
-        final TestableNetworkCallback wifiCb = new TestableNetworkCallback();
-        final TestableNetworkCallback defaultCb = new TestableNetworkCallback();
         final int previousAvoidBadWifi =
                 ConnectivitySettingsManager.getNetworkAvoidBadWifi(mContext);
 
@@ -3036,8 +3029,10 @@
             final Network cellNetwork = networkCallbackRule.requestCell();
             final Network wifiNetwork = prepareValidatedNetwork();
 
-            registerDefaultNetworkCallback(defaultCb);
-            registerNetworkCallback(makeWifiNetworkRequest(), wifiCb);
+            final TestableNetworkCallback defaultCb =
+                    networkCallbackRule.registerDefaultNetworkCallback();
+            final TestableNetworkCallback wifiCb = networkCallbackRule.registerNetworkCallback(
+                    makeWifiNetworkRequest());
 
             // Verify wifi is the default network.
             defaultCb.eventuallyExpect(CallbackEntry.AVAILABLE, NETWORK_CALLBACK_TIMEOUT_MS,
@@ -3100,20 +3095,11 @@
         });
     }
 
-    private Network expectNetworkHasCapability(Network network, int expectedNetCap, long timeout)
-            throws Exception {
-        final CompletableFuture<Network> future = new CompletableFuture();
-        final NetworkCallback cb = new NetworkCallback() {
-            @Override
-            public void onCapabilitiesChanged(Network n, NetworkCapabilities nc) {
-                if (n.equals(network) && nc.hasCapability(expectedNetCap)) {
-                    future.complete(network);
-                }
-            }
-        };
-
-        registerNetworkCallback(new NetworkRequest.Builder().build(), cb);
-        return future.get(timeout, TimeUnit.MILLISECONDS);
+    private Network expectNetworkHasCapability(Network network, int expectedNetCap, long timeout) {
+        return networkCallbackRule.registerNetworkCallback(new NetworkRequest.Builder().build())
+                .eventuallyExpect(CallbackEntry.NETWORK_CAPS_UPDATED, timeout,
+                        cb -> cb.getNetwork().equals(network)
+                                && cb.getCaps().hasCapability(expectedNetCap)).getNetwork();
     }
 
     private void prepareHttpServer() throws Exception {
@@ -3265,12 +3251,13 @@
         // For testing mobile data preferred uids feature, it needs both wifi and cell network.
         final Network wifiNetwork = mCtsNetUtils.ensureWifiConnected();
         final Network cellNetwork = networkCallbackRule.requestCell();
-        final TestableNetworkCallback defaultTrackingCb = new TestableNetworkCallback();
-        final TestableNetworkCallback systemDefaultCb = new TestableNetworkCallback();
         final Handler h = new Handler(Looper.getMainLooper());
-        runWithShellPermissionIdentity(() -> registerSystemDefaultNetworkCallback(
-                systemDefaultCb, h), NETWORK_SETTINGS);
-        registerDefaultNetworkCallback(defaultTrackingCb);
+        final TestableNetworkCallback systemDefaultCb = runWithShellPermissionIdentity(
+                () -> networkCallbackRule.registerSystemDefaultNetworkCallback(h),
+                NETWORK_SETTINGS);
+
+        final TestableNetworkCallback defaultTrackingCb =
+                networkCallbackRule.registerDefaultNetworkCallback();
 
         try {
             // CtsNetTestCases uid is not listed in MOBILE_DATA_PREFERRED_UIDS setting, so the
@@ -3339,7 +3326,7 @@
         // Create test network agent with restricted network.
         final NetworkCapabilities nc = new NetworkCapabilities.Builder()
                 .addTransportType(NetworkCapabilities.TRANSPORT_TEST)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
+                .removeCapability(NET_CAPABILITY_NOT_RESTRICTED)
                 .setNetworkSpecifier(CompatUtil.makeTestNetworkSpecifier(
                         TEST_RESTRICTED_NW_IFACE_NAME))
                 .build();
@@ -3373,23 +3360,23 @@
                 mContext, originalUidsAllowedOnRestrictedNetworks), NETWORK_SETTINGS);
 
         // File a restricted network request with permission first to hold the connection.
-        final TestableNetworkCallback testNetworkCb = new TestableNetworkCallback();
         final NetworkRequest testRequest = new NetworkRequest.Builder()
                 .addTransportType(NetworkCapabilities.TRANSPORT_TEST)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
+                .removeCapability(NET_CAPABILITY_TRUSTED)
+                .removeCapability(NET_CAPABILITY_NOT_RESTRICTED)
                 .setNetworkSpecifier(CompatUtil.makeTestNetworkSpecifier(
                         TEST_RESTRICTED_NW_IFACE_NAME))
                 .build();
-        runWithShellPermissionIdentity(() -> requestNetwork(testRequest, testNetworkCb),
+        final TestableNetworkCallback testNetworkCb = runWithShellPermissionIdentity(
+                () -> networkCallbackRule.requestNetwork(testRequest),
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS);
 
         // File another restricted network request without permission.
         final TestableNetworkCallback restrictedNetworkCb = new TestableNetworkCallback();
         final NetworkRequest restrictedRequest = new NetworkRequest.Builder()
                 .addTransportType(NetworkCapabilities.TRANSPORT_TEST)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
-                .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
+                .removeCapability(NET_CAPABILITY_TRUSTED)
+                .removeCapability(NET_CAPABILITY_NOT_RESTRICTED)
                 .setNetworkSpecifier(CompatUtil.makeTestNetworkSpecifier(
                         TEST_RESTRICTED_NW_IFACE_NAME))
                 .build();
@@ -3406,7 +3393,7 @@
                     NETWORK_CALLBACK_TIMEOUT_MS,
                     entry -> network.equals(entry.getNetwork())
                             && (!((CallbackEntry.CapabilitiesChanged) entry).getCaps()
-                            .hasCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)));
+                            .hasCapability(NET_CAPABILITY_NOT_RESTRICTED)));
             // CtsNetTestCases package doesn't hold CONNECTIVITY_USE_RESTRICTED_NETWORKS, so it
             // does not allow to bind socket to restricted network.
             assertThrows(IOException.class, () -> network.bindSocket(socket));
@@ -3424,13 +3411,13 @@
 
             if (TestUtils.shouldTestTApis()) {
                 // Uid is in allowed list. Try file network request again.
-                requestNetwork(restrictedRequest, restrictedNetworkCb);
+                networkCallbackRule.requestNetwork(restrictedRequest, restrictedNetworkCb);
                 // Verify that the network is restricted.
                 restrictedNetworkCb.eventuallyExpect(CallbackEntry.NETWORK_CAPS_UPDATED,
                         NETWORK_CALLBACK_TIMEOUT_MS,
                         entry -> network.equals(entry.getNetwork())
                                 && (!((CallbackEntry.CapabilitiesChanged) entry).getCaps()
-                                .hasCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)));
+                                .hasCapability(NET_CAPABILITY_NOT_RESTRICTED)));
             }
         } finally {
             agent.unregister();
@@ -3728,58 +3715,4 @@
         // shims, and @IgnoreUpTo does not check that.
         assumeTrue(TestUtils.shouldTestSApis());
     }
-
-    private void unregisterRegisteredCallbacks() {
-        for (NetworkCallback callback: mRegisteredCallbacks) {
-            mCm.unregisterNetworkCallback(callback);
-        }
-    }
-
-    private void registerDefaultNetworkCallback(NetworkCallback callback) {
-        mCm.registerDefaultNetworkCallback(callback);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void registerDefaultNetworkCallback(NetworkCallback callback, Handler handler) {
-        mCm.registerDefaultNetworkCallback(callback, handler);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void registerNetworkCallback(NetworkRequest request, NetworkCallback callback) {
-        mCm.registerNetworkCallback(request, callback);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void registerSystemDefaultNetworkCallback(NetworkCallback callback, Handler handler) {
-        mCmShim.registerSystemDefaultNetworkCallback(callback, handler);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void registerDefaultNetworkCallbackForUid(int uid, NetworkCallback callback,
-            Handler handler) throws Exception {
-        mCmShim.registerDefaultNetworkCallbackForUid(uid, callback, handler);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void requestNetwork(NetworkRequest request, NetworkCallback callback) {
-        mCm.requestNetwork(request, callback);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void requestNetwork(NetworkRequest request, NetworkCallback callback, int timeoutSec) {
-        mCm.requestNetwork(request, callback, timeoutSec);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void registerBestMatchingNetworkCallback(NetworkRequest request,
-            NetworkCallback callback, Handler handler) {
-        mCm.registerBestMatchingNetworkCallback(request, callback, handler);
-        mRegisteredCallbacks.add(callback);
-    }
-
-    private void requestBackgroundNetwork(NetworkRequest request, NetworkCallback callback,
-            Handler handler) throws Exception {
-        mCmShim.requestBackgroundNetwork(request, callback, handler);
-        mRegisteredCallbacks.add(callback);
-    }
 }
diff --git a/tests/cts/net/src/android/net/cts/IpSecManagerTest.java b/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
index a40ed0f..b703f77 100644
--- a/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
+++ b/tests/cts/net/src/android/net/cts/IpSecManagerTest.java
@@ -81,6 +81,7 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.modules.utils.build.SdkLevel;
+import com.android.testutils.ConnectivityModuleTest;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo;
 import com.android.testutils.SkipMainlinePresubmit;
@@ -101,6 +102,7 @@
 import java.util.Map.Entry;
 import java.util.Set;
 
+@ConnectivityModuleTest
 @RunWith(AndroidJUnit4.class)
 @AppModeFull(reason = "Socket cannot bind in instant app mode")
 public class IpSecManagerTest extends IpSecBaseTest {
@@ -444,6 +446,11 @@
             long uidTxDelta = 0;
             long uidRxDelta = 0;
             for (int i = 0; i < 100; i++) {
+                // Clear TrafficStats cache is needed to avoid rate-limit caching for
+                // TrafficStats API results on V+ devices.
+                if (SdkLevel.isAtLeastV()) {
+                    runAsShell(NETWORK_SETTINGS, () -> TrafficStats.clearRateLimitCaches());
+                }
                 uidTxDelta = TrafficStats.getUidTxPackets(Os.getuid()) - uidTxPackets;
                 uidRxDelta = TrafficStats.getUidRxPackets(Os.getuid()) - uidRxPackets;
 
@@ -518,6 +525,11 @@
         }
 
         private static void initStatsChecker() throws Exception {
+            // Clear TrafficStats cache is needed to avoid rate-limit caching for
+            // TrafficStats API results on V+ devices.
+            if (SdkLevel.isAtLeastV()) {
+                runAsShell(NETWORK_SETTINGS, () -> TrafficStats.clearRateLimitCaches());
+            }
             uidTxBytes = TrafficStats.getUidTxBytes(Os.getuid());
             uidRxBytes = TrafficStats.getUidRxBytes(Os.getuid());
             uidTxPackets = TrafficStats.getUidTxPackets(Os.getuid());
diff --git a/tests/integration/src/com/android/server/net/integrationtests/NetworkStatsIntegrationTest.kt b/tests/integration/src/com/android/server/net/integrationtests/NetworkStatsIntegrationTest.kt
index 52e502d..4780c5d 100644
--- a/tests/integration/src/com/android/server/net/integrationtests/NetworkStatsIntegrationTest.kt
+++ b/tests/integration/src/com/android/server/net/integrationtests/NetworkStatsIntegrationTest.kt
@@ -38,6 +38,7 @@
 import android.os.Build
 import android.os.Process
 import androidx.test.platform.app.InstrumentationRegistry
+import com.android.modules.utils.build.SdkLevel
 import com.android.server.net.integrationtests.NetworkStatsIntegrationTest.Direction.DOWNLOAD
 import com.android.server.net.integrationtests.NetworkStatsIntegrationTest.Direction.UPLOAD
 import com.android.testutils.DevSdkIgnoreRule.IgnoreUpTo
@@ -214,6 +215,11 @@
         // In practice, for one way 10k download payload, the download usage is about
         // 11222~12880 bytes, with 14~17 packets. And the upload usage is about 1279~1626 bytes
         // with 14~17 packets, which is majorly contributed by TCP ACK packets.
+        // Clear TrafficStats cache is needed to avoid rate-limit caching for
+        // TrafficStats API results on V+ devices.
+        if (SdkLevel.isAtLeastV()) {
+            TrafficStats.clearRateLimitCaches()
+        }
         val snapshotAfterDownload = StatsSnapshot(context, internalInterfaceName)
         val (expectedDownloadLower, expectedDownloadUpper) = getExpectedStatsBounds(
             TEST_DOWNLOAD_SIZE,
@@ -236,6 +242,9 @@
         )
 
         // Verify upload data usage accounting.
+        if (SdkLevel.isAtLeastV()) {
+            TrafficStats.clearRateLimitCaches()
+        }
         val snapshotAfterUpload = StatsSnapshot(context, internalInterfaceName)
         val (expectedUploadLower, expectedUploadUpper) = getExpectedStatsBounds(
             TEST_UPLOAD_SIZE,
diff --git a/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt b/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt
index 3f6e88d..aa28e5a 100644
--- a/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt
+++ b/tests/unit/java/com/android/metrics/NetworkNsdReportedMetricsTest.kt
@@ -231,8 +231,10 @@
         val clientId = 99
         val transactionId = 100
         val durationMs = 10L
+        val sentQueryCount = 10
         val metrics = NetworkNsdReportedMetrics(clientId, deps)
-        metrics.reportServiceResolutionStop(true /* isLegacy */, transactionId, durationMs)
+        metrics.reportServiceResolutionStop(
+                true /* isLegacy */, transactionId, durationMs, sentQueryCount)
 
         val eventCaptor = ArgumentCaptor.forClass(NetworkNsdReported::class.java)
         verify(deps).statsWrite(eventCaptor.capture())
@@ -243,6 +245,7 @@
             assertEquals(NsdEventType.NET_RESOLVE, it.type)
             assertEquals(MdnsQueryResult.MQR_SERVICE_RESOLUTION_STOP, it.queryResult)
             assertEquals(durationMs, it.eventDurationMillisec)
+            assertEquals(sentQueryCount, it.sentQueryCount)
         }
     }
 
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 7822fe0..ce49533 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -1721,8 +1721,6 @@
     private void mockUidNetworkingBlocked() {
         doAnswer(i -> isUidBlocked(mBlockedReasons, i.getArgument(1))
         ).when(mNetworkPolicyManager).isUidNetworkingBlocked(anyInt(), anyBoolean());
-        doAnswer(i -> isUidBlocked(mBlockedReasons, i.getArgument(1))
-        ).when(mBpfNetMaps).isUidNetworkingBlocked(anyInt(), anyBoolean());
     }
 
     private boolean isUidBlocked(int blockedReasons, boolean meteredNetwork) {
diff --git a/tests/unit/java/com/android/server/NsdServiceTest.java b/tests/unit/java/com/android/server/NsdServiceTest.java
index d91e29c..5731d01 100644
--- a/tests/unit/java/com/android/server/NsdServiceTest.java
+++ b/tests/unit/java/com/android/server/NsdServiceTest.java
@@ -904,7 +904,7 @@
                 request.getServiceName().equals(ns.getServiceName())
                         && request.getServiceType().equals(ns.getServiceType())));
         verify(mMetrics).reportServiceResolutionStop(
-                true /* isLegacy */, resolveId, 10L /* durationMs */);
+                true /* isLegacy */, resolveId, 10L /* durationMs */, 0 /* sentQueryCount */);
     }
 
     @Test
@@ -978,7 +978,7 @@
                 request.getServiceName().equals(ns.getServiceName())
                         && request.getServiceType().equals(ns.getServiceType())));
         verify(mMetrics).reportServiceResolutionStop(
-                true /* isLegacy */, getAddrId, 10L /* durationMs */);
+                true /* isLegacy */, getAddrId, 10L /* durationMs */,  0 /* sentQueryCount */);
     }
 
     private void verifyUpdatedServiceInfo(NsdServiceInfo info, String serviceName,
@@ -1679,20 +1679,23 @@
         // Subtypes are not used for resolution, only for discovery
         assertEquals(Collections.emptyList(), optionsCaptor.getValue().getSubtypes());
 
+        final MdnsListener listener = listenerCaptor.getValue();
+        // Callbacks for query sent.
+        listener.onDiscoveryQuerySent(Collections.emptyList(), 1 /* transactionId */);
+
         doReturn(TEST_TIME_MS + 10L).when(mClock).elapsedRealtime();
         client.stopServiceResolution(resolveListener);
         waitForIdle();
 
         // Verify the listener has been unregistered.
-        final MdnsListener listener = listenerCaptor.getValue();
         verify(mDiscoveryManager, timeout(TIMEOUT_MS))
                 .unregisterListener(eq(constructedServiceType), eq(listener));
         verify(resolveListener, timeout(TIMEOUT_MS)).onResolutionStopped(argThat(ns ->
                 request.getServiceName().equals(ns.getServiceName())
                         && request.getServiceType().equals(ns.getServiceType())));
         verify(mSocketProvider, timeout(CLEANUP_DELAY_MS + TIMEOUT_MS)).requestStopWhenInactive();
-        verify(mMetrics).reportServiceResolutionStop(
-                false /* isLegacy */, listener.mTransactionId, 10L /* durationMs */);
+        verify(mMetrics).reportServiceResolutionStop(false /* isLegacy */, listener.mTransactionId,
+                10L /* durationMs */, 1 /* sentQueryCount */);
     }
 
     @Test
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index d4f5619..6425daa 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -67,11 +67,14 @@
 import static com.android.server.net.NetworkStatsEventLogger.POLL_REASON_RAT_CHANGED;
 import static com.android.server.net.NetworkStatsEventLogger.PollEvent.pollReasonNameOf;
 import static com.android.server.net.NetworkStatsService.ACTION_NETWORK_STATS_POLL;
+import static com.android.server.net.NetworkStatsService.DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS;
+import static com.android.server.net.NetworkStatsService.DEFAULT_TRAFFIC_STATS_CACHE_MAX_ENTRIES;
 import static com.android.server.net.NetworkStatsService.NETSTATS_FASTDATAINPUT_FALLBACKS_COUNTER_NAME;
 import static com.android.server.net.NetworkStatsService.NETSTATS_FASTDATAINPUT_SUCCESSES_COUNTER_NAME;
 import static com.android.server.net.NetworkStatsService.NETSTATS_IMPORT_ATTEMPTS_COUNTER_NAME;
 import static com.android.server.net.NetworkStatsService.NETSTATS_IMPORT_FALLBACKS_COUNTER_NAME;
 import static com.android.server.net.NetworkStatsService.NETSTATS_IMPORT_SUCCESSES_COUNTER_NAME;
+import static com.android.server.net.NetworkStatsService.TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG;
 import static com.android.testutils.DevSdkIgnoreRuleKt.SC_V2;
 
 import static org.junit.Assert.assertEquals;
@@ -116,6 +119,7 @@
 import android.net.TestNetworkSpecifier;
 import android.net.TetherStatsParcel;
 import android.net.TetheringManager;
+import android.net.TrafficStats;
 import android.net.UnderlyingNetworkInfo;
 import android.net.netstats.provider.INetworkStatsProviderCallback;
 import android.net.wifi.WifiInfo;
@@ -125,6 +129,7 @@
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.PowerManager;
+import android.os.Process;
 import android.os.SimpleClock;
 import android.provider.Settings;
 import android.system.ErrnoException;
@@ -159,12 +164,15 @@
 import com.android.testutils.HandlerUtils;
 import com.android.testutils.TestBpfMap;
 import com.android.testutils.TestableNetworkStatsProviderBinder;
+import com.android.testutils.com.android.testutils.SetFeatureFlagsRule;
+import com.android.testutils.com.android.testutils.SetFeatureFlagsRule.FeatureFlag;
 
 import libcore.testing.io.TestIoUtils;
 
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Ignore;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -189,6 +197,7 @@
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Function;
 
 /**
  * Tests for {@link NetworkStatsService}.
@@ -202,6 +211,7 @@
 // NetworkStatsService is not updatable before T, so tests do not need to be backwards compatible
 @DevSdkIgnoreRule.IgnoreUpTo(SC_V2)
 public class NetworkStatsServiceTest extends NetworkStatsBaseTest {
+
     private static final String TAG = "NetworkStatsServiceTest";
 
     private static final long TEST_START = 1194220800000L;
@@ -295,6 +305,16 @@
     private Boolean mIsDebuggable;
     private HandlerThread mObserverHandlerThread;
     final TestDependencies mDeps = new TestDependencies();
+    final HashMap<String, Boolean> mFeatureFlags = new HashMap<>();
+
+    // This will set feature flags from @FeatureFlag annotations
+    // into the map before setUp() runs.
+    @Rule
+    public final SetFeatureFlagsRule mSetFeatureFlagsRule =
+            new SetFeatureFlagsRule((name, enabled) -> {
+                mFeatureFlags.put(name, enabled);
+                return null;
+            });
 
     private class MockContext extends BroadcastInterceptingContext {
         private final Context mBaseContext;
@@ -395,8 +415,6 @@
 
         mElapsedRealtime = 0L;
 
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
         mService.systemReady();
         // Verify that system ready fetches realtime stats
@@ -432,6 +450,7 @@
 
     class TestDependencies extends NetworkStatsService.Dependencies {
         private int mCompareStatsInvocation = 0;
+        private NetworkStats.Entry mMockedTrafficStatsNativeStat = null;
 
         @Override
         public File getLegacyStatsDir() {
@@ -573,6 +592,43 @@
         public boolean supportEventLogger(@NonNull Context cts) {
             return true;
         }
+
+        @Override
+        public boolean supportTrafficStatsRateLimitCache(Context ctx) {
+            return mFeatureFlags.getOrDefault(TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG, false);
+        }
+
+        @Override
+        public int getTrafficStatsRateLimitCacheExpiryDuration() {
+            return DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS;
+        }
+
+        @Override
+        public int getTrafficStatsRateLimitCacheMaxEntries() {
+            return DEFAULT_TRAFFIC_STATS_CACHE_MAX_ENTRIES;
+        }
+
+        @Nullable
+        @Override
+        public NetworkStats.Entry nativeGetTotalStat() {
+            return mMockedTrafficStatsNativeStat;
+        }
+
+        @Nullable
+        @Override
+        public NetworkStats.Entry nativeGetIfaceStat(String iface) {
+            return mMockedTrafficStatsNativeStat;
+        }
+
+        @Nullable
+        @Override
+        public NetworkStats.Entry nativeGetUidStat(int uid) {
+            return mMockedTrafficStatsNativeStat;
+        }
+
+        public void setNativeStat(NetworkStats.Entry entry) {
+            mMockedTrafficStatsNativeStat = entry;
+        }
     }
 
     @After
@@ -729,8 +785,6 @@
         assertStatsFilesExist(true);
 
         // boot through serviceReady() again
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
 
         mService.systemReady();
@@ -2111,8 +2165,6 @@
                 getLegacyCollection(PREFIX_UID_TAG, true /* includeTags */));
 
         // Mock zero usage and boot through serviceReady(), verify there is no imported data.
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
         mService.systemReady();
         assertStatsFilesExist(false);
@@ -2124,8 +2176,6 @@
         assertStatsFilesExist(false);
 
         // Boot through systemReady() again.
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
         mService.systemReady();
 
@@ -2199,8 +2249,6 @@
                 getLegacyCollection(PREFIX_UID_TAG, true /* includeTags */));
 
         // Mock zero usage and boot through serviceReady(), verify there is no imported data.
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
         mService.systemReady();
         assertStatsFilesExist(false);
@@ -2212,8 +2260,6 @@
         assertStatsFilesExist(false);
 
         // Boot through systemReady() again.
-        mockDefaultSettings();
-        mockNetworkStatsUidDetail(buildEmptyStats());
         prepareForSystemReady();
         mService.systemReady();
 
@@ -2365,6 +2411,68 @@
         assertUidTotal(sTemplateWifi, UID_GREEN, 64L, 3L, 1024L, 8L, 0);
     }
 
+    @FeatureFlag(name = TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG, enabled = false)
+    @Test
+    public void testTrafficStatsRateLimitCache_disabled() throws Exception {
+        doTestTrafficStatsRateLimitCache(false /* cacheEnabled */);
+    }
+
+    @FeatureFlag(name = TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG)
+    @Test
+    public void testTrafficStatsRateLimitCache_enabled() throws Exception {
+        doTestTrafficStatsRateLimitCache(true /* cacheEnabled */);
+    }
+
+    private void doTestTrafficStatsRateLimitCache(boolean cacheEnabled) throws Exception {
+        mockDefaultSettings();
+        // Calling uid is not injected into the service, use the real uid to pass the caller check.
+        final int myUid = Process.myUid();
+        mockTrafficStatsValues(64L, 3L, 1024L, 8L);
+        assertTrafficStatsValues(TEST_IFACE, myUid, 64L, 3L, 1024L, 8L);
+
+        // Verify the values are cached.
+        incrementCurrentTime(DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS / 2);
+        mockTrafficStatsValues(65L, 8L, 1055L, 9L);
+        if (cacheEnabled) {
+            assertTrafficStatsValues(TEST_IFACE, myUid, 64L, 3L, 1024L, 8L);
+        } else {
+            assertTrafficStatsValues(TEST_IFACE, myUid, 65L, 8L, 1055L, 9L);
+        }
+
+        // Verify the values are updated after cache expiry.
+        incrementCurrentTime(DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS);
+        assertTrafficStatsValues(TEST_IFACE, myUid, 65L, 8L, 1055L, 9L);
+    }
+
+    private void mockTrafficStatsValues(long rxBytes, long rxPackets,
+            long txBytes, long txPackets) {
+        // In practice, keys and operations are not used and filled with default values when
+        // returned by JNI layer.
+        final NetworkStats.Entry entry = new NetworkStats.Entry(IFACE_ALL, UID_ALL, SET_DEFAULT,
+                TAG_NONE, METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO,
+                rxBytes, rxPackets, txBytes, txPackets, 0L);
+        mDeps.setNativeStat(entry);
+    }
+
+    // Assert for 3 different API return values respectively.
+    private void assertTrafficStatsValues(String iface, int uid, long rxBytes, long rxPackets,
+            long txBytes, long txPackets) {
+        assertTrafficStatsValuesThat(rxBytes, rxPackets, txBytes, txPackets,
+                (type) -> mService.getTotalStats(type));
+        assertTrafficStatsValuesThat(rxBytes, rxPackets, txBytes, txPackets,
+                (type) -> mService.getIfaceStats(iface, type));
+        assertTrafficStatsValuesThat(rxBytes, rxPackets, txBytes, txPackets,
+                (type) -> mService.getUidStats(uid, type));
+    }
+
+    private void assertTrafficStatsValuesThat(long rxBytes, long rxPackets, long txBytes,
+            long txPackets, Function<Integer, Long> fetcher) {
+        assertEquals(rxBytes, (long) fetcher.apply(TrafficStats.TYPE_RX_BYTES));
+        assertEquals(rxPackets, (long) fetcher.apply(TrafficStats.TYPE_RX_PACKETS));
+        assertEquals(txBytes, (long) fetcher.apply(TrafficStats.TYPE_TX_BYTES));
+        assertEquals(txPackets, (long) fetcher.apply(TrafficStats.TYPE_TX_PACKETS));
+    }
+
     private void assertShouldRunComparison(boolean expected, boolean isDebuggable) {
         assertEquals("shouldRunComparison (debuggable=" + isDebuggable + "): ",
                 expected, mService.shouldRunComparison());
@@ -2429,6 +2537,8 @@
     }
 
     private void prepareForSystemReady() throws Exception {
+        mockDefaultSettings();
+        mockNetworkStatsUidDetail(buildEmptyStats());
         mockNetworkStatsSummary(buildEmptyStats());
     }
 
diff --git a/thread/tests/integration/src/android/net/thread/ServiceDiscoveryTest.java b/thread/tests/integration/src/android/net/thread/ServiceDiscoveryTest.java
index 5a8d21f..e10f134 100644
--- a/thread/tests/integration/src/android/net/thread/ServiceDiscoveryTest.java
+++ b/thread/tests/integration/src/android/net/thread/ServiceDiscoveryTest.java
@@ -265,6 +265,7 @@
     }
 
     @Test
+    @Ignore("TODO: b/333806992 - Enable when it's not flaky at all")
     public void advertisingProxy_srpClientUnregistersService_serviceIsNotDiscoverableByMdns()
             throws Exception {
         /*