fix: add RateLimiter

google-cloud-firestore/src/main/java/com/google/cloud/firestore/RateLimiter.java

@@ -0,0 +1,138 @@
+/*
+ * 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.
+ */
+
+package com.google.cloud.firestore;
+
+import java.util.Date;
+
+/**
+ * A helper that uses the Token Bucket algorithm to rate limit the number of operations that can be
+ * made in a second.
+ *
+ * <p>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.
+ *
+ * <p>RateLimiter can also implement a gradually increasing rate limit. This is used to enforce the
+ * 500/50/5 rule.
+ *
+ * @see "https://cloud.google.com/datastore/docs/best-practices#ramping_up_traffic"
+ */
+public class RateLimiter {
+  private final int initialCapacity;
+  private final double multiplier;
+  private final int multiplierMillis;
+  private final long startTimeMillis;
+
+  private int availableTokens;
+  private long lastRefillTimeMillis;
+
+  RateLimiter(int initialCapacity, int multiplier, int multiplierMillis) {
+    this(initialCapacity, multiplier, multiplierMillis, new Date().getTime());
+  }
+
+  /**
+   * @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 startTimeMillis The starting time in epoch milliseconds that the rate limit is based on.
+   *     Used for testing the limiter.
+   */
+  RateLimiter(int initialCapacity, double multiplier, int multiplierMillis, long startTimeMillis) {
+    this.initialCapacity = initialCapacity;
+    this.multiplier = multiplier;
+    this.multiplierMillis = multiplierMillis;
+    this.startTimeMillis = startTimeMillis;
+
+    this.availableTokens = initialCapacity;
+    this.lastRefillTimeMillis = startTimeMillis;
+  }
+
+  public boolean tryMakeRequest(int numOperations) {
+    return tryMakeRequest(numOperations, new Date(0).getTime());
+  }
+
+  /**
+   * Tries to make the number of operations. Returns true if the request succeeded and false
+   * otherwise.
+   *
+   * @param requestTimeMillis The time used to calculate the number of available tokens. Used for
+   *     testing the limiter.
+   */
+  public boolean tryMakeRequest(int numOperations, long requestTimeMillis) {
+    refillTokens(requestTimeMillis);
+    if (numOperations <= availableTokens) {
+      availableTokens -= numOperations;
+      return true;
+    }
+    return false;
+  }
+
+  public long getNextRequestDelayMs(int numOperations) {
+    return getNextRequestDelayMs(numOperations, new Date().getTime());
+  }
+
+  /**
+   * 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 requestTimeMillis The time used to calculate the number of available tokens. Used for
+   *     testing the limiter.
+   */
+  public long getNextRequestDelayMs(int numOperations, long requestTimeMillis) {
+    if (numOperations < availableTokens) {
+      return 0;
+    }
+
+    int capacity = calculateCapacity(requestTimeMillis);
+    if (capacity < numOperations) {
+      return -1;
+    }
+
+    int requiredTokens = numOperations - availableTokens;
+    return (long) Math.ceil((double) (requiredTokens * 1000) / capacity);
+  }
+
+  /**
+   * Refills the number of available tokens based on how much time has elapsed since the last time
+   * the tokens were refilled.
+   *
+   * @param requestTimeMillis The time used to calculate the number of available tokens. Used for
+   *     testing the limiter.
+   */
+  private void refillTokens(long requestTimeMillis) {
+    if (requestTimeMillis >= lastRefillTimeMillis) {
+      long elapsedTime = requestTimeMillis - lastRefillTimeMillis;
+      int capacity = calculateCapacity(requestTimeMillis);
+      int tokensToAdd = (int) ((elapsedTime * capacity) / 1000);
+      if (tokensToAdd > 0) {
+        availableTokens = Math.min(capacity, availableTokens + tokensToAdd);
+        lastRefillTimeMillis = requestTimeMillis;
+      }
+    } else {
+      throw new IllegalArgumentException(
+          "Request time should not be before the last token refill time");
+    }
+  }
+
+  public int calculateCapacity(long requestTimeMillis) {
+    long millisElapsed = requestTimeMillis - startTimeMillis;
+    int operationsPerSecond =
+        (int) (Math.pow(multiplier, (int) (millisElapsed / multiplierMillis)) * initialCapacity);
+    return operationsPerSecond;
+  }
+}

google-cloud-firestore/src/test/java/com/google/cloud/firestore/RateLimiterTest.java

@@ -0,0 +1,112 @@
+/*
+ * 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.
+ */
+
+package com.google.cloud.firestore;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import java.util.Date;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.runners.MockitoJUnitRunner;
+
+@RunWith(MockitoJUnitRunner.class)
+public class RateLimiterTest {
+  private RateLimiter limiter;
+
+  @Before
+  public void before() {
+    limiter =
+        new RateLimiter(
+            /* initialCapacity= */ 500,
+            /* multiplier= */ 1.5,
+            /* multiplierMillis= */ 5 * 60 * 1000,
+            /* startTime= */ new Date(0).getTime());
+  }
+
+  @Test
+  public void processRequestsFromCapacity() {
+    assertTrue(limiter.tryMakeRequest(250, new Date(0).getTime()));
+    assertTrue(limiter.tryMakeRequest(250, new Date(0).getTime()));
+
+    // Once tokens have been used, further requests should fail.
+    assertFalse(limiter.tryMakeRequest(1, new Date(0).getTime()));
+
+    // Tokens will only refill up to max capacity.
+    assertFalse(limiter.tryMakeRequest(501, new Date(1 * 1000).getTime()));
+    assertTrue(limiter.tryMakeRequest(500, new Date(1 * 1000).getTime()));
+
+    // Tokens will refill incrementally based on number of ms elapsed.
+    assertFalse(limiter.tryMakeRequest(250, new Date(1 * 1000 + 499).getTime()));
+    assertTrue(limiter.tryMakeRequest(249, new Date(1 * 1000 + 500).getTime()));
+
+    // Scales with multiplier.
+    assertFalse(limiter.tryMakeRequest(751, new Date((5 * 60 - 1) * 1000).getTime()));
+    assertFalse(limiter.tryMakeRequest(751, new Date(5 * 60 * 1000).getTime()));
+    assertTrue(limiter.tryMakeRequest(750, new Date(5 * 60 * 1000).getTime()));
+
+    // Tokens will never exceed capacity.
+    assertFalse(limiter.tryMakeRequest(751, new Date((5 * 60 + 3) * 1000).getTime()));
+
+    // Rejects requests made before lastRefillTime.
+    try {
+      limiter.tryMakeRequest(751, new Date((5 * 60 + 2) * 1000).getTime());
+      fail();
+    } catch (IllegalArgumentException e) {
+      assertEquals("Request time should not be before the last token refill time", e.getMessage());
+    }
+  }
+
+  @Test
+  public void calculatesMsForNextRequest() {
+    // Should return 0 if there are enough tokens for the request to be made.
+    long timestamp = new Date(0).getTime();
+    assertEquals(0, limiter.getNextRequestDelayMs(500, timestamp));
+
+    // Should factor in remaining tokens when calculating the time.
+    assertTrue(limiter.tryMakeRequest(250, timestamp));
+    assertEquals(500, limiter.getNextRequestDelayMs(500, timestamp));
+
+    // Once tokens have been used, should calculate time before next request.
+    timestamp = new Date(1 * 1000).getTime();
+    assertTrue(limiter.tryMakeRequest(500, timestamp));
+    assertEquals(200, limiter.getNextRequestDelayMs(100, timestamp));
+    assertEquals(500, limiter.getNextRequestDelayMs(250, timestamp));
+    assertEquals(1000, limiter.getNextRequestDelayMs(500, timestamp));
+    assertEquals(-1, limiter.getNextRequestDelayMs(501, timestamp));
+
+    // Scales with multiplier.
+    timestamp = new Date(5 * 60 * 1000).getTime();
+    assertTrue(limiter.tryMakeRequest(750, timestamp));
+    assertEquals(334, limiter.getNextRequestDelayMs(250, timestamp));
+    assertEquals(667, limiter.getNextRequestDelayMs(500, timestamp));
+    assertEquals(1000, limiter.getNextRequestDelayMs(750, timestamp));
+    assertEquals(-1, limiter.getNextRequestDelayMs(751, timestamp));
+  }
+
+  @Test
+  public void calculatesMaxOperations() {
+    assertEquals(500, limiter.calculateCapacity(new Date(0).getTime()));
+    assertEquals(750, limiter.calculateCapacity(new Date(5 * 60 * 1000).getTime()));
+    assertEquals(1125, limiter.calculateCapacity(new Date(10 * 60 * 1000).getTime()));
+    assertEquals(1687, limiter.calculateCapacity(new Date(15 * 60 * 1000).getTime()));
+    assertEquals(738945, limiter.calculateCapacity(new Date(90 * 60 * 1000).getTime()));
+  }
+}