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

Commit 66b8a14a authored by Jan Tomljanovic's avatar Jan Tomljanovic
Browse files

Implement multi rate limiter. 

Implement a class that can keep track of multiple quotas at the same time.

Test: atest MultiRateLimiterTest
Bug: 154198299
Change-Id: If2428f0a2c21cf91d250fec4ecf182550e00977f
parent 06f75983

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

0 → 100644
+183 −0
Original line number Diff line number Diff line 
/*

 * Copyright (C) 2020 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.utils.quota;



import android.annotation.NonNull;

import android.annotation.Nullable;

import android.content.Context;



import com.android.internal.annotations.GuardedBy;

import com.android.internal.annotations.VisibleForTesting;



import java.time.Duration;

import java.util.ArrayList;

import java.util.List;



/**

 * Can be used to rate limit events per app based on multiple rates at the same time. For example,

 * it can limit an event to happen only:

 *

 * <li>5 times in 20 seconds</li>

 * and

 * <li>6 times in 40 seconds</li>

 * and

 * <li>10 times in 1 hour</li>

 *

 * <p><br>

 * All listed rates apply at the same time, and the UPTC will be out of quota if it doesn't satisfy

 * all the given rates. The underlying mechanism used is

 * {@link com.android.server.utils.quota.CountQuotaTracker}, so all its conditions apply, as well

 * as an additional constraint: all the user-package-tag combinations (UPTC) are considered to be in

 * the same {@link com.android.server.utils.quota.Category}.

 * </p>

 *

 * @hide

 */

public class MultiRateLimiter {



    private static final CountQuotaTracker[] EMPTY_TRACKER_ARRAY = {};



    private final Object mLock = new Object();

    @GuardedBy("mLock")

    private final CountQuotaTracker[] mQuotaTrackers;



    private MultiRateLimiter(List<CountQuotaTracker> quotaTrackers) {

        mQuotaTrackers = quotaTrackers.toArray(EMPTY_TRACKER_ARRAY);

    }



    /** Record that an event happened and count it towards the given quota. */

    public void noteEvent(int userId, @NonNull String packageName, @Nullable String tag) {

        synchronized (mLock) {

            noteEventLocked(userId, packageName, tag);

        }

    }



    /** Check whether the given UPTC is allowed to trigger an event. */

    public boolean isWithinQuota(int userId, @NonNull String packageName, @Nullable String tag) {

        synchronized (mLock) {

            return isWithinQuotaLocked(userId, packageName, tag);

        }

    }



    @GuardedBy("mLock")

    private void noteEventLocked(int userId, @NonNull String packageName, @Nullable String tag) {

        for (CountQuotaTracker quotaTracker : mQuotaTrackers) {

            quotaTracker.noteEvent(userId, packageName, tag);

        }

    }



    @GuardedBy("mLock")

    private boolean isWithinQuotaLocked(int userId, @NonNull String packageName,

            @Nullable String tag) {

        for (CountQuotaTracker quotaTracker : mQuotaTrackers) {

            if (!quotaTracker.isWithinQuota(userId, packageName, tag)) {

                return false;

            }

        }

        return true;

    }



    /** Can create a new {@link MultiRateLimiter}. */

    public static class Builder {



        private final List<CountQuotaTracker> mQuotaTrackers;

        private final Context mContext;

        private final Categorizer mCategorizer;

        private final Category mCategory;

        @Nullable private final QuotaTracker.Injector mInjector;



        /**

         * Creates a new builder and allows to inject an object that can be used

         * to manipulate elapsed time in tests.

         */

        @VisibleForTesting

        Builder(Context context, QuotaTracker.Injector injector) {

            this.mQuotaTrackers = new ArrayList<>();

            this.mContext = context;

            this.mInjector = injector;

            this.mCategorizer = Categorizer.SINGLE_CATEGORIZER;

            this.mCategory = Category.SINGLE_CATEGORY;

        }



        /** Creates a new builder for {@link MultiRateLimiter}. */

        public Builder(Context context) {

            this(context, null);

        }



        /**

         * Adds another rate limit to be used in {@link MultiRateLimiter}.

         *

         * @param limit The maximum event count an app can have in the rolling time window.

         * @param windowSize The rolling time window to use when checking quota usage.

         */

        public Builder addRateLimit(int limit, Duration windowSize) {

            CountQuotaTracker countQuotaTracker;

            if (mInjector != null) {

                countQuotaTracker = new CountQuotaTracker(mContext, mCategorizer, mInjector);

            } else {

                countQuotaTracker = new CountQuotaTracker(mContext, mCategorizer);

            }

            countQuotaTracker.setCountLimit(mCategory, limit, windowSize.toMillis());

            mQuotaTrackers.add(countQuotaTracker);

            return this;

        }



        /** Adds another rate limit to be used in {@link MultiRateLimiter}. */

        public Builder addRateLimit(@NonNull RateLimit rateLimit) {

            return addRateLimit(rateLimit.mLimit, rateLimit.mWindowSize);

        }



        /** Adds all given rate limits that will be used in {@link MultiRateLimiter}. */

        public Builder addRateLimits(@NonNull RateLimit[] rateLimits) {

            for (RateLimit rateLimit : rateLimits) {

                addRateLimit(rateLimit);

            }

            return this;

        }



        /**

         * Return a new {@link com.android.server.utils.quota.MultiRateLimiter} using set rate

         * limit.

         */

        public MultiRateLimiter build() {

            return new MultiRateLimiter(mQuotaTrackers);

        }

    }



    /** Helper class that describes a rate limit. */

    public static class RateLimit {

        public final int mLimit;

        public final Duration mWindowSize;



        /**

         * @param limit The maximum count of some occurrence in the rolling time window.

         * @param windowSize The rolling time window to use when checking quota usage.

         */

        private RateLimit(int limit, Duration windowSize) {

            this.mLimit = limit;

            this.mWindowSize = windowSize;

        }



        /**

         * @param limit The maximum count of some occurrence in the rolling time window.

         * @param windowSize The rolling time window to use when checking quota usage.

         */

        public static RateLimit create(int limit, Duration windowSize) {

            return new RateLimit(limit, windowSize);

        }

    }

}

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

0 → 100644
+197 −0
Original line number Diff line number Diff line 
/*

 * Copyright (C) 2020 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.utils.quota;



import static com.google.common.truth.Truth.assertThat;



import android.testing.TestableContext;



import androidx.test.InstrumentationRegistry;

import androidx.test.filters.SmallTest;



import org.junit.Rule;

import org.junit.Test;



import java.time.Duration;



@SmallTest

public class MultiRateLimiterTest {



    private static final int USER_ID = 1;

    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 =

            new TestableContext(InstrumentationRegistry.getContext(), null);



    private final InjectorForTest mInjector = new InjectorForTest();



    private static class InjectorForTest extends QuotaTracker.Injector {

        Duration mElapsedTime = Duration.ZERO;



        @Override

        public long getElapsedRealtime() {

            return mElapsedTime.toMillis();

        }



        @Override

        public boolean isAlarmManagerReady() {

            return true;

        }

    }



    @Test

    public void testSingleRateLimit_belowLimit_isWithinQuota() {

        MultiRateLimiter multiRateLimiter =  new MultiRateLimiter.Builder(mContext, mInjector)

                .addRateLimit(3, Duration.ofSeconds(20))

                .build();



        // Three quick events are within quota.

        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.ofMillis(50);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(100);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

    }



    @Test

    public void testSingleRateLimit_aboveLimit_isNotWithinQuota() {

        MultiRateLimiter multiRateLimiter =  new MultiRateLimiter.Builder(mContext, mInjector)

                .addRateLimit(3, Duration.ofSeconds(20))

                .build();



        mInjector.mElapsedTime = Duration.ZERO;

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(50);

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(100);

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(150);

        // We hit the limit, 4th event in under 20 seconds is not within quota.

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isFalse();

    }



    @Test

    public void testSingleRateLimit_afterGoingAboveQuotaAndWaitingWindow_isBackWithinQuota() {

        MultiRateLimiter multiRateLimiter =  new MultiRateLimiter.Builder(mContext, mInjector)

                .addRateLimit(3, Duration.ofSeconds(20))

                .build();



        mInjector.mElapsedTime = Duration.ZERO;

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(50);

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(100);

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(150);

        // We hit the limit, 4th event in under 20 seconds is not within quota.

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isFalse();



        mInjector.mElapsedTime = Duration.ofSeconds(21);

        // 20 seconds have passed, we're again within quota.

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

    }



    @Test

    public void createMultipleRateLimits_testTheyLimitsAsExpected() {

        MultiRateLimiter multiRateLimiter = new MultiRateLimiter.Builder(mContext, mInjector)

                .addRateLimit(3, Duration.ofSeconds(20)) // 1st limit

                .addRateLimit(4, Duration.ofSeconds(40)) // 2nd limit

                .addRateLimit(5, Duration.ofSeconds(60)) // 3rd limit

                .build();



        // Testing the 1st limit

        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.ofMillis(50);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(100);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(150);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isFalse();



        mInjector.mElapsedTime = Duration.ofSeconds(21);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        // Testing the 2nd limit

        mInjector.mElapsedTime = Duration.ofSeconds(35);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isFalse();



        mInjector.mElapsedTime = Duration.ofSeconds(42);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        // Testing the 3rd limit.

        mInjector.mElapsedTime = Duration.ofSeconds(43);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isFalse();



        mInjector.mElapsedTime = Duration.ofSeconds(62);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);

    }



    @Test

    public void createSingleRateLimit_testItLimitsOnlyGivenUptc() {

        MultiRateLimiter multiRateLimiter =  new MultiRateLimiter.Builder(mContext, mInjector)

                .addRateLimit(3, Duration.ofSeconds(20))

                .build();



        mInjector.mElapsedTime = Duration.ZERO;

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_2, TAG)).isTrue();

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(50);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_2, TAG)).isTrue();

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(100);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_2, TAG)).isTrue();

        multiRateLimiter.noteEvent(USER_ID, PACKAGE_NAME_1, TAG);



        mInjector.mElapsedTime = Duration.ofMillis(150);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isFalse();

        // Different userId - packageName - tag combination is still allowed.

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_2, TAG)).isTrue();



        mInjector.mElapsedTime = Duration.ofSeconds(21);

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_1, TAG)).isTrue();

        assertThat(multiRateLimiter.isWithinQuota(USER_ID, PACKAGE_NAME_2, TAG)).isTrue();

    }

}