feat: introduce MaximumRequestCallbackWaitTimeExceededException (#2401)

agrawal-siddharth · web-flow · commit 0dbbfb80632e · 2024-02-06T21:54:11.000-08:00
exception
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java
@@ -844,11 +844,8 @@ private void appendLoop() {
   private void throwIfWaitCallbackTooLong(Instant timeToCheck) {
     Duration milliSinceLastCallback = Duration.between(timeToCheck, Instant.now());
     if (milliSinceLastCallback.compareTo(MAXIMUM_REQUEST_CALLBACK_WAIT_TIME) > 0) {
-      throw new RuntimeException(
-          String.format(
-              "Request has waited in inflight queue for %sms for writer %s, "
-                  + "which is over maximum wait time %s",
-              milliSinceLastCallback, writerId, MAXIMUM_REQUEST_CALLBACK_WAIT_TIME.toString()));
+      throw new Exceptions.MaximumRequestCallbackWaitTimeExceededException(
+          milliSinceLastCallback, writerId, MAXIMUM_REQUEST_CALLBACK_WAIT_TIME);
     }
   }
 
diff --git a/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/Exceptions.java b/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/Exceptions.java
@@ -22,6 +22,7 @@
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import io.grpc.protobuf.StatusProto;
+import java.time.Duration;
 import java.util.Map;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -416,5 +417,38 @@ public String getFieldName() {
     }
   }
 
+  /**
+   * The connection was shut down because a callback was not received within the maximum wait time.
+   */
+  public static class MaximumRequestCallbackWaitTimeExceededException extends RuntimeException {
+    private final Duration callbackWaitTime;
+    private final String writerId;
+    private final Duration callbackWaitTimeLimit;
+
+    public MaximumRequestCallbackWaitTimeExceededException(
+        Duration callbackWaitTime, String writerId, Duration callbackWaitTimeLimit) {
+      super(
+          String.format(
+              "Request has waited in inflight queue for %sms for writer %s, "
+                  + "which is over maximum wait time %s",
+              callbackWaitTime, writerId, callbackWaitTimeLimit.toString()));
+      this.callbackWaitTime = callbackWaitTime;
+      this.writerId = writerId;
+      this.callbackWaitTimeLimit = callbackWaitTimeLimit;
+    }
+
+    public Duration getCallbackWaitTime() {
+      return callbackWaitTime;
+    }
+
+    public String getWriterId() {
+      return writerId;
+    }
+
+    public Duration getCallbackWaitTimeLimit() {
+      return callbackWaitTimeLimit;
+    }
+  }
+
   private Exceptions() {}
 }
diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/StreamWriterTest.java
@@ -1235,6 +1235,8 @@ public void testThrowExceptionWhileWithinAppendLoop_MaxWaitTimeExceed() throws E
               () -> futures.get(finalI).get().getAppendResult().getOffset().getValue());
       if (i == 0) {
         assertThat(ex.getCause()).hasMessageThat().contains("Request has waited in inflight queue");
+        assertThat(ex.getCause())
+            .isInstanceOf(Exceptions.MaximumRequestCallbackWaitTimeExceededException.class);
       } else {
         assertThat(ex.getCause())
             .hasMessageThat()
