Add a method to reset quota for the given app.

Test: atest MultiRateLimterTest
Bug: 174841445
Change-Id: I6a08e36d8e18bb504a0597fe6134fe5a9103a5be

services/core/java/com/android/server/utils/quota/MultiRateLimiter.java

@@ -19,6 +19,8 @@ package com.android.server.utils.quota;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
+import android.content.pm.PackageManager;
+import android.util.Slog;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
@@ -48,15 +50,18 @@ import java.util.List;
  * @hide
  */
 public class MultiRateLimiter {
+    private static final String TAG = "MultiRateLimiter";
 
     private static final CountQuotaTracker[] EMPTY_TRACKER_ARRAY = {};
 
     private final Object mLock = new Object();
     @GuardedBy("mLock")
     private final CountQuotaTracker[] mQuotaTrackers;
+    private final PackageManager mPackageManager;
 
-    private MultiRateLimiter(List<CountQuotaTracker> quotaTrackers) {
+    private MultiRateLimiter(List<CountQuotaTracker> quotaTrackers, PackageManager packageManager) {
         mQuotaTrackers = quotaTrackers.toArray(EMPTY_TRACKER_ARRAY);
+        mPackageManager = packageManager;
     }
 
     /** Record that an event happened and count it towards the given quota. */
@@ -73,6 +78,13 @@ public class MultiRateLimiter {
         }
     }
 
+    /** Remove all saved events from the rate limiter for the given app (reset it). */
+    public void clear(int userId, @NonNull String packageName) {
+        synchronized (mLock) {
+            clearLocked(userId, packageName);
+        }
+    }
+
     @GuardedBy("mLock")
     private void noteEventLocked(int userId, @NonNull String packageName, @Nullable String tag) {
         for (CountQuotaTracker quotaTracker : mQuotaTrackers) {
@@ -91,6 +103,22 @@ public class MultiRateLimiter {
         return true;
     }
 
+    @GuardedBy("mLock")
+    private void clearLocked(int userId, @NonNull String packageName) {
+        try {
+            int uid = mPackageManager.getApplicationInfoAsUser(packageName, 0, userId).uid;
+            for (CountQuotaTracker quotaTracker : mQuotaTrackers) {
+                // This method behaves as if the package has been removed from the device, which
+                // isn't the case here, but it does similar clean-up to what we are aiming for here,
+                // so it works for this use case.
+                quotaTracker.onAppRemovedLocked(packageName, uid);
+            }
+        } catch (PackageManager.NameNotFoundException e) {
+            Slog.e(TAG, "clear(userId, packageName) called with unrecognized arguments, no "
+                    + "action taken");
+        }
+    }
+
     /** Can create a new {@link MultiRateLimiter}. */
     public static class Builder {
 
@@ -154,7 +182,7 @@ public class MultiRateLimiter {
          * limit.
          */
         public MultiRateLimiter build() {
-            return new MultiRateLimiter(mQuotaTrackers);
+            return new MultiRateLimiter(mQuotaTrackers, mContext.getPackageManager());
         }
     }
 

services/core/java/com/android/server/utils/quota/QuotaTracker.java

@@ -58,7 +58,7 @@ import java.util.PriorityQueue;
  * of quota until it is below that limit again. Limits are applied according to the category
  * the UPTC is placed in. Categories are basic constructs to apply different limits to
  * different groups of UPTCs. For example, standby buckets can be a set of categories, or
- * foreground & background could be two categories. If every UPTC should have the limits
+ * foreground & background could be two categories. If every UPTC should have the same limits
  * applied, then only one category is needed.
  *
  * Note: all limits are enforced per category unless explicitly stated otherwise.
@@ -361,7 +361,7 @@ abstract class QuotaTracker {
     abstract void handleRemovedAppLocked(String packageName, int uid);
 
     @GuardedBy("mLock")
-    private void onAppRemovedLocked(String packageName, int uid) {
+    void onAppRemovedLocked(String packageName, int uid) {
         if (packageName == null) {
             Slog.wtf(TAG, "Told app removed but given null package name.");
             return;

services/tests/mockingservicestests/src/com/android/server/utils/quota/MultiRateLimiterTest.java

@@ -16,15 +16,29 @@
 
 package com.android.server.utils.quota;
 
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
+
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.Mockito.when;
+
+import android.content.ContextWrapper;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
+import android.os.UserHandle;
 import android.testing.TestableContext;
 
 import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
 
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoSession;
+import org.mockito.quality.Strictness;
 
 import java.time.Duration;
 
@@ -32,15 +46,20 @@ import java.time.Duration;
 public class MultiRateLimiterTest {
 
     private static final int USER_ID = 1;
+    private static final int UID_1 = 10_001;
     private static final String PACKAGE_NAME_1 = "com.android.package.one";
     private static final String PACKAGE_NAME_2 = "com.android.package.two";
     private static final String TAG = "tag";
 
     @Rule
-    public final TestableContext mContext =
+    public final TestableContext mTestableContext =
             new TestableContext(InstrumentationRegistry.getContext(), null);
 
     private final InjectorForTest mInjector = new InjectorForTest();
+    private MockitoSession mMockingSession;
+    private ContextWrapper mContext;
+
+    @Mock private PackageManager mPackageManager;
 
     private static class InjectorForTest extends QuotaTracker.Injector {
         Duration mElapsedTime = Duration.ZERO;
@@ -56,6 +75,35 @@ public class MultiRateLimiterTest {
         }
     }
 
+    @Before
+    public void setup() throws Exception {
+        mMockingSession = mockitoSession()
+                .initMocks(this)
+                .strictness(Strictness.LENIENT)
+                .mockStatic(UserHandle.class)
+                .startMocking();
+        doReturn(USER_ID).when(() -> UserHandle.getUserId(UID_1));
+
+        ApplicationInfo applicationInfo = new ApplicationInfo();
+        applicationInfo.uid = UID_1;
+        when(mPackageManager.getApplicationInfoAsUser(PACKAGE_NAME_1, 0, USER_ID))
+                .thenReturn(applicationInfo);
+
+        mContext = new ContextWrapper(mTestableContext) {
+            @Override
+            public PackageManager getPackageManager() {
+                return mPackageManager;
+            }
+        };
+    }
+
+    @After
+    public void tearDown() {
+        if (mMockingSession != null) {
+            mMockingSession.finishMocking();
+        }
+    }
+
     @Test
     public void testSingleRateLimit_belowLimit_isWithinQuota() {
         MultiRateLimiter multiRateLimiter =  new MultiRateLimiter.Builder(mContext, mInjector)
@@ -194,4 +242,48 @@ public class MultiRateLimiterTest {
         assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();
         assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_2, TAG)).isTrue();
     }
+
+    @Test
+    public void clearRateLimiterForPackage_afterReachingQuota_quotaIsReset() {
+        MultiRateLimiter multiRateLimiter =  new MultiRateLimiter.Builder(mContext, mInjector)
+                .addRateLimit(1, Duration.ofSeconds(100))
+                .build();
+
+        mInjector.mElapsedTime = Duration.ZERO;
+        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();
+        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);
+
+        mInjector.mElapsedTime = Duration.ofSeconds(1);
+        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isFalse();
+
+        multiRateLimiter.clear(USER_ID, PACKAGE_NAME_1);
+
+        // Quota for that package is reset.
+        mInjector.mElapsedTime = Duration.ofSeconds(1);
+        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();
+
+        // Quota is enforced again.
+        mInjector.mElapsedTime = Duration.ofSeconds(1);
+        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);
+        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isFalse();
+    }
+
+    @Test
+    public void clearRateLimiterForPackage_doesntAffectOtherPackages() {
+        MultiRateLimiter multiRateLimiter =  new MultiRateLimiter.Builder(mContext, mInjector)
+                .addRateLimit(1, Duration.ofSeconds(100))
+                .build();
+
+        mInjector.mElapsedTime = Duration.ZERO;
+        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_2, TAG)).isTrue();
+        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_2, TAG);
+
+        mInjector.mElapsedTime = Duration.ofSeconds(1);
+        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_2, TAG)).isFalse();
+
+        multiRateLimiter.clear(USER_ID, PACKAGE_NAME_1);
+
+        // Doesn't affect the other package.
+        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_2, TAG)).isFalse();
+    }
 }