fix: update grpc WriteObject response handling to provide context when a failure happens (#2532)

Author: BenWhitehead

google-cloud-storage/src/main/java/com/google/cloud/storage/AsyncStorageTaskException.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.storage;
+
+/**
+ * This exception is used to preserve the caller's stacktrace when invoking an async task in a sync
+ * context. It will be added as a suppressed exception when propagating the async exception. This
+ * allows callers to catch ApiException thrown in an async operation, while still maintaining the
+ * call site.
+ */
+public final class AsyncStorageTaskException extends RuntimeException {
+  // mimic of com.google.api.gax.rpc.AsyncTaskException which doesn't have a public constructor
+  // if that class is ever made public, make this class extend it
+  AsyncStorageTaskException() {
+    super("Asynchronous task failed");
+  }
+}

google-cloud-storage/src/main/java/com/google/cloud/storage/GapicUnbufferedChunkedResumableWritableByteChannel.java

@@ -21,6 +21,7 @@
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.api.gax.retrying.ResultRetryAlgorithm;
+import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.ApiStreamObserver;
 import com.google.api.gax.rpc.ClientStreamingCallable;
 import com.google.api.gax.rpc.OutOfRangeException;
@@ -208,14 +209,19 @@ private void flush(
         deps,
         alg,
         () -> {
-          Observer observer = new Observer(content, finalizing);
+          Observer observer = new Observer(content, finalizing, segments, internalContext);
           ApiStreamObserver<WriteObjectRequest> write = callable.clientStreamingCall(observer);
 
           for (WriteObjectRequest message : segments) {
             write.onNext(message);
           }
           write.onCompleted();
-          observer.await();
+          try {
+            observer.await();
+          } catch (Throwable t) {
+            t.addSuppressed(new AsyncStorageTaskException());
+            throw t;
+          }
           return null;
         },
         Decoder.identity());
@@ -230,13 +236,21 @@ class Observer implements ApiStreamObserver<WriteObjectResponse> {
 
     private final RewindableContent content;
     private final boolean finalizing;
+    private final List<WriteObjectRequest> segments;
+    private final GrpcCallContext context;
 
     private final SettableApiFuture<Void> invocationHandle;
     private volatile WriteObjectResponse last;
 
-    Observer(@Nullable RewindableContent content, boolean finalizing) {
+    Observer(
+        @Nullable RewindableContent content,
+        boolean finalizing,
+        @NonNull List<WriteObjectRequest> segments,
+        GrpcCallContext context) {
       this.content = content;
       this.finalizing = finalizing;
+      this.segments = segments;
+      this.context = context;
       this.invocationHandle = SettableApiFuture.create();
     }
 
@@ -250,10 +264,20 @@ public void onError(Throwable t) {
       if (t instanceof OutOfRangeException) {
         OutOfRangeException oore = (OutOfRangeException) t;
         open = false;
-        invocationHandle.setException(
-            ResumableSessionFailureScenario.SCENARIO_5.toStorageException());
-      } else {
-        invocationHandle.setException(t);
+        StorageException storageException =
+            ResumableSessionFailureScenario.SCENARIO_5.toStorageException(
+                segments, null, context, oore);
+        invocationHandle.setException(storageException);
+      } else if (t instanceof ApiException) {
+        // use StorageExceptions logic to translate from ApiException to our status codes ensuring
+        // things fall in line with our retry handlers.
+        // This is suboptimal, as it will initialize a second exception, however this is the
+        // unusual case, and it should not cause a significant overhead given its rarity.
+        StorageException tmp = StorageException.asStorageException((ApiException) t);
+        StorageException storageException =
+            ResumableSessionFailureScenario.toStorageException(
+                tmp.getCode(), tmp.getMessage(), tmp.getReason(), segments, null, context, t);
+        invocationHandle.setException(storageException);
       }
     }
 
@@ -276,7 +300,8 @@ public void onCompleted() {
             writeCtx.getTotalSentBytes().set(persistedSize);
             writeCtx.getConfirmedBytes().set(persistedSize);
           } else {
-            throw ResumableSessionFailureScenario.SCENARIO_7.toStorageException();
+            throw ResumableSessionFailureScenario.SCENARIO_7.toStorageException(
+                segments, last, context, null);
           }
         } else if (finalizing && last.hasResource()) {
           long totalSentBytes = writeCtx.getTotalSentBytes().get();
@@ -285,22 +310,28 @@ public void onCompleted() {
             writeCtx.getConfirmedBytes().set(finalSize);
             resultFuture.set(last);
           } else if (finalSize < totalSentBytes) {
-            throw ResumableSessionFailureScenario.SCENARIO_4_1.toStorageException();
+            throw ResumableSessionFailureScenario.SCENARIO_4_1.toStorageException(
+                segments, last, context, null);
           } else {
-            throw ResumableSessionFailureScenario.SCENARIO_4_2.toStorageException();
+            throw ResumableSessionFailureScenario.SCENARIO_4_2.toStorageException(
+                segments, last, context, null);
           }
         } else if (!finalizing && last.hasResource()) {
-          throw ResumableSessionFailureScenario.SCENARIO_1.toStorageException();
+          throw ResumableSessionFailureScenario.SCENARIO_1.toStorageException(
+              segments, last, context, null);
         } else if (finalizing && last.hasPersistedSize()) {
           long totalSentBytes = writeCtx.getTotalSentBytes().get();
           long persistedSize = last.getPersistedSize();
           if (persistedSize < totalSentBytes) {
-            throw ResumableSessionFailureScenario.SCENARIO_3.toStorageException();
+            throw ResumableSessionFailureScenario.SCENARIO_3.toStorageException(
+                segments, last, context, null);
           } else {
-            throw ResumableSessionFailureScenario.SCENARIO_2.toStorageException();
+            throw ResumableSessionFailureScenario.SCENARIO_2.toStorageException(
+                segments, last, context, null);
           }
         } else {
-          throw ResumableSessionFailureScenario.SCENARIO_0.toStorageException();
+          throw ResumableSessionFailureScenario.SCENARIO_0.toStorageException(
+              segments, last, context, null);
         }
       } catch (Throwable se) {
         open = false;

google-cloud-storage/src/main/java/com/google/cloud/storage/ResumableSessionFailureScenario.java

@@ -16,19 +16,31 @@
 
 package com.google.cloud.storage;
 
+import static com.google.cloud.storage.Utils.ifNonNull;
+
 import com.google.api.client.http.HttpHeaders;
 import com.google.api.client.http.HttpResponse;
 import com.google.api.client.http.HttpResponseException;
+import com.google.api.gax.grpc.GrpcCallContext;
+import com.google.api.gax.rpc.ApiException;
 import com.google.cloud.BaseServiceException;
 import com.google.cloud.storage.StorageException.IOExceptionCallable;
 import com.google.common.io.CharStreams;
+import com.google.protobuf.MessageOrBuilder;
+import com.google.storage.v2.ChecksummedData;
+import com.google.storage.v2.ObjectChecksums;
+import com.google.storage.v2.WriteObjectRequest;
+import com.google.storage.v2.WriteObjectResponse;
+import io.grpc.StatusRuntimeException;
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.function.Consumer;
 import java.util.function.Predicate;
 import javax.annotation.ParametersAreNonnullByDefault;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 @ParametersAreNonnullByDefault
@@ -69,6 +81,10 @@ enum ResumableSessionFailureScenario {
   private static final String PREFIX_I = "\t|< ";
   private static final String PREFIX_O = "\t|> ";
   private static final String PREFIX_X = "\t|  ";
+  // define some constants for tab widths that are more compressed that the literals
+  private static final String T1 = "\t";
+  private static final String T2 = "\t\t";
+  private static final String T3 = "\t\t\t";
 
   private static final Predicate<String> includedHeaders =
       matches("Content-Length")
@@ -78,6 +94,7 @@ enum ResumableSessionFailureScenario {
           .or(matches("Range"))
           .or(startsWith("X-Goog-Stored-"))
           .or(matches("X-Goog-GCS-Idempotency-Token"))
+          .or(matches("X-Goog-request-params"))
           .or(matches("X-GUploader-UploadID"));
 
   private static final Predicate<Map.Entry<String, ?>> includeHeader =
@@ -116,8 +133,12 @@ StorageException toStorageException(
     return toStorageException(code, message, reason, uploadId, resp, cause, contentCallable);
   }
 
-  StorageException toStorageException() {
-    return new StorageException(code, message, reason, null);
+  StorageException toStorageException(
+      @NonNull List<@NonNull WriteObjectRequest> reqs,
+      @Nullable WriteObjectResponse resp,
+      @NonNull GrpcCallContext context,
+      @Nullable Throwable cause) {
+    return toStorageException(code, message, reason, reqs, resp, context, cause);
   }
 
   static StorageException toStorageException(
@@ -136,6 +157,102 @@ static StorageException toStorageException(
     return se;
   }
 
+  static StorageException toStorageException(
+      int code,
+      String message,
+      @Nullable String reason,
+      @NonNull List<@NonNull WriteObjectRequest> reqs,
+      @Nullable WriteObjectResponse resp,
+      @NonNull GrpcCallContext context,
+      @Nullable Throwable cause) {
+    final StringBuilder sb = new StringBuilder();
+    sb.append(message);
+    // request context
+    Map<String, List<String>> extraHeaders = context.getExtraHeaders();
+    recordHeadersTo(extraHeaders, PREFIX_O, sb);
+    int length = reqs.size();
+    for (int i = 0; i < length; i++) {
+      if (i == 0) {
+        sb.append("\n").append(PREFIX_O).append("[");
+      } else {
+        sb.append(",");
+      }
+      WriteObjectRequest req = reqs.get(i);
+      sb.append("\n").append(PREFIX_O).append(T1).append(req.getClass().getName()).append("{");
+      if (req.hasUploadId()) {
+        sb.append("\n").append(PREFIX_O).append(T2).append("upload_id: ").append(req.getUploadId());
+      }
+      long writeOffset = req.getWriteOffset();
+      if (req.hasChecksummedData()) {
+        ChecksummedData checksummedData = req.getChecksummedData();
+        sb.append("\n").append(PREFIX_O).append(T2);
+        sb.append(
+            String.format(
+                "checksummed_data: {range: [%d:%d]",
+                writeOffset, writeOffset + checksummedData.getContent().size()));
+        if (checksummedData.hasCrc32C()) {
+          sb.append(", crc32c: ").append(checksummedData.getCrc32C());
+        }
+        sb.append("}");
+      } else {
+        sb.append("\n").append(PREFIX_O).append(T2).append("write_offset: ").append(writeOffset);
+      }
+      if (req.getFinishWrite()) {
+        sb.append("\n").append(PREFIX_O).append(T2).append("finish_write: true");
+      }
+      if (req.hasObjectChecksums()) {
+        ObjectChecksums objectChecksums = req.getObjectChecksums();
+        sb.append("\n").append(PREFIX_O).append(T2).append("object_checksums: ").append("{");
+        fmt(objectChecksums, PREFIX_O, T3, sb);
+        sb.append("\n").append(PREFIX_O).append(T2).append("}");
+      }
+      sb.append("\n").append(PREFIX_O).append("\t}");
+      if (i == length - 1) {
+        sb.append("\n").append(PREFIX_O).append("]");
+      }
+    }
+
+    sb.append("\n").append(PREFIX_X);
+
+    // response context
+    if (resp != null) {
+      sb.append("\n").append(PREFIX_I).append(resp.getClass().getName()).append("{");
+      fmt(resp, PREFIX_I, T1, sb);
+      sb.append("\n").append(PREFIX_I).append("}");
+      sb.append("\n").append(PREFIX_X);
+    }
+
+    if (cause != null) {
+      if (cause instanceof ApiException) {
+        ApiException apiException = (ApiException) cause;
+        Throwable cause1 = apiException.getCause();
+        if (cause1 instanceof StatusRuntimeException) {
+          StatusRuntimeException statusRuntimeException = (StatusRuntimeException) cause1;
+          sb.append("\n").append(PREFIX_I).append(statusRuntimeException.getStatus());
+          ifNonNull(
+              statusRuntimeException.getTrailers(),
+              t -> sb.append("\n").append(PREFIX_I).append(t));
+        } else {
+          sb.append("\n")
+              .append(PREFIX_I)
+              .append("code: ")
+              .append(apiException.getStatusCode().toString());
+          ifNonNull(
+              apiException.getReason(),
+              r -> sb.append("\n").append(PREFIX_I).append("reason: ").append(r));
+          ifNonNull(
+              apiException.getDomain(),
+              d -> sb.append("\n").append(PREFIX_I).append("domain: ").append(d));
+          ifNonNull(
+              apiException.getErrorDetails(),
+              e -> sb.append("\n").append(PREFIX_I).append("errorDetails: ").append(e));
+        }
+        sb.append("\n").append(PREFIX_X);
+      }
+    }
+    return new StorageException(code, sb.toString(), reason, cause);
+  }
+
   static StorageException toStorageException(
       int overrideCode,
       String message,
@@ -213,14 +330,21 @@ private static Predicate<String> startsWith(String prefix) {
   }
 
   private static void recordHeaderTo(HttpHeaders h, String prefix, StringBuilder sb) {
-    h.entrySet().stream()
-        .filter(includeHeader)
-        .forEach(
-            e -> {
-              String key = e.getKey();
-              String value = headerValueToString(e.getValue());
-              sb.append("\n").append(prefix).append(key).append(": ").append(value);
-            });
+    h.entrySet().stream().filter(includeHeader).forEach(writeHeaderValue(prefix, sb));
+  }
+
+  private static void recordHeadersTo(
+      Map<String, List<String>> headers, String prefix, StringBuilder sb) {
+    headers.entrySet().stream().filter(includeHeader).forEach(writeHeaderValue(prefix, sb));
+  }
+
+  private static <V> Consumer<Map.Entry<String, V>> writeHeaderValue(
+      String prefix, StringBuilder sb) {
+    return e -> {
+      String key = e.getKey();
+      String value = headerValueToString(e.getValue());
+      sb.append("\n").append(prefix).append(key).append(": ").append(value);
+    };
   }
 
   private static String headerValueToString(Object o) {
@@ -233,4 +357,18 @@ private static String headerValueToString(Object o) {
 
     return o.toString();
   }
+
+  private static void fmt(
+      MessageOrBuilder msg,
+      @SuppressWarnings("SameParameterValue") String prefix,
+      String indentation,
+      StringBuilder sb) {
+    String string = msg.toString();
+    // drop the final new line before prefixing
+    string = string.replaceAll("\n$", "");
+    sb.append("\n")
+        .append(prefix)
+        .append(indentation)
+        .append(string.replaceAll("\r?\n", "\n" + prefix + indentation));
+  }
 }