fix: Add a throwException behavior when the StreamWriter inflight queue is full

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/JsonStreamWriter.java

@@ -78,6 +78,8 @@ private JsonStreamWriter(Builder builder)
     this.protoSchema = ProtoSchemaConverter.convert(this.descriptor);
     this.totalMessageSize = protoSchema.getSerializedSize();
     streamWriterBuilder.setWriterSchema(protoSchema);
+    streamWriterBuilder.setLimitExceededBehavior(
+        builder.flowControlSettings.getLimitExceededBehavior());
     setStreamWriterSettings(
         builder.channelProvider,
         builder.credentialsProvider,
@@ -214,6 +216,10 @@ private void setStreamWriterSettings(
         streamWriterBuilder.setMaxInflightRequests(
             flowControlSettings.getMaxOutstandingElementCount());
       }
+      if (flowControlSettings.getLimitExceededBehavior() != null) {
+        streamWriterBuilder.setLimitExceededBehavior(
+            flowControlSettings.getLimitExceededBehavior());
+      }
     }
   }
 
@@ -335,7 +341,6 @@ public Builder setCredentialsProvider(CredentialsProvider credentialsProvider) {
      * @return Builder
      */
     public Builder setFlowControlSettings(FlowControlSettings flowControlSettings) {
-      Preconditions.checkNotNull(flowControlSettings, "FlowControlSettings is null.");
       this.flowControlSettings =
           Preconditions.checkNotNull(flowControlSettings, "FlowControlSettings is null.");
       return this;

google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/StreamWriter.java

@@ -17,6 +17,7 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.SettableApiFuture;
+import com.google.api.gax.batching.FlowController;
 import com.google.api.gax.core.CredentialsProvider;
 import com.google.api.gax.rpc.FixedHeaderProvider;
 import com.google.api.gax.rpc.TransportChannelProvider;
@@ -74,6 +75,11 @@ public class StreamWriter implements AutoCloseable {
    */
   private final long maxInflightBytes;
 
+  /*
+   * Behavior when inflight queue is exceeded. Only supports Block or Throw, default is Block.
+   */
+  private final FlowController.LimitExceededBehavior limitExceededBehavior;
+
   /*
    * TraceId for debugging purpose.
    */
@@ -190,6 +196,7 @@ private StreamWriter(Builder builder) throws IOException {
     this.writerSchema = builder.writerSchema;
     this.maxInflightRequests = builder.maxInflightRequest;
     this.maxInflightBytes = builder.maxInflightBytes;
+    this.limitExceededBehavior = builder.limitExceededBehavior;
     this.traceId = builder.traceId;
     this.waitingRequestQueue = new LinkedList<AppendRequestAndResponse>();
     this.inflightRequestQueue = new LinkedList<AppendRequestAndResponse>();
@@ -332,18 +339,29 @@ private void maybeWaitForInflightQuota() {
     long start_time = System.currentTimeMillis();
     while (this.inflightRequests >= this.maxInflightRequests
         || this.inflightBytes >= this.maxInflightBytes) {
-      try {
-        inflightReduced.await(100, TimeUnit.MILLISECONDS);
-      } catch (InterruptedException e) {
-        log.warning(
-            "Interrupted while waiting for inflight quota. Stream: "
-                + streamName
-                + " Error: "
-                + e.toString());
+      if (this.limitExceededBehavior == FlowController.LimitExceededBehavior.ThrowException) {
+        throw new StatusRuntimeException(
+            Status.fromCode(Code.RESOURCE_EXHAUSTED)
+                .withDescription(
+                    "Exceeds client side inflight buffer, consider add more buffer or open more connections."));
+      } else if (this.limitExceededBehavior == FlowController.LimitExceededBehavior.Ignore) {
         throw new StatusRuntimeException(
-            Status.fromCode(Code.CANCELLED)
-                .withCause(e)
-                .withDescription("Interrupted while waiting for quota."));
+            Status.fromCode(Code.INVALID_ARGUMENT)
+                .withDescription("LimitExceededBehavior.Ignore is not supported on StreamWriter."));
+      } else {
+        try {
+          inflightReduced.await(100, TimeUnit.MILLISECONDS);
+        } catch (InterruptedException e) {
+          log.warning(
+              "Interrupted while waiting for inflight quota. Stream: "
+                  + streamName
+                  + " Error: "
+                  + e.toString());
+          throw new StatusRuntimeException(
+              Status.fromCode(Code.CANCELLED)
+                  .withCause(e)
+                  .withDescription("Interrupted while waiting for quota."));
+        }
       }
     }
     inflightWaitSec.set((System.currentTimeMillis() - start_time) / 1000);
@@ -714,6 +732,9 @@ public static final class Builder {
     private CredentialsProvider credentialsProvider =
         BigQueryWriteSettings.defaultCredentialsProviderBuilder().build();
 
+    private FlowController.LimitExceededBehavior limitExceededBehavior =
+        FlowController.LimitExceededBehavior.Block;
+
     private String traceId = null;
 
     private TableSchema updatedTableSchema = null;
@@ -784,6 +805,18 @@ public Builder setTraceId(String traceId) {
       return this;
     }
 
+    /**
+     * Sets the limit exceeded behavior.
+     *
+     * @param limitExceededBehavior
+     * @return
+     */
+    public Builder setLimitExceededBehavior(
+        FlowController.LimitExceededBehavior limitExceededBehavior) {
+      this.limitExceededBehavior = limitExceededBehavior;
+      return this;
+    }
+
     /** Builds the {@code StreamWriterV2}. */
     public StreamWriter build() throws IOException {
       return new StreamWriter(this);

google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java

@@ -15,12 +15,10 @@
  */
 package com.google.cloud.bigquery.storage.v1;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.*;
 
 import com.google.api.core.ApiFuture;
+import com.google.api.gax.batching.FlowController;
 import com.google.api.gax.core.NoCredentialsProvider;
 import com.google.api.gax.grpc.testing.MockGrpcService;
 import com.google.api.gax.grpc.testing.MockServiceHelper;
@@ -543,6 +541,31 @@ public void testAppendsWithTinyMaxInflightBytes() throws Exception {
     writer.close();
   }
 
+  @Test
+  public void testAppendsWithTinyMaxInflightBytesThrow() throws Exception {
+    StreamWriter writer =
+        StreamWriter.newBuilder(TEST_STREAM, client)
+            .setWriterSchema(createProtoSchema())
+            .setMaxInflightBytes(1)
+            .setLimitExceededBehavior(FlowController.LimitExceededBehavior.ThrowException)
+            .build();
+    // Server will sleep 100ms before every response.
+    testBigQueryWrite.setResponseSleep(Duration.ofMillis(100));
+    long appendCount = 10;
+    for (int i = 0; i < appendCount; i++) {
+      testBigQueryWrite.addResponse(createAppendResponse(i));
+    }
+    try {
+      writer.append(createProtoRows(new String[] {String.valueOf(10)}), -1);
+      fail("Expect exception to throw");
+    } catch (StatusRuntimeException ex) {
+      assertEquals(
+          "RESOURCE_EXHAUSTED: Exceeds client side inflight buffer, consider add more buffer or open more connections.",
+          ex.getMessage());
+    }
+    writer.close();
+  }
+
   @Test
   public void testMessageTooLarge() throws Exception {
     StreamWriter writer = getTestStreamWriter();