feat: add max/min throttling options to BulkWriterOptions

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

@@ -43,7 +43,7 @@
 
 final class BulkWriter implements AutoCloseable {
   /** The maximum number of writes that can be in a single batch. */
-  public static final int MAX_BATCH_SIZE = 500;
+  public static final int MAX_BATCH_SIZE = 20;
 
   public static final int MAX_RETRY_ATTEMPTS = 10;
 
@@ -53,7 +53,7 @@ final class BulkWriter implements AutoCloseable {
    * @see <a href=https://cloud.google.com/datastore/docs/best-practices#ramping_up_traffic>Ramping
    *     up traffic</a>
    */
-  private static final int STARTING_MAXIMUM_OPS_PER_SECOND = 500;
+  static final int DEFAULT_STARTING_MAXIMUM_OPS_PER_SECOND = 500;
 
   /**
    * The rate by which to increase the capacity as specified by the 500/50/5 rule.
@@ -97,22 +97,51 @@ final class BulkWriter implements AutoCloseable {
   private final ExponentialRetryAlgorithm backoff;
   private TimedAttemptSettings nextAttempt;
 
-  BulkWriter(FirestoreImpl firestore, boolean enableThrottling) {
+  BulkWriter(FirestoreImpl firestore, BulkWriterOptions options) {
+    validateBulkWriterOptions(options);
     this.firestore = firestore;
     this.backoff =
         new ExponentialRetryAlgorithm(
             firestore.getOptions().getRetrySettings(), CurrentMillisClock.getDefaultClock());
     this.nextAttempt = backoff.createFirstAttempt();
     this.firestoreExecutor = firestore.getClient().getExecutor();
 
-    if (enableThrottling) {
-      rateLimiter =
+    if (!options.isThrottlingEnabled()) {
+      this.rateLimiter =
           new RateLimiter(
-              STARTING_MAXIMUM_OPS_PER_SECOND,
-              RATE_LIMITER_MULTIPLIER,
-              RATE_LIMITER_MULTIPLIER_MILLIS);
+              Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE);
     } else {
-      rateLimiter = new RateLimiter(Integer.MAX_VALUE, Double.MAX_VALUE, Integer.MAX_VALUE);
+      double startingRate = DEFAULT_STARTING_MAXIMUM_OPS_PER_SECOND;
+      double maxRate = Integer.MAX_VALUE;
+
+      if (options.getInitialOpsPerSecond() != BulkWriterOptions.DEFAULT_UNSET_VALUE) {
+        startingRate = options.getInitialOpsPerSecond();
+      }
+
+      if (options.getMaxOpsPerSecond() != BulkWriterOptions.DEFAULT_UNSET_VALUE) {
+        maxRate = options.getMaxOpsPerSecond();
+      }
+
+      // The initial validation step ensures that the maxOpsPerSecond is
+      // greater than initialOpsPerSecond. If this inequality is true, that
+      // means initialOpsPerSecond was not set and maxOpsPerSecond is less
+      // than the default starting rate.
+      if (maxRate < startingRate) {
+        startingRate = maxRate;
+      }
+
+      // Ensure that the batch size is not larger than the number of allowed
+      // operations per second.
+      if (startingRate < maxBatchSize) {
+        this.maxBatchSize = (int) startingRate;
+      }
+
+      this.rateLimiter =
+          new RateLimiter(
+              (int) startingRate,
+              RATE_LIMITER_MULTIPLIER,
+              RATE_LIMITER_MULTIPLIER_MILLIS,
+              (int) maxRate);
     }
   }
 
@@ -679,4 +708,33 @@ public ApiFuture<Void> apply(List<BatchWriteResult> results) {
   void setMaxBatchSize(int size) {
     maxBatchSize = size;
   }
+
+  @VisibleForTesting
+  RateLimiter getRateLimiter() {
+    return rateLimiter;
+  }
+
+  private void validateBulkWriterOptions(BulkWriterOptions options) {
+    double initialRate = options.getInitialOpsPerSecond();
+    double maxRate = options.getMaxOpsPerSecond();
+
+    if (initialRate < 1) {
+      throw FirestoreException.invalidState(
+          "Value for argument 'initialOpsPerSecond' must be an integer within [1, Infinity], but was: "
+              + (int) initialRate);
+    }
+
+    if (maxRate < 1) {
+      throw FirestoreException.invalidState(
+          "Value for argument 'maxOpsPerSecond' must be an integer within [1, Infinity], but was: "
+              + (int) maxRate);
+    }
+
+    if (initialRate != BulkWriterOptions.DEFAULT_UNSET_VALUE
+        && maxRate != BulkWriterOptions.DEFAULT_UNSET_VALUE
+        && initialRate > maxRate) {
+      throw FirestoreException.invalidState(
+          "'maxOpsPerSecond' cannot be less than 'initialOpsPerSecond'.");
+    }
+  }
 }

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

@@ -20,15 +20,31 @@
 
 /** Options used to disable request throttling in BulkWriter. */
 final class BulkWriterOptions {
+  static final double DEFAULT_UNSET_VALUE = 1.1;
+  private final boolean throttling;
+  private double initialOpsPerSecond = DEFAULT_UNSET_VALUE;
+  private double maxOpsPerSecond = DEFAULT_UNSET_VALUE;
 
-  private final boolean enableThrottling;
+  BulkWriterOptions(boolean enableThrottling) {
+    this.throttling = enableThrottling;
+  }
 
-  private BulkWriterOptions(boolean enableThrottling) {
-    this.enableThrottling = enableThrottling;
+  private BulkWriterOptions(double initialOpsPerSecond, double maxOpsPerSecond) {
+    this.throttling = true;
+    this.initialOpsPerSecond = initialOpsPerSecond;
+    this.maxOpsPerSecond = maxOpsPerSecond;
   }
 
   boolean isThrottlingEnabled() {
-    return enableThrottling;
+    return throttling;
+  }
+
+  double getInitialOpsPerSecond() {
+    return initialOpsPerSecond;
+  }
+
+  double getMaxOpsPerSecond() {
+    return maxOpsPerSecond;
   }
 
   /**
@@ -40,4 +56,48 @@ boolean isThrottlingEnabled() {
   public static BulkWriterOptions withThrottlingDisabled() {
     return new BulkWriterOptions(false);
   }
+
+  /**
+   * An options object that will enable throttling in the created BulkWriter with the provided
+   * starting rate.
+   *
+   * @param initialOpsPerSecond The initial maximum number of operations per second allowed by the
+   *     throttler.
+   * @return The BulkWriterOptions object.
+   */
+  @Nonnull
+  public static BulkWriterOptions withInitialOpsPerSecondThrottling(int initialOpsPerSecond) {
+    return new BulkWriterOptions(initialOpsPerSecond, DEFAULT_UNSET_VALUE);
+  }
+
+  /**
+   * An options object that will enable throttling in the created BulkWriter with the provided
+   * maximum rate.
+   *
+   * @param maxOpsPerSecond The maximum number of operations per second allowed by the throttler.
+   *     The throttler's allowed operations per second does not ramp up past the specified
+   *     operations per second.
+   * @return The BulkWriterOptions object.
+   */
+  @Nonnull
+  public static BulkWriterOptions withMaxOpsPerSecondThrottling(int maxOpsPerSecond) {
+    return new BulkWriterOptions(DEFAULT_UNSET_VALUE, maxOpsPerSecond);
+  }
+
+  /**
+   * An options object that will enable throttling in the created BulkWriter with the provided
+   * starting and maximum rate.
+   *
+   * @param initialOpsPerSecond The initial maximum number of operations per second allowed by the
+   *     throttler.
+   * @param maxOpsPerSecond The maximum number of operations per second allowed by the throttler.
+   *     The throttler's allowed operations per second does not ramp up past the specified
+   *     operations per second.
+   * @return The BulkWriterOptions object.
+   */
+  @Nonnull
+  public static BulkWriterOptions withCustomThrottling(
+      int initialOpsPerSecond, int maxOpsPerSecond) {
+    return new BulkWriterOptions(initialOpsPerSecond, maxOpsPerSecond);
+  }
 }

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

@@ -94,12 +94,12 @@ public WriteBatch batch() {
 
   @Nonnull
   BulkWriter bulkWriter() {
-    return new BulkWriter(this, /* enableThrottling= */ true);
+    return new BulkWriter(this, new BulkWriterOptions(true));
   }
 
   @Nonnull
   BulkWriter bulkWriter(BulkWriterOptions options) {
-    return new BulkWriter(this, options.isThrottlingEnabled());
+    return new BulkWriter(this, options);
   }
 
   @Nonnull

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

@@ -38,31 +38,48 @@ class RateLimiter {
   private final double multiplier;
   private final int multiplierMillis;
   private final long startTimeMillis;
+  private final int maximumCapacity;
 
   private int availableTokens;
   private long lastRefillTimeMillis;
 
-  RateLimiter(int initialCapacity, double multiplier, int multiplierMillis) {
-    this(initialCapacity, multiplier, multiplierMillis, new Date().getTime());
+  RateLimiter(int initialCapacity, double multiplier, int multiplierMillis, int maximumCapacity) {
+    this(initialCapacity, multiplier, multiplierMillis, maximumCapacity, 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 maximumCapacity Maximum number of allowed operations per second. The number of tokens
+   *     added per second will never exceed this number.
    * @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) {
+  RateLimiter(
+      int initialCapacity,
+      double multiplier,
+      int multiplierMillis,
+      int maximumCapacity,
+      long startTimeMillis) {
     this.initialCapacity = initialCapacity;
     this.multiplier = multiplier;
     this.multiplierMillis = multiplierMillis;
+    this.maximumCapacity = maximumCapacity;
     this.startTimeMillis = startTimeMillis;
 
     this.availableTokens = initialCapacity;
     this.lastRefillTimeMillis = startTimeMillis;
   }
 
+  public int getInitialCapacity() {
+    return initialCapacity;
+  }
+
+  public int getMaximumCapacity() {
+    return maximumCapacity;
+  }
+
   public boolean tryMakeRequest(int numOperations) {
     return tryMakeRequest(numOperations, new Date().getTime());
   }
@@ -132,7 +149,10 @@ private void refillTokens(long requestTimeMillis) {
   public int calculateCapacity(long requestTimeMillis) {
     long millisElapsed = requestTimeMillis - startTimeMillis;
     int operationsPerSecond =
-        (int) (Math.pow(multiplier, (int) (millisElapsed / multiplierMillis)) * initialCapacity);
+        Math.min(
+            (int)
+                (Math.pow(multiplier, (int) (millisElapsed / multiplierMillis)) * initialCapacity),
+            maximumCapacity);
     return operationsPerSecond;
   }
 }

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

@@ -50,7 +50,6 @@
 import javax.annotation.Nonnull;
 import org.junit.Assert;
 import org.junit.Before;
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.Timeout;
@@ -549,8 +548,6 @@ public void flushCompletesWhenAllWritesComplete() throws Exception {
   }
 
   @Test
-  @Ignore
-  // TODO(chenbrian): Fix this test after throttling options are added.
   public void doesNotSendBatchesIfDoingSoExceedsRateLimit() {
     final boolean[] timeoutCalled = {false};
     final ScheduledExecutorService timeoutExecutor =
@@ -565,7 +562,8 @@ public ScheduledFuture<?> schedule(Runnable command, long delay, TimeUnit unit)
           }
         };
     doReturn(timeoutExecutor).when(firestoreRpc).getExecutor();
-    BulkWriter bulkWriter = firestoreMock.bulkWriter();
+    BulkWriter bulkWriter =
+        firestoreMock.bulkWriter(BulkWriterOptions.withInitialOpsPerSecondThrottling(5));
 
     for (int i = 0; i < 600; ++i) {
       bulkWriter.set(firestoreMock.document("coll/doc" + i), LocalFirestoreHelper.SINGLE_FIELD_MAP);
@@ -706,4 +704,65 @@ public ApiFuture<Object> answer(InvocationOnMock mock) {
       assertEquals(retryAttempts[0], BulkWriter.MAX_RETRY_ATTEMPTS + 1);
     }
   }
+
+  @Test
+  public void optionsRequiresPositiveInteger() throws Exception {
+    try {
+      firestoreMock.bulkWriter(BulkWriterOptions.withInitialOpsPerSecondThrottling(-1));
+      fail("bulkWriter() call should have failed");
+    } catch (Exception e) {
+      assertEquals(
+          e.getMessage(),
+          "Value for argument 'initialOpsPerSecond' must be an integer within [1, Infinity], but was: -1");
+    }
+
+    try {
+      firestoreMock.bulkWriter(BulkWriterOptions.withMaxOpsPerSecondThrottling(-1));
+      fail("bulkWriter() call should have failed");
+    } catch (Exception e) {
+      assertEquals(
+          e.getMessage(),
+          "Value for argument 'maxOpsPerSecond' must be an integer within [1, Infinity], but was: -1");
+    }
+  }
+
+  @Test
+  public void optionsRequiresMaxGreaterThanInitial() throws Exception {
+    try {
+      firestoreMock.bulkWriter(BulkWriterOptions.withCustomThrottling(550, 500));
+      fail("bulkWriter() call should have failed");
+    } catch (Exception e) {
+      assertEquals(e.getMessage(), "'maxOpsPerSecond' cannot be less than 'initialOpsPerSecond'.");
+    }
+  }
+
+  @Test
+  public void optionsInitialAndMaxRatesAreProperlySet() throws Exception {
+    BulkWriter bulkWriter =
+        firestoreMock.bulkWriter(BulkWriterOptions.withCustomThrottling(500, 550));
+    assertEquals(bulkWriter.getRateLimiter().getInitialCapacity(), 500);
+    assertEquals(bulkWriter.getRateLimiter().getMaximumCapacity(), 550);
+
+    bulkWriter = firestoreMock.bulkWriter(BulkWriterOptions.withMaxOpsPerSecondThrottling(1000));
+    assertEquals(bulkWriter.getRateLimiter().getInitialCapacity(), 500);
+    assertEquals(bulkWriter.getRateLimiter().getMaximumCapacity(), 1000);
+
+    bulkWriter = firestoreMock.bulkWriter(BulkWriterOptions.withInitialOpsPerSecondThrottling(100));
+    assertEquals(bulkWriter.getRateLimiter().getInitialCapacity(), 100);
+    assertEquals(bulkWriter.getRateLimiter().getMaximumCapacity(), Integer.MAX_VALUE);
+
+    bulkWriter = firestoreMock.bulkWriter(BulkWriterOptions.withMaxOpsPerSecondThrottling(100));
+    assertEquals(bulkWriter.getRateLimiter().getInitialCapacity(), 100);
+    assertEquals(bulkWriter.getRateLimiter().getMaximumCapacity(), 100);
+
+    bulkWriter = firestoreMock.bulkWriter();
+    assertEquals(
+        bulkWriter.getRateLimiter().getInitialCapacity(),
+        BulkWriter.DEFAULT_STARTING_MAXIMUM_OPS_PER_SECOND);
+    assertEquals(bulkWriter.getRateLimiter().getMaximumCapacity(), Integer.MAX_VALUE);
+
+    bulkWriter = firestoreMock.bulkWriter(BulkWriterOptions.withThrottlingDisabled());
+    assertEquals(bulkWriter.getRateLimiter().getInitialCapacity(), Integer.MAX_VALUE);
+    assertEquals(bulkWriter.getRateLimiter().getMaximumCapacity(), Integer.MAX_VALUE);
+  }
 }

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

@@ -38,6 +38,7 @@ public void before() {
             /* initialCapacity= */ 500,
             /* multiplier= */ 1.5,
             /* multiplierMillis= */ 5 * 60 * 1000,
+            /* maximumCapacity= */ 1000000,
             /* startTime= */ new Date(0).getTime());
   }
 
@@ -108,5 +109,8 @@ public void calculatesMaxOperations() {
     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()));
+
+    // Check that maximum rate limit is enforced.
+    assertEquals(1000000, limiter.calculateCapacity(new Date(1000 * 60 * 1000).getTime()));
   }
 }