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', () => {