From 8c1430a0ff0a4fa1b3ea5d204c4b61cc02cfab54 Mon Sep 17 00:00:00 2001 From: Brian Chen <chenbrian@google.com> Date: Wed, 15 Apr 2020 11:19:42 -0700 Subject: [PATCH 1/4] add rate limiting --- dev/src/rate-limiter.ts | 145 +++++++++++++++++++++++++++++++++++++++ dev/test/rate-limiter.ts | 94 +++++++++++++++++++++++++ 2 files changed, 239 insertions(+) create mode 100644 dev/src/rate-limiter.ts create mode 100644 dev/test/rate-limiter.ts diff --git a/dev/src/rate-limiter.ts b/dev/src/rate-limiter.ts new file mode 100644 index 000000000..b949ee4b9 --- /dev/null +++ b/dev/src/rate-limiter.ts @@ -0,0 +1,145 @@ +/*! + * Copyright 2020 Google LLC + * + * 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. + */ +import * as assert from 'assert'; + +import {Timestamp} from './timestamp'; + +/** + * A helper that uses the Token Bucket algorithm to rate limit the number of + * operations that can be made in a second. + * + * Before a given request containing a number of operations can proceed, + * RateLimiter determines doing so stays under the provided rate limits. It can + * also determine how much time is required before a request can be made. + * + * RateLimiter can also implement a gradually increasing rate limit. This is + * used to enforce the 500/50/5 rule + * (https://cloud.google.com/datastore/docs/best-practices#ramping_up_traffic). + * + * @private + */ +export class RateLimiter { + // Number of tokens available. Each operation consumes one token. + availableTokens: number; + + // When the token bucket was last refilled. + lastRefillTime = Timestamp.now(); + + /** + * @param initialCapacity Initial maximum number of operations per second. + * @param multiplier Rate by which to increase the capacity. + * @param multiplierMillis How often the capacity should increase in + * milliseconds. + * @param startTime Used for testing the limiter. + */ + constructor( + private readonly initialCapacity: number, + private readonly multiplier: number, + private readonly multiplierMillis: number, + private readonly startTime = Timestamp.now() + ) { + this.availableTokens = initialCapacity; + this.lastRefillTime = startTime; + } + + /** + * Tries to make the number of operations. Returns true if the request succeeded + * and false otherwise. + * + * @param currentTime Used for testing the limiter. + * @private + */ + tryMakeRequest( + numOperations: number, + currentTime = Timestamp.now() + ): boolean { + this.refillTokens(currentTime); + if (numOperations <= this.availableTokens) { + this.availableTokens -= numOperations; + return true; + } + return false; + } + + /** + * Returns the number of ms needed to make a request with the provided number + * of operations. Returns 0 if the request can be made with the existing + * capacity. Returns -1 if the request is not possible with the current + * capacity. + * + * @param currentTime Used for testing the limiter. + * @private + */ + getNextRequestDelayMs( + numOperations: number, + currentTime = Timestamp.now() + ): number { + if (numOperations < this.availableTokens) { + return 0; + } + + const capacity = this.calculateCapacity(currentTime, this.startTime); + if (capacity < numOperations) { + return -1; + } + + const requiredTokens = numOperations - this.availableTokens; + return Math.ceil((requiredTokens * 1000) / capacity); + } + + /** + * Refills the number of available tokens based on how much time has elapsed + * since the last time the tokens were refilled. + * + * @private + */ + private refillTokens(currentTime = Timestamp.now()): void { + if (currentTime.toMillis() > this.lastRefillTime.toMillis()) { + const elapsedTime = + currentTime.toMillis() - this.lastRefillTime.toMillis(); + const capacity = this.calculateCapacity(currentTime, this.startTime); + const tokensToAdd = Math.floor((elapsedTime * capacity) / 1000); + if (tokensToAdd > 0) { + this.availableTokens = Math.min( + capacity, + this.availableTokens + tokensToAdd + ); + this.lastRefillTime = currentTime; + } + } + } + + /** + * Calculates the maximum capacity based on the two provided timestamps. + * + * @private + */ + // Visible for testing. + calculateCapacity(currentTime: Timestamp, startTime: Timestamp): number { + assert( + currentTime.valueOf() >= startTime.valueOf(), + 'startTime cannot be after currentTime' + ); + const millisElapsed = currentTime.toMillis() - startTime.toMillis(); + const operationsPerSecond = Math.floor( + Math.pow( + this.multiplier, + Math.floor(millisElapsed / this.multiplierMillis) + ) * this.initialCapacity + ); + return operationsPerSecond; + } +} diff --git a/dev/test/rate-limiter.ts b/dev/test/rate-limiter.ts new file mode 100644 index 000000000..ff7693090 --- /dev/null +++ b/dev/test/rate-limiter.ts @@ -0,0 +1,94 @@ +// Copyright 2020 Google LLC +// +// 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. + +import {expect} from 'chai'; + +import {Timestamp} from '../src'; +import {RateLimiter} from '../src/rate-limiter'; + +describe('RateLimiter', () => { + let limiter: RateLimiter; + + beforeEach(() => { + limiter = new RateLimiter(500, 1.5, 5 * 60 * 1000, new Timestamp(0, 0)); + }); + + it('accepts and rejects requests based on capacity', () => { + expect(limiter.tryMakeRequest(250, new Timestamp(0, 0))).to.be.true; + expect(limiter.tryMakeRequest(250, new Timestamp(0, 0))).to.be.true; + + // Once tokens have been used, further requests should fail. + expect(limiter.tryMakeRequest(1, new Timestamp(0, 0))).to.be.false; + + // Tokens will only refill up to max capacity. + expect(limiter.tryMakeRequest(501, new Timestamp(1, 0))).to.be.false; + expect(limiter.tryMakeRequest(500, new Timestamp(1, 0))).to.be.true; + + // Tokens will refill incrementally based on the number of ms elapsed. + expect(limiter.tryMakeRequest(251, new Timestamp(1, 500 * 1e6))).to.be + .false; + expect(limiter.tryMakeRequest(250, new Timestamp(1, 500 * 1e6))).to.be.true; + + // Scales with multiplier. + expect(limiter.tryMakeRequest(751, new Timestamp(5 * 60 - 1, 0))).to.be + .false; + expect(limiter.tryMakeRequest(751, new Timestamp(5 * 60, 0))).to.be.false; + expect(limiter.tryMakeRequest(750, new Timestamp(5 * 60, 0))).to.be.true; + }); + + it('calculates the number of ms needed to place the next request', () => { + // Should return 0 if there are enough tokens for the request to be made. + let timestamp = new Timestamp(0, 0); + expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(0); + + // Should factor in remaining tokens when calculating the time. + expect(limiter.tryMakeRequest(250, timestamp)); + expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(500); + + // Once tokens have been used, should calculate time before next request. + timestamp = new Timestamp(1, 0); + expect(limiter.tryMakeRequest(500, timestamp)).to.be.true; + expect(limiter.getNextRequestDelayMs(100, timestamp)).to.equal(200); + expect(limiter.getNextRequestDelayMs(250, timestamp)).to.equal(500); + expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(1000); + expect(limiter.getNextRequestDelayMs(501, timestamp)).to.equal(-1); + + // Scales with multiplier. + timestamp = new Timestamp(5 * 60, 0); + expect(limiter.tryMakeRequest(750, timestamp)).to.be.true; + expect(limiter.getNextRequestDelayMs(250, timestamp)).to.equal(334); + expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(667); + expect(limiter.getNextRequestDelayMs(750, timestamp)).to.equal(1000); + expect(limiter.getNextRequestDelayMs(751, timestamp)).to.equal(-1); + }); + + it('calculates the maximum number of operations correctly', async () => { + const currentTime = new Timestamp(0, 0); + expect( + limiter.calculateCapacity(new Timestamp(1, 0), currentTime) + ).to.equal(500); + expect( + limiter.calculateCapacity(new Timestamp(5 * 60, 0), currentTime) + ).to.equal(750); + expect( + limiter.calculateCapacity(new Timestamp(10 * 60, 0), currentTime) + ).to.equal(1125); + expect( + limiter.calculateCapacity(new Timestamp(15 * 60, 0), currentTime) + ).to.equal(1687); + expect( + limiter.calculateCapacity(new Timestamp(90 * 60, 0), currentTime) + ).to.equal(738945); + }); +}); From 554172a5cc7726820ed44e68ccfb2b2c2c372387 Mon Sep 17 00:00:00 2001 From: Brian Chen <chenbrian@google.com> Date: Wed, 15 Apr 2020 11:31:34 -0700 Subject: [PATCH 2/4] add internal methods --- dev/src/rate-limiter.ts | 30 +++++++++++++++++++------ dev/test/rate-limiter.ts | 47 ++++++++++++++++++++-------------------- 2 files changed, 47 insertions(+), 30 deletions(-) diff --git a/dev/src/rate-limiter.ts b/dev/src/rate-limiter.ts index b949ee4b9..d2473a18e 100644 --- a/dev/src/rate-limiter.ts +++ b/dev/src/rate-limiter.ts @@ -62,10 +62,17 @@ export class RateLimiter { * @param currentTime Used for testing the limiter. * @private */ - tryMakeRequest( - numOperations: number, - currentTime = Timestamp.now() - ): boolean { + tryMakeRequest(numOperations: number) { + return this._tryMakeRequest(numOperations, Timestamp.now()); + } + + /** + * Used in testing for different timestamps. + * + * @private + */ + // Visible for testing. + _tryMakeRequest(numOperations: number, currentTime: Timestamp): boolean { this.refillTokens(currentTime); if (numOperations <= this.availableTokens) { this.availableTokens -= numOperations; @@ -80,12 +87,21 @@ export class RateLimiter { * capacity. Returns -1 if the request is not possible with the current * capacity. * - * @param currentTime Used for testing the limiter. * @private */ - getNextRequestDelayMs( + getNextRequestDelayMs(numOperations: number) { + return this._getNextRequestDelayMs(numOperations, Timestamp.now()); + } + + /** + * Used in testing for different timestamps. + * + * @private + */ + // Visible for testing. + _getNextRequestDelayMs( numOperations: number, - currentTime = Timestamp.now() + currentTime: Timestamp ): number { if (numOperations < this.availableTokens) { return 0; diff --git a/dev/test/rate-limiter.ts b/dev/test/rate-limiter.ts index ff7693090..023404d8b 100644 --- a/dev/test/rate-limiter.ts +++ b/dev/test/rate-limiter.ts @@ -25,52 +25,53 @@ describe('RateLimiter', () => { }); it('accepts and rejects requests based on capacity', () => { - expect(limiter.tryMakeRequest(250, new Timestamp(0, 0))).to.be.true; - expect(limiter.tryMakeRequest(250, new Timestamp(0, 0))).to.be.true; + expect(limiter._tryMakeRequest(250, new Timestamp(0, 0))).to.be.true; + expect(limiter._tryMakeRequest(250, new Timestamp(0, 0))).to.be.true; // Once tokens have been used, further requests should fail. - expect(limiter.tryMakeRequest(1, new Timestamp(0, 0))).to.be.false; + expect(limiter._tryMakeRequest(1, new Timestamp(0, 0))).to.be.false; // Tokens will only refill up to max capacity. - expect(limiter.tryMakeRequest(501, new Timestamp(1, 0))).to.be.false; - expect(limiter.tryMakeRequest(500, new Timestamp(1, 0))).to.be.true; + expect(limiter._tryMakeRequest(501, new Timestamp(1, 0))).to.be.false; + expect(limiter._tryMakeRequest(500, new Timestamp(1, 0))).to.be.true; // Tokens will refill incrementally based on the number of ms elapsed. - expect(limiter.tryMakeRequest(251, new Timestamp(1, 500 * 1e6))).to.be + expect(limiter._tryMakeRequest(251, new Timestamp(1, 500 * 1e6))).to.be .false; - expect(limiter.tryMakeRequest(250, new Timestamp(1, 500 * 1e6))).to.be.true; + expect(limiter._tryMakeRequest(250, new Timestamp(1, 500 * 1e6))).to.be + .true; // Scales with multiplier. - expect(limiter.tryMakeRequest(751, new Timestamp(5 * 60 - 1, 0))).to.be + expect(limiter._tryMakeRequest(751, new Timestamp(5 * 60 - 1, 0))).to.be .false; - expect(limiter.tryMakeRequest(751, new Timestamp(5 * 60, 0))).to.be.false; - expect(limiter.tryMakeRequest(750, new Timestamp(5 * 60, 0))).to.be.true; + expect(limiter._tryMakeRequest(751, new Timestamp(5 * 60, 0))).to.be.false; + expect(limiter._tryMakeRequest(750, new Timestamp(5 * 60, 0))).to.be.true; }); it('calculates the number of ms needed to place the next request', () => { // Should return 0 if there are enough tokens for the request to be made. let timestamp = new Timestamp(0, 0); - expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(0); + expect(limiter._getNextRequestDelayMs(500, timestamp)).to.equal(0); // Should factor in remaining tokens when calculating the time. - expect(limiter.tryMakeRequest(250, timestamp)); - expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(500); + expect(limiter._tryMakeRequest(250, timestamp)); + expect(limiter._getNextRequestDelayMs(500, timestamp)).to.equal(500); // Once tokens have been used, should calculate time before next request. timestamp = new Timestamp(1, 0); - expect(limiter.tryMakeRequest(500, timestamp)).to.be.true; - expect(limiter.getNextRequestDelayMs(100, timestamp)).to.equal(200); - expect(limiter.getNextRequestDelayMs(250, timestamp)).to.equal(500); - expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(1000); - expect(limiter.getNextRequestDelayMs(501, timestamp)).to.equal(-1); + expect(limiter._tryMakeRequest(500, timestamp)).to.be.true; + expect(limiter._getNextRequestDelayMs(100, timestamp)).to.equal(200); + expect(limiter._getNextRequestDelayMs(250, timestamp)).to.equal(500); + expect(limiter._getNextRequestDelayMs(500, timestamp)).to.equal(1000); + expect(limiter._getNextRequestDelayMs(501, timestamp)).to.equal(-1); // Scales with multiplier. timestamp = new Timestamp(5 * 60, 0); - expect(limiter.tryMakeRequest(750, timestamp)).to.be.true; - expect(limiter.getNextRequestDelayMs(250, timestamp)).to.equal(334); - expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(667); - expect(limiter.getNextRequestDelayMs(750, timestamp)).to.equal(1000); - expect(limiter.getNextRequestDelayMs(751, timestamp)).to.equal(-1); + expect(limiter._tryMakeRequest(750, timestamp)).to.be.true; + expect(limiter._getNextRequestDelayMs(250, timestamp)).to.equal(334); + expect(limiter._getNextRequestDelayMs(500, timestamp)).to.equal(667); + expect(limiter._getNextRequestDelayMs(750, timestamp)).to.equal(1000); + expect(limiter._getNextRequestDelayMs(751, timestamp)).to.equal(-1); }); it('calculates the maximum number of operations correctly', async () => { From 60c187a2b3f975f9c4ce0954f930ba3a747ecb3e Mon Sep 17 00:00:00 2001 From: Brian Chen <chenbrian@google.com> Date: Wed, 15 Apr 2020 18:56:51 -0700 Subject: [PATCH 3/4] change to use date --- dev/src/rate-limiter.ts | 71 ++++++++++++++------------------- dev/test/rate-limiter.ts | 85 ++++++++++++++++++++++------------------ 2 files changed, 77 insertions(+), 79 deletions(-) diff --git a/dev/src/rate-limiter.ts b/dev/src/rate-limiter.ts index d2473a18e..e6b793631 100644 --- a/dev/src/rate-limiter.ts +++ b/dev/src/rate-limiter.ts @@ -36,44 +36,39 @@ export class RateLimiter { availableTokens: number; // When the token bucket was last refilled. - lastRefillTime = Timestamp.now(); + lastRefillTimeMillis: number; /** * @param initialCapacity Initial maximum number of operations per second. * @param multiplier Rate by which to increase the capacity. * @param multiplierMillis How often the capacity should increase in * milliseconds. - * @param startTime Used for testing the limiter. + * @param startTimeMillis The starting time in epoch milliseconds that the rate limit is based on. Used + * for testing the limiter. */ constructor( private readonly initialCapacity: number, private readonly multiplier: number, private readonly multiplierMillis: number, - private readonly startTime = Timestamp.now() + private readonly startTimeMillis = Date.now() ) { this.availableTokens = initialCapacity; - this.lastRefillTime = startTime; + this.lastRefillTimeMillis = startTimeMillis; } /** - * Tries to make the number of operations. Returns true if the request succeeded - * and false otherwise. + * Tries to make the number of operations. Returns true if the request + * succeeded and false otherwise. * - * @param currentTime Used for testing the limiter. + * @param currentTimeMillis The date used to calculate the number of available + * tokens. Used for testing the limiter. * @private */ - tryMakeRequest(numOperations: number) { - return this._tryMakeRequest(numOperations, Timestamp.now()); - } - - /** - * Used in testing for different timestamps. - * - * @private - */ - // Visible for testing. - _tryMakeRequest(numOperations: number, currentTime: Timestamp): boolean { - this.refillTokens(currentTime); + tryMakeRequest( + numOperations: number, + currentTimeMillis = Date.now() + ): boolean { + this.refillTokens(currentTimeMillis); if (numOperations <= this.availableTokens) { this.availableTokens -= numOperations; return true; @@ -87,27 +82,20 @@ export class RateLimiter { * capacity. Returns -1 if the request is not possible with the current * capacity. * + * @param currentTimeMillis The date used to calculate the number of available + * tokens. Used for testing the limiter. * @private */ - getNextRequestDelayMs(numOperations: number) { - return this._getNextRequestDelayMs(numOperations, Timestamp.now()); - } - /** - * Used in testing for different timestamps. - * - * @private - */ - // Visible for testing. - _getNextRequestDelayMs( + getNextRequestDelayMs( numOperations: number, - currentTime: Timestamp + currentTimeMillis = Date.now() ): number { if (numOperations < this.availableTokens) { return 0; } - const capacity = this.calculateCapacity(currentTime, this.startTime); + const capacity = this.calculateCapacity(currentTimeMillis); if (capacity < numOperations) { return -1; } @@ -120,36 +108,37 @@ export class RateLimiter { * Refills the number of available tokens based on how much time has elapsed * since the last time the tokens were refilled. * + * @param currentTimeMillis The date used to calculate the number of available + * tokens. Used for testing the limiter. * @private */ - private refillTokens(currentTime = Timestamp.now()): void { - if (currentTime.toMillis() > this.lastRefillTime.toMillis()) { - const elapsedTime = - currentTime.toMillis() - this.lastRefillTime.toMillis(); - const capacity = this.calculateCapacity(currentTime, this.startTime); + private refillTokens(currentTimeMillis = Date.now()): void { + if (currentTimeMillis > this.lastRefillTimeMillis) { + const elapsedTime = currentTimeMillis - this.lastRefillTimeMillis; + const capacity = this.calculateCapacity(currentTimeMillis); const tokensToAdd = Math.floor((elapsedTime * capacity) / 1000); if (tokensToAdd > 0) { this.availableTokens = Math.min( capacity, this.availableTokens + tokensToAdd ); - this.lastRefillTime = currentTime; + this.lastRefillTimeMillis = currentTimeMillis; } } } /** - * Calculates the maximum capacity based on the two provided timestamps. + * Calculates the maximum capacity based on the provided date. * * @private */ // Visible for testing. - calculateCapacity(currentTime: Timestamp, startTime: Timestamp): number { + calculateCapacity(currentTimeMillis: number): number { assert( - currentTime.valueOf() >= startTime.valueOf(), + currentTimeMillis >= this.startTimeMillis, 'startTime cannot be after currentTime' ); - const millisElapsed = currentTime.toMillis() - startTime.toMillis(); + const millisElapsed = currentTimeMillis - this.startTimeMillis; const operationsPerSecond = Math.floor( Math.pow( this.multiplier, diff --git a/dev/test/rate-limiter.ts b/dev/test/rate-limiter.ts index 023404d8b..9024149e8 100644 --- a/dev/test/rate-limiter.ts +++ b/dev/test/rate-limiter.ts @@ -14,82 +14,91 @@ import {expect} from 'chai'; -import {Timestamp} from '../src'; import {RateLimiter} from '../src/rate-limiter'; describe('RateLimiter', () => { let limiter: RateLimiter; beforeEach(() => { - limiter = new RateLimiter(500, 1.5, 5 * 60 * 1000, new Timestamp(0, 0)); + limiter = new RateLimiter( + /* initialCapacity= */ 500, + /* multiplier= */ 1.5, + /* multiplierMillis= */ 5 * 60 * 1000, + /* startTime= */ new Date(0).getTime() + ); }); it('accepts and rejects requests based on capacity', () => { - expect(limiter._tryMakeRequest(250, new Timestamp(0, 0))).to.be.true; - expect(limiter._tryMakeRequest(250, new Timestamp(0, 0))).to.be.true; + expect(limiter.tryMakeRequest(250, new Date(0).getTime())).to.be.true; + expect(limiter.tryMakeRequest(250, new Date(0).getTime())).to.be.true; // Once tokens have been used, further requests should fail. - expect(limiter._tryMakeRequest(1, new Timestamp(0, 0))).to.be.false; + expect(limiter.tryMakeRequest(1, new Date(0).getTime())).to.be.false; // Tokens will only refill up to max capacity. - expect(limiter._tryMakeRequest(501, new Timestamp(1, 0))).to.be.false; - expect(limiter._tryMakeRequest(500, new Timestamp(1, 0))).to.be.true; - - // Tokens will refill incrementally based on the number of ms elapsed. - expect(limiter._tryMakeRequest(251, new Timestamp(1, 500 * 1e6))).to.be + expect(limiter.tryMakeRequest(501, new Date(1 * 1000).getTime())).to.be .false; - expect(limiter._tryMakeRequest(250, new Timestamp(1, 500 * 1e6))).to.be + expect(limiter.tryMakeRequest(500, new Date(1 * 1000).getTime())).to.be .true; + // Tokens will refill incrementally based on the number of ms elapsed. + expect(limiter.tryMakeRequest(250, new Date(1 * 1000 + 499).getTime())).to + .be.false; + expect(limiter.tryMakeRequest(249, new Date(1 * 1000 + 500).getTime())).to + .be.true; + // Scales with multiplier. - expect(limiter._tryMakeRequest(751, new Timestamp(5 * 60 - 1, 0))).to.be + expect(limiter.tryMakeRequest(751, new Date((5 * 60 - 1) * 1000).getTime())) + .to.be.false; + expect(limiter.tryMakeRequest(751, new Date(5 * 60 * 1000).getTime())).to.be .false; - expect(limiter._tryMakeRequest(751, new Timestamp(5 * 60, 0))).to.be.false; - expect(limiter._tryMakeRequest(750, new Timestamp(5 * 60, 0))).to.be.true; + expect(limiter.tryMakeRequest(750, new Date(5 * 60 * 1000).getTime())).to.be + .true; + + // Tokens will never exceed capacity. + expect(limiter.tryMakeRequest(751, new Date((5 * 60 + 3) * 1000).getTime())) + .to.be.false; }); it('calculates the number of ms needed to place the next request', () => { // Should return 0 if there are enough tokens for the request to be made. - let timestamp = new Timestamp(0, 0); - expect(limiter._getNextRequestDelayMs(500, timestamp)).to.equal(0); + let timestamp = new Date(0).getTime(); + expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(0); // Should factor in remaining tokens when calculating the time. - expect(limiter._tryMakeRequest(250, timestamp)); - expect(limiter._getNextRequestDelayMs(500, timestamp)).to.equal(500); + expect(limiter.tryMakeRequest(250, timestamp)); + expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(500); // Once tokens have been used, should calculate time before next request. - timestamp = new Timestamp(1, 0); - expect(limiter._tryMakeRequest(500, timestamp)).to.be.true; - expect(limiter._getNextRequestDelayMs(100, timestamp)).to.equal(200); - expect(limiter._getNextRequestDelayMs(250, timestamp)).to.equal(500); - expect(limiter._getNextRequestDelayMs(500, timestamp)).to.equal(1000); - expect(limiter._getNextRequestDelayMs(501, timestamp)).to.equal(-1); + timestamp = new Date(1 * 1000).getTime(); + expect(limiter.tryMakeRequest(500, timestamp)).to.be.true; + expect(limiter.getNextRequestDelayMs(100, timestamp)).to.equal(200); + expect(limiter.getNextRequestDelayMs(250, timestamp)).to.equal(500); + expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(1000); + expect(limiter.getNextRequestDelayMs(501, timestamp)).to.equal(-1); // Scales with multiplier. - timestamp = new Timestamp(5 * 60, 0); - expect(limiter._tryMakeRequest(750, timestamp)).to.be.true; - expect(limiter._getNextRequestDelayMs(250, timestamp)).to.equal(334); - expect(limiter._getNextRequestDelayMs(500, timestamp)).to.equal(667); - expect(limiter._getNextRequestDelayMs(750, timestamp)).to.equal(1000); - expect(limiter._getNextRequestDelayMs(751, timestamp)).to.equal(-1); + timestamp = new Date(5 * 60 * 1000).getTime(); + expect(limiter.tryMakeRequest(750, timestamp)).to.be.true; + expect(limiter.getNextRequestDelayMs(250, timestamp)).to.equal(334); + expect(limiter.getNextRequestDelayMs(500, timestamp)).to.equal(667); + expect(limiter.getNextRequestDelayMs(750, timestamp)).to.equal(1000); + expect(limiter.getNextRequestDelayMs(751, timestamp)).to.equal(-1); }); it('calculates the maximum number of operations correctly', async () => { - const currentTime = new Timestamp(0, 0); - expect( - limiter.calculateCapacity(new Timestamp(1, 0), currentTime) - ).to.equal(500); + expect(limiter.calculateCapacity(new Date(0).getTime())).to.equal(500); expect( - limiter.calculateCapacity(new Timestamp(5 * 60, 0), currentTime) + limiter.calculateCapacity(new Date(5 * 60 * 1000).getTime()) ).to.equal(750); expect( - limiter.calculateCapacity(new Timestamp(10 * 60, 0), currentTime) + limiter.calculateCapacity(new Date(10 * 60 * 1000).getTime()) ).to.equal(1125); expect( - limiter.calculateCapacity(new Timestamp(15 * 60, 0), currentTime) + limiter.calculateCapacity(new Date(15 * 60 * 1000).getTime()) ).to.equal(1687); expect( - limiter.calculateCapacity(new Timestamp(90 * 60, 0), currentTime) + limiter.calculateCapacity(new Date(90 * 60 * 1000).getTime()) ).to.equal(738945); }); }); From a008efd793fc1ee7e66c84ac02ae81d61562fca8 Mon Sep 17 00:00:00 2001 From: Brian Chen <chenbrian@google.com> Date: Thu, 16 Apr 2020 14:43:52 -0700 Subject: [PATCH 4/4] resolve comments, add error for older timestamps --- dev/src/rate-limiter.ts | 39 +++++++++++++++++++++------------------ dev/test/rate-limiter.ts | 5 +++++ 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/dev/src/rate-limiter.ts b/dev/src/rate-limiter.ts index e6b793631..e2d7be390 100644 --- a/dev/src/rate-limiter.ts +++ b/dev/src/rate-limiter.ts @@ -43,8 +43,8 @@ export class RateLimiter { * @param multiplier Rate by which to increase the capacity. * @param multiplierMillis How often the capacity should increase in * milliseconds. - * @param startTimeMillis The starting time in epoch milliseconds that the rate limit is based on. Used - * for testing the limiter. + * @param startTimeMillis The starting time in epoch milliseconds that the + * rate limit is based on. Used for testing the limiter. */ constructor( private readonly initialCapacity: number, @@ -60,15 +60,15 @@ export class RateLimiter { * Tries to make the number of operations. Returns true if the request * succeeded and false otherwise. * - * @param currentTimeMillis The date used to calculate the number of available + * @param requestTimeMillis The date used to calculate the number of available * tokens. Used for testing the limiter. * @private */ tryMakeRequest( numOperations: number, - currentTimeMillis = Date.now() + requestTimeMillis = Date.now() ): boolean { - this.refillTokens(currentTimeMillis); + this.refillTokens(requestTimeMillis); if (numOperations <= this.availableTokens) { this.availableTokens -= numOperations; return true; @@ -82,20 +82,19 @@ export class RateLimiter { * capacity. Returns -1 if the request is not possible with the current * capacity. * - * @param currentTimeMillis The date used to calculate the number of available + * @param requestTimeMillis The date used to calculate the number of available * tokens. Used for testing the limiter. * @private */ - getNextRequestDelayMs( numOperations: number, - currentTimeMillis = Date.now() + requestTimeMillis = Date.now() ): number { if (numOperations < this.availableTokens) { return 0; } - const capacity = this.calculateCapacity(currentTimeMillis); + const capacity = this.calculateCapacity(requestTimeMillis); if (capacity < numOperations) { return -1; } @@ -108,22 +107,26 @@ export class RateLimiter { * Refills the number of available tokens based on how much time has elapsed * since the last time the tokens were refilled. * - * @param currentTimeMillis The date used to calculate the number of available + * @param requestTimeMillis The date used to calculate the number of available * tokens. Used for testing the limiter. * @private */ - private refillTokens(currentTimeMillis = Date.now()): void { - if (currentTimeMillis > this.lastRefillTimeMillis) { - const elapsedTime = currentTimeMillis - this.lastRefillTimeMillis; - const capacity = this.calculateCapacity(currentTimeMillis); + private refillTokens(requestTimeMillis = Date.now()): void { + if (requestTimeMillis >= this.lastRefillTimeMillis) { + const elapsedTime = requestTimeMillis - this.lastRefillTimeMillis; + const capacity = this.calculateCapacity(requestTimeMillis); const tokensToAdd = Math.floor((elapsedTime * capacity) / 1000); if (tokensToAdd > 0) { this.availableTokens = Math.min( capacity, this.availableTokens + tokensToAdd ); - this.lastRefillTimeMillis = currentTimeMillis; + this.lastRefillTimeMillis = requestTimeMillis; } + } else { + throw new Error( + 'Request time should not be before the last token refill time.' + ); } } @@ -133,12 +136,12 @@ export class RateLimiter { * @private */ // Visible for testing. - calculateCapacity(currentTimeMillis: number): number { + calculateCapacity(requestTimeMillis: number): number { assert( - currentTimeMillis >= this.startTimeMillis, + requestTimeMillis >= this.startTimeMillis, 'startTime cannot be after currentTime' ); - const millisElapsed = currentTimeMillis - this.startTimeMillis; + const millisElapsed = requestTimeMillis - this.startTimeMillis; const operationsPerSecond = Math.floor( Math.pow( this.multiplier, diff --git a/dev/test/rate-limiter.ts b/dev/test/rate-limiter.ts index 9024149e8..e3271303d 100644 --- a/dev/test/rate-limiter.ts +++ b/dev/test/rate-limiter.ts @@ -58,6 +58,11 @@ describe('RateLimiter', () => { // Tokens will never exceed capacity. expect(limiter.tryMakeRequest(751, new Date((5 * 60 + 3) * 1000).getTime())) .to.be.false; + + // Rejects requests made before lastRefillTime + expect(() => + limiter.tryMakeRequest(751, new Date((5 * 60 + 2) * 1000).getTime()) + ).to.throw('Request time should not be before the last token refill time.'); }); it('calculates the number of ms needed to place the next request', () => {