diff --git a/service-t/src/com/android/server/net/TrafficStatsRateLimitCache.java b/service-t/src/com/android/server/net/TrafficStatsRateLimitCache.java new file mode 100644 index 0000000000..8598ac4271 --- /dev/null +++ b/service-t/src/com/android/server/net/TrafficStatsRateLimitCache.java @@ -0,0 +1,134 @@ +/* + * 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.server.net; + +import android.annotation.NonNull; +import android.annotation.Nullable; +import android.net.NetworkStats; + +import com.android.internal.annotations.GuardedBy; + +import java.time.Clock; +import java.util.HashMap; +import java.util.Objects; + +/** + * A thread-safe cache for storing and retrieving {@link NetworkStats.Entry} objects, + * with an adjustable expiry duration to manage data freshness. + */ +class TrafficStatsRateLimitCache { + private final Clock mClock; + private final long mExpiryDurationMs; + + /** + * Constructs a new {@link TrafficStatsRateLimitCache} with the specified expiry duration. + * + * @param clock The {@link Clock} to use for determining timestamps. + * @param expiryDurationMs The expiry duration in milliseconds. + */ + TrafficStatsRateLimitCache(@NonNull Clock clock, long expiryDurationMs) { + mClock = clock; + mExpiryDurationMs = expiryDurationMs; + } + + private static class TrafficStatsCacheKey { + @Nullable + public final String iface; + public final int uid; + + TrafficStatsCacheKey(@Nullable String iface, int uid) { + this.iface = iface; + this.uid = uid; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (!(o instanceof TrafficStatsCacheKey)) return false; + TrafficStatsCacheKey that = (TrafficStatsCacheKey) o; + return uid == that.uid && Objects.equals(iface, that.iface); + } + + @Override + public int hashCode() { + return Objects.hash(iface, uid); + } + } + + private static class TrafficStatsCacheValue { + public final long timestamp; + @NonNull + public final NetworkStats.Entry entry; + + TrafficStatsCacheValue(long timestamp, NetworkStats.Entry entry) { + this.timestamp = timestamp; + this.entry = entry; + } + } + + @GuardedBy("mMap") + private final HashMap mMap = new HashMap<>(); + + /** + * Retrieves a {@link NetworkStats.Entry} from the cache, associated with the given key. + * + * @param iface The interface name to include in the cache key. Null if not applicable. + * @param uid The UID to include in the cache key. {@code UID_ALL} if not applicable. + * @return The cached {@link NetworkStats.Entry}, or null if not found or expired. + */ + @Nullable + NetworkStats.Entry get(String iface, int uid) { + final TrafficStatsCacheKey key = new TrafficStatsCacheKey(iface, uid); + synchronized (mMap) { // Synchronize for thread-safety + final TrafficStatsCacheValue value = mMap.get(key); + if (value != null && !isExpired(value.timestamp)) { + return value.entry; + } else { + mMap.remove(key); // Remove expired entries + return null; + } + } + } + + /** + * Stores a {@link NetworkStats.Entry} in the cache, associated with the given key. + * + * @param iface The interface name to include in the cache key. Null if not applicable. + * @param uid The UID to include in the cache key. {@code UID_ALL} if not applicable. + * @param entry The {@link NetworkStats.Entry} to store in the cache. + */ + void put(String iface, int uid, @NonNull final NetworkStats.Entry entry) { + Objects.requireNonNull(entry); + final TrafficStatsCacheKey key = new TrafficStatsCacheKey(iface, uid); + synchronized (mMap) { // Synchronize for thread-safety + mMap.put(key, new TrafficStatsCacheValue(mClock.millis(), entry)); + } + } + + /** + * Clear the cache. + */ + void clear() { + synchronized (mMap) { + mMap.clear(); + } + } + + private boolean isExpired(long timestamp) { + return mClock.millis() > timestamp + mExpiryDurationMs; + } +} diff --git a/tests/unit/java/com/android/server/net/TrafficStatsRateLimitCacheTest.kt b/tests/unit/java/com/android/server/net/TrafficStatsRateLimitCacheTest.kt new file mode 100644 index 0000000000..27e6f96f9e --- /dev/null +++ b/tests/unit/java/com/android/server/net/TrafficStatsRateLimitCacheTest.kt @@ -0,0 +1,89 @@ +/* + * 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.server.net + +import android.net.NetworkStats +import com.android.testutils.DevSdkIgnoreRunner +import java.time.Clock +import kotlin.test.assertEquals +import kotlin.test.assertNull +import org.junit.Test +import org.junit.runner.RunWith +import org.mockito.Mockito.mock +import org.mockito.Mockito.`when` + +@RunWith(DevSdkIgnoreRunner::class) +class TrafficStatsRateLimitCacheTest { + companion object { + private const val expiryDurationMs = 1000L + } + + private val clock = mock(Clock::class.java) + private val entry = mock(NetworkStats.Entry::class.java) + private val cache = TrafficStatsRateLimitCache(clock, expiryDurationMs) + + @Test + fun testGet_returnsEntryIfNotExpired() { + cache.put("iface", 2, entry) + `when`(clock.millis()).thenReturn(500L) // Set clock to before expiry + val result = cache.get("iface", 2) + assertEquals(entry, result) + } + + @Test + fun testGet_returnsNullIfExpired() { + cache.put("iface", 2, entry) + `when`(clock.millis()).thenReturn(2000L) // Set clock to after expiry + assertNull(cache.get("iface", 2)) + } + + @Test + fun testGet_returnsNullForNonExistentKey() { + val result = cache.get("otherIface", 99) + assertNull(result) + } + + @Test + fun testPutAndGet_retrievesCorrectEntryForDifferentKeys() { + val entry1 = mock(NetworkStats.Entry::class.java) + val entry2 = mock(NetworkStats.Entry::class.java) + + cache.put("iface1", 2, entry1) + cache.put("iface2", 4, entry2) + + assertEquals(entry1, cache.get("iface1", 2)) + assertEquals(entry2, cache.get("iface2", 4)) + } + + @Test + fun testPut_overridesExistingEntry() { + val entry1 = mock(NetworkStats.Entry::class.java) + val entry2 = mock(NetworkStats.Entry::class.java) + + cache.put("iface", 2, entry1) + cache.put("iface", 2, entry2) // Put with the same key + + assertEquals(entry2, cache.get("iface", 2)) + } + + @Test + fun testClear() { + cache.put("iface", 2, entry) + cache.clear() + assertNull(cache.get("iface", 2)) + } +}