fix: add RateLimiter (#230)

Author: Brian Chen

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 com.google.common.base.Preconditions;
+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 <a href=https://cloud.google.com/datastore/docs/best-practices#ramping_up_traffic>Ramping up
+ * traffic</a>
+ */
+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) {
+ Preconditions.checkArgument(
+ requestTimeMillis >= lastRefillTimeMillis,
+ "Request time should not be before the last token refill time");
+ 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;
+ }
+ }
+
+ 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()));
+ }
+}