-
Notifications
You must be signed in to change notification settings - Fork 27
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #658 from brizental/1712920-rate-limit
Bug 1712920 - Implement rate limiting in ping upload
- Loading branch information
Showing
5 changed files
with
263 additions
and
9 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,125 @@ | ||
| /* This Source Code Form is subject to the terms of the Mozilla Public | ||
| * License, v. 2.0. If a copy of the MPL was not distributed with this | ||
| * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ | ||
|
|
||
| import { isUndefined, getMonotonicNow } from "../utils.js"; | ||
|
|
||
| /** | ||
| * An enum to represent the current state of the RateLimiter. | ||
| */ | ||
| export const enum RateLimiterState { | ||
| // The RateLimiter has not reached the maximum count and is still incrementing. | ||
| Incrementing, | ||
| // The RateLimiter has not reached the maximum count, but it is also not incrementing. | ||
| Stopped, | ||
| // The RateLimiter has reached the maximum count for the current interval. | ||
| Throttled, | ||
| } | ||
|
|
||
| class RateLimiter { | ||
| // Whether or not the RateLimiter is not counting any further for the current interval. | ||
| // This is different from the RateLimiter being throttled, because it may happen | ||
| // even if max count for the current interval has not been reached. | ||
| private stopped = false; | ||
|
|
||
| constructor( | ||
| // The duration of each interval, in millisecods. | ||
| private interval: number, | ||
| // The maximum count per interval. | ||
| private maxCount: number, | ||
| // The count for the current interval. | ||
| private count: number = 0, | ||
| // The instant the current interval has started, in milliseconds. | ||
| private started?: number, | ||
| ) {} | ||
|
|
||
| get elapsed(): number { | ||
| if (isUndefined(this.started)) { | ||
| return NaN; | ||
| } | ||
|
|
||
| const now = getMonotonicNow(); | ||
| const elapsed = now - this.started; | ||
|
|
||
| // It's very unlikely elapsed will be a negative number since we are using a monotonic timer | ||
| // here, but just to be extra sure, we account for it. | ||
| if (elapsed < 0) { | ||
| return NaN; | ||
| } | ||
|
|
||
| return elapsed; | ||
| } | ||
|
|
||
| private reset(): void { | ||
| this.started = getMonotonicNow(); | ||
| this.count = 0; | ||
| this.stopped = false; | ||
| } | ||
|
|
||
| /** | ||
| * The rate limiter should reset if | ||
| * | ||
| * 1. It has never started i.e. `started` is still `undefined`; | ||
| * 2. It has been started more than the interval time ago; | ||
| * 3. Something goes wrong while trying to calculate the elapsed time since the last reset. | ||
| * | ||
| * @returns Whether or not this rate limiter should reset. | ||
| */ | ||
| private shouldReset(): boolean { | ||
| if (isUndefined(this.started)) { | ||
| return true; | ||
| } | ||
|
|
||
| if (isNaN(this.elapsed) || this.elapsed > this.interval) { | ||
| return true; | ||
| } | ||
|
|
||
| return false; | ||
| } | ||
|
|
||
| /** | ||
| * Tries to increment the internal counter. | ||
| * | ||
| * @returns The current state of the RateLimiter plus the remaining time | ||
| * (in milliseconds) until the end of the current window. | ||
| */ | ||
| getState(): { | ||
| state: RateLimiterState, | ||
| remainingTime?: number, | ||
| } { | ||
| if (this.shouldReset()) { | ||
| this.reset(); | ||
| } | ||
|
|
||
| const remainingTime = this.interval - this.elapsed; | ||
| if (this.stopped) { | ||
| return { | ||
| state: RateLimiterState.Stopped, | ||
| remainingTime, | ||
| }; | ||
| } | ||
|
|
||
| if (this.count >= this.maxCount) { | ||
| return { | ||
| state: RateLimiterState.Throttled, | ||
| remainingTime, | ||
| }; | ||
| } | ||
|
|
||
| this.count++; | ||
| return { | ||
| state: RateLimiterState.Incrementing | ||
| }; | ||
| } | ||
|
|
||
| /** | ||
| * Stops counting for the current interval, regardless of the max count being reached. | ||
| * | ||
| * The RateLimiter will still be reset when time interval is over. | ||
| */ | ||
| stop(): void { | ||
| this.stopped = true; | ||
| } | ||
| } | ||
|
|
||
| export default RateLimiter; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,91 @@ | ||
| /* This Source Code Form is subject to the terms of the Mozilla Public | ||
| * License, v. 2.0. If a copy of the MPL was not distributed with this | ||
| * file, You can obtain one at http://mozilla.org/MPL/2.0/. */ | ||
|
|
||
| import assert from "assert"; | ||
| import type { SinonFakeTimers } from "sinon"; | ||
| import sinon from "sinon"; | ||
|
|
||
| import RateLimiter, { RateLimiterState } from "../../../../src/core/upload/rate_limiter"; | ||
|
|
||
| const sandbox = sinon.createSandbox(); | ||
| const now = new Date(); | ||
|
|
||
|
|
||
| describe("RateLimiter", function() { | ||
| let clock: SinonFakeTimers; | ||
|
|
||
| beforeEach(function() { | ||
| clock = sandbox.useFakeTimers(now.getTime()); | ||
| }); | ||
|
|
||
| afterEach(function () { | ||
| clock.restore(); | ||
| }); | ||
|
|
||
| it("rate limiter correctly resets in case elapsed time return an error", function () { | ||
| const rateLimiter = new RateLimiter( | ||
| 1000, /* interval */ | ||
| 3, /* maxCount */ | ||
| ); | ||
|
|
||
| // Reach the count for the current interval. | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
|
|
||
| sinon.replaceGetter(rateLimiter, "elapsed", () => NaN); | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
| }); | ||
|
|
||
| it("rate limiter correctly resets in case interval is over", function () { | ||
| const rateLimiter = new RateLimiter( | ||
| 1000, /* interval */ | ||
| 3, /* maxCount */ | ||
| ); | ||
|
|
||
| // Reach the count for the current interval. | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
|
|
||
| // Fake the time passing over the current interval | ||
| sinon.replaceGetter(rateLimiter, "elapsed", () => 1000 * 2); | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
| }); | ||
|
|
||
| it("rate limiter returns throttled state when it is throttled", function () { | ||
| const rateLimiter = new RateLimiter( | ||
| 1000, /* interval */ | ||
| 3, /* maxCount */ | ||
| ); | ||
|
|
||
| // Reach the count for the current interval. | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
|
|
||
| // Try one more time and we should be throttled. | ||
| const nextState = rateLimiter.getState(); | ||
| assert.strictEqual(nextState.state, RateLimiterState.Throttled); | ||
| assert.ok(nextState.remainingTime as number <= 1000 && nextState.remainingTime as number > 0); | ||
| }); | ||
|
|
||
| it("rate limiter returns stopped state when it is stopped", function () { | ||
| const rateLimiter = new RateLimiter( | ||
| 1000, /* interval */ | ||
| 3, /* maxCount */ | ||
| ); | ||
|
|
||
| // Don't reach the count for the current interval. | ||
| assert.deepStrictEqual(rateLimiter.getState(), { state: RateLimiterState.Incrementing }); | ||
|
|
||
| // Stop the rate limiter | ||
| rateLimiter.stop(); | ||
|
|
||
| // Try one more time and we should be stopped. | ||
| const nextState = rateLimiter.getState(); | ||
| assert.strictEqual(nextState.state, RateLimiterState.Stopped); | ||
| assert.ok(nextState.remainingTime as number <= 1000 && nextState.remainingTime as number > 0); | ||
| }); | ||
| }); |