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. + * + * @see Ramping up + * traffic + */ +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; + } +} diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/RateLimiterTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/RateLimiterTest.java new file mode 100644 index 000000000..291f60733 --- /dev/null +++ b/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())); + } +}