Donate to e Foundation | Murena handsets with /e/OS | Own a part of Murena! Learn more

Commit cc2aedf5 authored by Jan Tomljanovic's avatar Jan Tomljanovic
Browse files

Add a method to reset quota for the given app. 

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

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

+30 −2
Original line number Diff line number Diff line
@@ -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

+2 −2
Original line number Diff line number Diff line
@@ -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

+93 −1
Original line number Diff line number Diff line
@@ -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();

    }

}