Check compat change flag for rate-limit cache

This allows apps built against V but run on older platform
can adopt this feature. This keep compatibility for older
apps which does not built against newest platform.

Test: atest ConnectivityCoverageTests:android.net.connectivity.com.android.server.net.NetworkStatsServiceTest
Fix: 339154075
Change-Id: Ia12997595579190ae0d65ad36ce39d011dde37ce
diff --git a/framework/src/android/net/connectivity/ConnectivityCompatChanges.java b/framework/src/android/net/connectivity/ConnectivityCompatChanges.java
index 56537d9..733cc17 100644
--- a/framework/src/android/net/connectivity/ConnectivityCompatChanges.java
+++ b/framework/src/android/net/connectivity/ConnectivityCompatChanges.java
@@ -118,6 +118,20 @@
     @EnabledAfter(targetSdkVersion = Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
     public static final long NETWORKINFO_WITHOUT_INTERNET_BLOCKED = 333340911L;
 
+    /**
+     * Enable caching for TrafficStats#get* APIs.
+     *
+     * Apps targeting Android V or later or running on Android V or later may take up to several
+     * seconds to see the updated results.
+     * Apps targeting lower android SDKs do not see cached result for backward compatibility,
+     * results of TrafficStats#get* APIs are reflecting network statistics immediately.
+     *
+     * @hide
+     */
+    @ChangeId
+    @EnabledAfter(targetSdkVersion = Build.VERSION_CODES.UPSIDE_DOWN_CAKE)
+    public static final long ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE = 74210811L;
+
     private ConnectivityCompatChanges() {
     }
 }
diff --git a/service-t/src/com/android/server/net/NetworkStatsService.java b/service-t/src/com/android/server/net/NetworkStatsService.java
index d43c3da..9a51597 100644
--- a/service-t/src/com/android/server/net/NetworkStatsService.java
+++ b/service-t/src/com/android/server/net/NetworkStatsService.java
@@ -57,6 +57,7 @@
 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.connectivity.ConnectivityCompatChanges.ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE;
 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;
@@ -91,6 +92,7 @@
 import android.app.AlarmManager;
 import android.app.BroadcastOptions;
 import android.app.PendingIntent;
+import android.app.compat.CompatChanges;
 import android.app.usage.NetworkStatsManager;
 import android.content.ApexEnvironment;
 import android.content.BroadcastReceiver;
@@ -472,7 +474,7 @@
     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 boolean mAlwaysUseTrafficStatsRateLimitCache;
 
     private final Object mOpenSessionCallsLock = new Object();
 
@@ -665,7 +667,8 @@
 
         final long cacheExpiryDurationMs = mDeps.getTrafficStatsRateLimitCacheExpiryDuration();
         final int cacheMaxEntries = mDeps.getTrafficStatsRateLimitCacheMaxEntries();
-        mSupportTrafficStatsRateLimitCache = mDeps.supportTrafficStatsRateLimitCache(mContext);
+        mAlwaysUseTrafficStatsRateLimitCache =
+                mDeps.alwaysUseTrafficStatsRateLimitCache(mContext);
         mTrafficStatsTotalCache = new TrafficStatsRateLimitCache(mClock,
                 cacheExpiryDurationMs, cacheMaxEntries);
         mTrafficStatsIfaceCache = new TrafficStatsRateLimitCache(mClock,
@@ -920,12 +923,12 @@
         }
 
         /**
-         * Get whether TrafficStats rate-limit cache is supported.
+         * Get whether TrafficStats rate-limit cache is always applied.
          *
          * 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) {
+        public boolean alwaysUseTrafficStatsRateLimitCache(@NonNull Context ctx) {
             return SdkLevel.isAtLeastV() && DeviceConfigUtils.isTetheringFeatureNotChickenedOut(
                     ctx, TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG);
         }
@@ -955,6 +958,13 @@
         }
 
         /**
+         * Wrapper method for {@link CompatChanges#isChangeEnabled(long, int)}
+         */
+        public boolean isChangeEnabled(final long changeId, final int uid) {
+            return CompatChanges.isChangeEnabled(changeId, uid);
+        }
+
+        /**
          * Retrieves native network total statistics.
          *
          * @return A NetworkStats.Entry containing the native statistics, or
@@ -2085,14 +2095,14 @@
         }
         if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
 
-        if (!mSupportTrafficStatsRateLimitCache) {
-            return getEntryValueForType(mDeps.nativeGetUidStat(uid), type);
+        if (mAlwaysUseTrafficStatsRateLimitCache
+                || mDeps.isChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, callingUid)) {
+            final NetworkStats.Entry entry = mTrafficStatsUidCache.getOrCompute(IFACE_ALL, uid,
+                    () -> mDeps.nativeGetUidStat(uid));
+            return getEntryValueForType(entry, type);
         }
 
-        final NetworkStats.Entry entry = mTrafficStatsUidCache.getOrCompute(IFACE_ALL, uid,
-                () -> mDeps.nativeGetUidStat(uid));
-
-        return getEntryValueForType(entry, type);
+        return getEntryValueForType(mDeps.nativeGetUidStat(uid), type);
     }
 
     @Nullable
@@ -2114,14 +2124,15 @@
         Objects.requireNonNull(iface);
         if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
 
-        if (!mSupportTrafficStatsRateLimitCache) {
-            return getEntryValueForType(getIfaceStatsInternal(iface), type);
+        if (mAlwaysUseTrafficStatsRateLimitCache
+                || mDeps.isChangeEnabled(
+                        ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, Binder.getCallingUid())) {
+            final NetworkStats.Entry entry = mTrafficStatsIfaceCache.getOrCompute(iface, UID_ALL,
+                    () -> getIfaceStatsInternal(iface));
+            return getEntryValueForType(entry, type);
         }
 
-        final NetworkStats.Entry entry = mTrafficStatsIfaceCache.getOrCompute(iface, UID_ALL,
-                () -> getIfaceStatsInternal(iface));
-
-        return getEntryValueForType(entry, type);
+        return getEntryValueForType(getIfaceStatsInternal(iface), type);
     }
 
     private long getEntryValueForType(@Nullable NetworkStats.Entry entry, int type) {
@@ -2167,14 +2178,15 @@
     @Override
     public long getTotalStats(int type) {
         if (!isEntryValueTypeValid(type)) return UNSUPPORTED;
-        if (!mSupportTrafficStatsRateLimitCache) {
-            return getEntryValueForType(getTotalStatsInternal(), type);
+        if (mAlwaysUseTrafficStatsRateLimitCache
+                || mDeps.isChangeEnabled(
+                        ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, Binder.getCallingUid())) {
+            final NetworkStats.Entry entry = mTrafficStatsTotalCache.getOrCompute(
+                    IFACE_ALL, UID_ALL, () -> getTotalStatsInternal());
+            return getEntryValueForType(entry, type);
         }
 
-        final NetworkStats.Entry entry = mTrafficStatsTotalCache.getOrCompute(IFACE_ALL, UID_ALL,
-                () -> getTotalStatsInternal());
-
-        return getEntryValueForType(entry, type);
+        return getEntryValueForType(getTotalStatsInternal(), type);
     }
 
     @Override
@@ -2938,7 +2950,7 @@
             } catch (IOException e) {
                 pw.println("(failed to dump FastDataInput counters)");
             }
-            pw.print("trafficstats.cache.supported", mSupportTrafficStatsRateLimitCache);
+            pw.print("trafficstats.cache.alwaysuse", mAlwaysUseTrafficStatsRateLimitCache);
             pw.println();
             pw.print(TRAFFIC_STATS_CACHE_EXPIRY_DURATION_NAME,
                     mDeps.getTrafficStatsRateLimitCacheExpiryDuration());
diff --git a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
index 6425daa..06a70d3 100644
--- a/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/unit/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -56,6 +56,7 @@
 import static android.net.TrafficStats.MB_IN_BYTES;
 import static android.net.TrafficStats.UID_REMOVED;
 import static android.net.TrafficStats.UID_TETHERING;
+import static android.net.connectivity.ConnectivityCompatChanges.ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE;
 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;
@@ -306,6 +307,7 @@
     private HandlerThread mObserverHandlerThread;
     final TestDependencies mDeps = new TestDependencies();
     final HashMap<String, Boolean> mFeatureFlags = new HashMap<>();
+    final HashMap<Long, Boolean> mCompatChanges = new HashMap<>();
 
     // This will set feature flags from @FeatureFlag annotations
     // into the map before setUp() runs.
@@ -594,7 +596,7 @@
         }
 
         @Override
-        public boolean supportTrafficStatsRateLimitCache(Context ctx) {
+        public boolean alwaysUseTrafficStatsRateLimitCache(Context ctx) {
             return mFeatureFlags.getOrDefault(TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG, false);
         }
 
@@ -608,6 +610,14 @@
             return DEFAULT_TRAFFIC_STATS_CACHE_MAX_ENTRIES;
         }
 
+        @Override
+        public boolean isChangeEnabled(long changeId, int uid) {
+            return mCompatChanges.getOrDefault(changeId, true);
+        }
+
+        public void setChangeEnabled(long changeId, boolean enabled) {
+            mCompatChanges.put(changeId, enabled);
+        }
         @Nullable
         @Override
         public NetworkStats.Entry nativeGetTotalStat() {
@@ -2413,17 +2423,33 @@
 
     @FeatureFlag(name = TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG, enabled = false)
     @Test
-    public void testTrafficStatsRateLimitCache_disabled() throws Exception {
-        doTestTrafficStatsRateLimitCache(false /* cacheEnabled */);
+    public void testTrafficStatsRateLimitCache_disabledWithCompatChangeEnabled() throws Exception {
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, true);
+        doTestTrafficStatsRateLimitCache(true /* expectCached */);
     }
 
     @FeatureFlag(name = TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG)
     @Test
-    public void testTrafficStatsRateLimitCache_enabled() throws Exception {
-        doTestTrafficStatsRateLimitCache(true /* cacheEnabled */);
+    public void testTrafficStatsRateLimitCache_enabledWithCompatChangeEnabled() throws Exception {
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, true);
+        doTestTrafficStatsRateLimitCache(true /* expectCached */);
     }
 
-    private void doTestTrafficStatsRateLimitCache(boolean cacheEnabled) throws Exception {
+    @FeatureFlag(name = TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG, enabled = false)
+    @Test
+    public void testTrafficStatsRateLimitCache_disabledWithCompatChangeDisabled() throws Exception {
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, false);
+        doTestTrafficStatsRateLimitCache(false /* expectCached */);
+    }
+
+    @FeatureFlag(name = TRAFFICSTATS_RATE_LIMIT_CACHE_ENABLED_FLAG)
+    @Test
+    public void testTrafficStatsRateLimitCache_enabledWithCompatChangeDisabled() throws Exception {
+        mDeps.setChangeEnabled(ENABLE_TRAFFICSTATS_RATE_LIMIT_CACHE, false);
+        doTestTrafficStatsRateLimitCache(true /* expectCached */);
+    }
+
+    private void doTestTrafficStatsRateLimitCache(boolean expectCached) 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();
@@ -2433,7 +2459,7 @@
         // Verify the values are cached.
         incrementCurrentTime(DEFAULT_TRAFFIC_STATS_CACHE_EXPIRY_DURATION_MS / 2);
         mockTrafficStatsValues(65L, 8L, 1055L, 9L);
-        if (cacheEnabled) {
+        if (expectCached) {
             assertTrafficStatsValues(TEST_IFACE, myUid, 64L, 3L, 1024L, 8L);
         } else {
             assertTrafficStatsValues(TEST_IFACE, myUid, 65L, 8L, 1055L, 9L);