Merge changes from topic "rate-limit-trafficstats" into main

* changes:
  Make cache duration and max entries configurable
  Create TestRule: SetFeatureFlagsRule
  Rate-limit on TrafficStats#get*[Bytes|Packets] API calls
  Clear TrafficStats cache before getting readings from TrafficStats
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/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/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/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/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());
     }