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

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

@@ -37,6 +37,8 @@
 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.LimitExceededBehavior;
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
@@ -55,6 +57,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,6 +90,7 @@ 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
@@ -102,15 +106,55 @@ 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 flowControllerToUse a {@link FlowController} for throttling requests. If it's not set,
+   *     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 flowControllerToUse) {
+
     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 (flowControllerToUse == null) {
+      flowControllerToUse = new FlowController(batchingSettings.getFlowControlSettings());
+    }
+    // If throttling is enabled, make sure flow control limits are greater or equal to batch sizes
+    // to avoid deadlocking
+    if (flowControllerToUse.getLimitExceededBehavior() != LimitExceededBehavior.Ignore) {
+      Preconditions.checkArgument(
+          flowControllerToUse.getMaxOutstandingElementCount() == null
+              || batchingSettings.getElementCountThreshold() == null
+              || flowControllerToUse.getMaxOutstandingElementCount()
+                  >= batchingSettings.getElementCountThreshold(),
+          "if throttling and batching on element count are enabled, FlowController#maxOutstandingElementCount must be greater or equal to elementCountThreshold");
+      Preconditions.checkArgument(
+          flowControllerToUse.getMaxOutstandingRequestBytes() == null
+              || batchingSettings.getRequestByteThreshold() == null
+              || flowControllerToUse.getMaxOutstandingRequestBytes()
+                  >= batchingSettings.getRequestByteThreshold(),
+          "if throttling and batching on request bytes are enabled, FlowController#maxOutstandingRequestBytes must be greater or equal to requestByteThreshold");
+    }
+    this.flowController = flowControllerToUse;
     currentOpenBatch = new Batch<>(prototype, batchingDescriptor, batchingSettings, batcherStats);
-
     if (batchingSettings.getDelayThreshold() != null) {
       long delay = batchingSettings.getDelayThreshold().toMillis();
       PushCurrentBatchRunnable<ElementT, ElementResultT, RequestT, ResponseT> runnable =
@@ -127,8 +171,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 makes it seemed 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 new RuntimeException(e);
+    }
 
+    SettableApiFuture<ElementResultT> result = SettableApiFuture.create();
     synchronized (elementLock) {
       currentOpenBatch.add(element, result);
     }
@@ -169,6 +234,7 @@ public void sendOutstanding() {
           @Override
           public void onSuccess(ResponseT response) {
             try {
+              flowController.release(accumulatedBatch.elementCounter, accumulatedBatch.byteCounter);
               accumulatedBatch.onBatchSuccess(response);
             } finally {
               onBatchCompletion();
@@ -178,6 +244,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 +291,10 @@ public void close() throws InterruptedException {
     }
   }
 
+  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/BatchingSettings.java

@@ -174,6 +174,21 @@ public BatchingSettings build() {
           settings.getDelayThreshold() == null
               || settings.getDelayThreshold().compareTo(Duration.ZERO) > 0,
           "delayThreshold must be either unset or positive");
+      if (settings.getFlowControlSettings().getLimitExceededBehavior()
+          != LimitExceededBehavior.Ignore) {
+        Preconditions.checkArgument(
+            settings.getFlowControlSettings().getMaxOutstandingElementCount() == null
+                || settings.getElementCountThreshold() == null
+                || settings.getFlowControlSettings().getMaxOutstandingElementCount()
+                    >= settings.getElementCountThreshold(),
+            "if throttling and batching on element count are enabled, FlowController#maxOutstandingElementCount must be greater or equal to elementCountThreshold");
+        Preconditions.checkArgument(
+            settings.getFlowControlSettings().getMaxOutstandingRequestBytes() == null
+                || settings.getRequestByteThreshold() == null
+                || settings.getFlowControlSettings().getMaxOutstandingRequestBytes()
+                    >= settings.getRequestByteThreshold(),
+            "if throttling and batching on request bytes are enabled, FlowController#maxOutstandingRequestBytes must be greater or equal to requestByteThreshold");
+      }
       return settings;
     }
   }

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;
+  }
 }

gax/src/test/java/com/google/api/gax/batching/BatcherImplTest.java

@@ -38,6 +38,7 @@
 import com.google.api.core.ApiFutures;
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.batching.BatcherImpl.BatcherReference;
+import com.google.api.gax.batching.FlowController.LimitExceededBehavior;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.UnaryCallable;
 import com.google.api.gax.rpc.testing.FakeBatchableApi.LabeledIntList;
@@ -711,6 +712,134 @@ public void run() {
     }
   }
 
+  @Test
+  public void testConstructors() throws InterruptedException {
+    BatcherImpl batcher1 =
+        new BatcherImpl<>(
+            SQUARER_BATCHING_DESC_V2,
+            callLabeledIntSquarer,
+            labeledIntList,
+            batchingSettings,
+            EXECUTOR);
+    try {
+      assertThat(batcher1.getFlowController()).isNotNull();
+      assertThat(batcher1.getFlowController().getLimitExceededBehavior())
+          .isEqualTo(batchingSettings.getFlowControlSettings().getLimitExceededBehavior());
+      assertThat(batcher1.getFlowController().getMaxOutstandingElementCount())
+          .isEqualTo(batchingSettings.getFlowControlSettings().getMaxOutstandingElementCount());
+      assertThat(batcher1.getFlowController().getMaxOutstandingRequestBytes())
+          .isEqualTo(batchingSettings.getFlowControlSettings().getMaxOutstandingRequestBytes());
+    } finally {
+      batcher1.close();
+    }
+
+    FlowController flowController =
+        new FlowController(
+            FlowControlSettings.newBuilder()
+                .setLimitExceededBehavior(LimitExceededBehavior.ThrowException.ThrowException)
+                .setMaxOutstandingRequestBytes(6000L)
+                .build());
+    BatcherImpl batcher2 =
+        new BatcherImpl<>(
+            SQUARER_BATCHING_DESC_V2,
+            callLabeledIntSquarer,
+            labeledIntList,
+            batchingSettings,
+            EXECUTOR,
+            flowController);
+    try {
+      assertThat(batcher2.getFlowController()).isSameInstanceAs(flowController);
+    } finally {
+      batcher2.close();
+    }
+  }
+
+  @Test
+  public void testThrottlingBlocking() throws Exception {
+    BatchingSettings settings =
+        BatchingSettings.newBuilder()
+            .setElementCountThreshold(1L)
+            .setRequestByteThreshold(1L)
+            .build();
+    FlowController flowController =
+        new FlowController(
+            FlowControlSettings.newBuilder()
+                .setLimitExceededBehavior(LimitExceededBehavior.Block)
+                .setMaxOutstandingElementCount(1L)
+                .build());
+    final Batcher<Integer, Integer> batcher =
+        new BatcherImpl<>(
+            SQUARER_BATCHING_DESC_V2,
+            callLabeledIntSquarer,
+            labeledIntList,
+            settings,
+            EXECUTOR,
+            flowController);
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    try {
+      flowController.reserve(1, 1);
+      Future future =
+          executor.submit(
+              new Runnable() {
+                @Override
+                public void run() {
+                  batcher.add(1);
+                }
+              });
+      try {
+        future.get(100, TimeUnit.MILLISECONDS);
+        assertWithMessage("adding elements to batcher should be blocked by FlowControlled").fail();
+      } catch (TimeoutException e) {
+        // expected
+      }
+      flowController.release(1, 1);
+      try {
+        future.get(100, TimeUnit.MILLISECONDS);
+      } catch (TimeoutException e) {
+        assertWithMessage("adding elements to batcher should not be blocked").fail();
+      }
+    } finally {
+      batcher.close();
+      executor.shutdownNow();
+    }
+  }
+
+  @Test
+  public void testThrottlingNonBlocking() throws Exception {
+    BatchingSettings settings =
+        BatchingSettings.newBuilder()
+            .setElementCountThreshold(1L)
+            .setRequestByteThreshold(1L)
+            .build();
+    FlowController flowController =
+        new FlowController(
+            FlowControlSettings.newBuilder()
+                .setLimitExceededBehavior(LimitExceededBehavior.ThrowException)
+                .setMaxOutstandingElementCount(1L)
+                .build());
+    final Batcher<Integer, Integer> batcher =
+        new BatcherImpl<>(
+            SQUARER_BATCHING_DESC_V2,
+            callLabeledIntSquarer,
+            labeledIntList,
+            settings,
+            EXECUTOR,
+            flowController);
+    try {
+      flowController.reserve(1, 1);
+      try {
+        batcher.add(1);
+        assertWithMessage("Should throw exception because it exceeded FlowController limit").fail();
+      } catch (Exception e) {
+        assertThat(e.getMessage()).contains("The maximum number of batch elements");
+      }
+      flowController.release(1, 1);
+      batcher.add(1);
+    } finally {
+      batcher.close();
+    }
+  }
+
   private void testElementTriggers(BatchingSettings settings) throws Exception {
     underTest =
         new BatcherImpl<>(