feat: dynamic flow control part 1 - add FlowController to Batcher (#1289)

* feat: dynamic flow control part 1 - add FlowController to Batcher

* update test

* make sure to close batchers in the test

* Add internal api descriptions

* updates based on review

* fix formatting

* fixing doc

* updates based on code review

Co-authored-by: Igor Bernstein <igorbernstein@google.com>

Author: mutianf

gax/src/main/java/com/google/api/gax/batching/BatcherImpl.java

@@ -37,6 +37,9 @@
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
 import com.google.api.core.SettableApiFuture;
+import com.google.api.gax.batching.FlowController.FlowControlException;
+import com.google.api.gax.batching.FlowController.FlowControlRuntimeException;
+import com.google.api.gax.batching.FlowController.LimitExceededBehavior;
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -55,6 +58,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import javax.annotation.Nullable;
 
 /**
  * Queues up the elements until {@link #flush()} is called; once batching is over, returned future
@@ -87,13 +91,14 @@ public class BatcherImpl<ElementT, ElementResultT, RequestT, ResponseT>
   private final Future<?> scheduledFuture;
   private volatile boolean isClosed = false;
   private final BatcherStats batcherStats = new BatcherStats();
+  private final FlowController flowController;
 
   /**
    * @param batchingDescriptor a {@link BatchingDescriptor} for transforming individual elements
-   *     into wrappers request and response.
-   * @param unaryCallable a {@link UnaryCallable} object.
-   * @param prototype a {@link RequestT} object.
-   * @param batchingSettings a {@link BatchingSettings} with configuration of thresholds.
+   *     into wrappers request and response
+   * @param unaryCallable a {@link UnaryCallable} object
+   * @param prototype a {@link RequestT} object
+   * @param batchingSettings a {@link BatchingSettings} with configuration of thresholds
    */
   public BatcherImpl(
       BatchingDescriptor<ElementT, ElementResultT, RequestT, ResponseT> batchingDescriptor,
@@ -102,15 +107,56 @@ public BatcherImpl(
       BatchingSettings batchingSettings,
       ScheduledExecutorService executor) {
 
+    this(batchingDescriptor, unaryCallable, prototype, batchingSettings, executor, null);
+  }
+
+  /**
+   * @param batchingDescriptor a {@link BatchingDescriptor} for transforming individual elements
+   *     into wrappers request and response
+   * @param unaryCallable a {@link UnaryCallable} object
+   * @param prototype a {@link RequestT} object
+   * @param batchingSettings a {@link BatchingSettings} with configuration of thresholds
+   * @param flowController a {@link FlowController} for throttling requests. If it's null, create a
+   *     {@link FlowController} object from {@link BatchingSettings#getFlowControlSettings()}.
+   */
+  public BatcherImpl(
+      BatchingDescriptor<ElementT, ElementResultT, RequestT, ResponseT> batchingDescriptor,
+      UnaryCallable<RequestT, ResponseT> unaryCallable,
+      RequestT prototype,
+      BatchingSettings batchingSettings,
+      ScheduledExecutorService executor,
+      @Nullable FlowController flowController) {
+
     this.batchingDescriptor =
         Preconditions.checkNotNull(batchingDescriptor, "batching descriptor cannot be null");
     this.unaryCallable = Preconditions.checkNotNull(unaryCallable, "callable cannot be null");
     this.prototype = Preconditions.checkNotNull(prototype, "request prototype cannot be null");
     this.batchingSettings =
         Preconditions.checkNotNull(batchingSettings, "batching setting cannot be null");
     Preconditions.checkNotNull(executor, "executor cannot be null");
+    if (flowController == null) {
+      flowController = new FlowController(batchingSettings.getFlowControlSettings());
+    }
+    // If throttling is enabled, make sure flow control limits are greater or equal to batch sizes
+    // to avoid deadlocking
+    if (flowController.getLimitExceededBehavior() != LimitExceededBehavior.Ignore) {
+      Preconditions.checkArgument(
+          flowController.getMaxOutstandingElementCount() == null
+              || batchingSettings.getElementCountThreshold() == null
+              || flowController.getMaxOutstandingElementCount()
+                  >= batchingSettings.getElementCountThreshold(),
+          "If throttling and batching on element count are enabled, FlowController"
+              + "#maxOutstandingElementCount must be greater or equal to elementCountThreshold");
+      Preconditions.checkArgument(
+          flowController.getMaxOutstandingRequestBytes() == null
+              || batchingSettings.getRequestByteThreshold() == null
+              || flowController.getMaxOutstandingRequestBytes()
+                  >= batchingSettings.getRequestByteThreshold(),
+          "If throttling and batching on request bytes are enabled, FlowController"
+              + "#maxOutstandingRequestBytes must be greater or equal to requestByteThreshold");
+    }
+    this.flowController = flowController;
     currentOpenBatch = new Batch<>(prototype, batchingDescriptor, batchingSettings, batcherStats);
-
     if (batchingSettings.getDelayThreshold() != null) {
       long delay = batchingSettings.getDelayThreshold().toMillis();
       PushCurrentBatchRunnable<ElementT, ElementResultT, RequestT, ResponseT> runnable =
@@ -127,8 +173,29 @@ public BatcherImpl(
   @Override
   public ApiFuture<ElementResultT> add(ElementT element) {
     Preconditions.checkState(!isClosed, "Cannot add elements on a closed batcher");
-    SettableApiFuture<ElementResultT> result = SettableApiFuture.create();
+    // This is not the optimal way of throttling. It does not send out partial batches, which
+    // means that the Batcher might not use up all the resources allowed by FlowController.
+    // The more efficient implementation should look like:
+    // if (!flowController.tryReserve(1, bytes)) {
+    //   sendOutstanding();
+    //   reserve(1, bytes);
+    // }
+    // where tryReserve() will return false if there isn't enough resources, or reserve and return
+    // true.
+    // However, with the current FlowController implementation, adding a tryReserve() could be
+    // confusing. FlowController will end up having 3 different reserve behaviors: blocking,
+    // non blocking and try reserve. And we'll also need to add a tryAcquire() to the Semaphore64
+    // class, which made it seem unnecessary to have blocking and non-blocking semaphore
+    // implementations. Some refactoring may be needed for the optimized implementation. So we'll
+    // defer it till we decide on if refactoring FlowController is necessary.
+    try {
+      flowController.reserve(1, batchingDescriptor.countBytes(element));
+    } catch (FlowControlException e) {
+      // This exception will only be thrown if the FlowController is set to ThrowException behavior
+      throw FlowControlRuntimeException.fromFlowControlException(e);
+    }
 
+    SettableApiFuture<ElementResultT> result = SettableApiFuture.create();
     synchronized (elementLock) {
       currentOpenBatch.add(element, result);
     }
@@ -169,6 +236,7 @@ public void sendOutstanding() {
           @Override
           public void onSuccess(ResponseT response) {
             try {
+              flowController.release(accumulatedBatch.elementCounter, accumulatedBatch.byteCounter);
               accumulatedBatch.onBatchSuccess(response);
             } finally {
               onBatchCompletion();
@@ -178,6 +246,7 @@ public void onSuccess(ResponseT response) {
           @Override
           public void onFailure(Throwable throwable) {
             try {
+              flowController.release(accumulatedBatch.elementCounter, accumulatedBatch.byteCounter);
               accumulatedBatch.onBatchFailure(throwable);
             } finally {
               onBatchCompletion();
@@ -224,6 +293,12 @@ public void close() throws InterruptedException {
     }
   }
 
+  /** Package-private for use in testing. */
+  @VisibleForTesting
+  FlowController getFlowController() {
+    return flowController;
+  }
+
   /**
    * This class represent one logical Batch. It accumulates all the elements and their corresponding
    * future results for one batch.

gax/src/main/java/com/google/api/gax/batching/FlowController.java

@@ -30,6 +30,7 @@
 package com.google.api.gax.batching;
 
 import com.google.api.core.BetaApi;
+import com.google.api.core.InternalApi;
 import com.google.common.base.Preconditions;
 import javax.annotation.Nullable;
 
@@ -143,9 +144,10 @@ public enum LimitExceededBehavior {
   @Nullable private final Semaphore64 outstandingByteCount;
   @Nullable private final Long maxOutstandingElementCount;
   @Nullable private final Long maxOutstandingRequestBytes;
+  private final LimitExceededBehavior limitExceededBehavior;
 
   public FlowController(FlowControlSettings settings) {
-    boolean failOnLimits;
+    this.limitExceededBehavior = settings.getLimitExceededBehavior();
     switch (settings.getLimitExceededBehavior()) {
       case ThrowException:
       case Block:
@@ -216,4 +218,20 @@ public void release(long elements, long bytes) {
       outstandingByteCount.release(permitsToReturn);
     }
   }
+
+  LimitExceededBehavior getLimitExceededBehavior() {
+    return limitExceededBehavior;
+  }
+
+  @InternalApi("For internal use by google-cloud-java clients only")
+  @Nullable
+  public Long getMaxOutstandingElementCount() {
+    return maxOutstandingElementCount;
+  }
+
+  @InternalApi("For internal use by google-cloud-java clients only")
+  @Nullable
+  public Long getMaxOutstandingRequestBytes() {
+    return maxOutstandingRequestBytes;
+  }
 }