fix: update BlobReadChannelV2 handling to correctly restart for decompressed object (#1867)

BenWhitehead · web-flow · commit 93e8ed484224 · 2023-01-20T16:23:31.000-05:00
* fix: update BlobReadChannelV2 handling to correctly restart for decompressed object

When downloading bytes from gcs and decompressing them, a restart of the stream needs to pickup from the offset within the compressed bytes not the decompressed bytes.

Prior to this change http-client was automatically applying gzip decompression to the stream it returns to us thereby causing our tracking to be off.

This change updates our interaction with http client to always request the raw bytes without any transform applied to them, we then at a higher level can handle whether gzip decompression needs to be plumbed in.

Fix up a couple channel closed state subtleties when buffering is used with decompression.

Add a new test leveraging the testbench which forces a broken stream on some compressed object bytes. _NOTE_ This test depends on the next release of testbench and will be failing until we get that release.

* deps(test): update testbench to v0.33.0
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/ApiaryUnbufferedReadableByteChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/ApiaryUnbufferedReadableByteChannel.java
@@ -17,6 +17,7 @@
 package com.google.cloud.storage;
 
 import static com.google.cloud.storage.Utils.ifNonNull;
+import static java.util.Objects.requireNonNull;
 
 import com.google.api.client.http.HttpHeaders;
 import com.google.api.client.http.HttpResponse;
@@ -49,6 +50,7 @@
 import java.nio.channels.ReadableByteChannel;
 import java.nio.channels.ScatteringByteChannel;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.function.Function;
 import javax.annotation.concurrent.Immutable;
@@ -82,13 +84,9 @@ class ApiaryUnbufferedReadableByteChannel implements UnbufferedReadableByteChann
     this.options = options;
     this.resultRetryAlgorithm = resultRetryAlgorithm;
     this.open = true;
-    this.position =
-        apiaryReadRequest.getByteRangeSpec() != null
-            ? apiaryReadRequest.getByteRangeSpec().beginOffset()
-            : 0;
+    this.position = apiaryReadRequest.getByteRangeSpec().beginOffset();
   }
 
-  @SuppressWarnings("UnnecessaryContinue")
   @Override
   public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
     do {
@@ -113,12 +111,10 @@ public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
           // if our retry algorithm COULD allow a retry, continue the loop and allow trying to
           // open the stream again.
           sbc = null;
-          continue;
         } else if (t instanceof IOException) {
           IOException ioE = (IOException) t;
           if (resultRetryAlgorithm.shouldRetry(StorageException.translate(ioE), null)) {
             sbc = null;
-            continue;
           } else {
             throw ioE;
           }
@@ -148,11 +144,8 @@ private void setXGoogGeneration(long xGoogGeneration) {
 
   private ScatteringByteChannel open() {
     try {
-      Boolean b =
-          (Boolean) apiaryReadRequest.options.get(StorageRpc.Option.RETURN_RAW_INPUT_STREAM);
-      boolean returnRawInputStream = b != null ? b : true;
       ApiaryReadRequest request = apiaryReadRequest.withNewBeginOffset(position);
-      Get get = createGetRequest(request, storage.objects(), xGoogGeneration, returnRawInputStream);
+      Get get = createGetRequest(request, storage.objects(), xGoogGeneration);
 
       HttpResponse media = get.executeMedia();
       InputStream content = media.getContent();
@@ -215,10 +208,7 @@ private ScatteringByteChannel open() {
 
   @VisibleForTesting
   static Get createGetRequest(
-      ApiaryReadRequest apiaryReadRequest,
-      Objects objects,
-      Long xGoogGeneration,
-      boolean returnRawInputStream)
+      ApiaryReadRequest apiaryReadRequest, Objects objects, Long xGoogGeneration)
       throws IOException {
     StorageObject from = apiaryReadRequest.getObject();
     Map<StorageRpc.Option, ?> options = apiaryReadRequest.getOptions();
@@ -262,7 +252,9 @@ static Get createGetRequest(
               base64.encode(hashFunction.hashBytes(base64.decode(key)).asBytes()));
         });
 
-    get.setReturnRawInputStream(returnRawInputStream);
+    // gzip handling is performed upstream of here. Ensure we always get the raw input stream from
+    // the request
+    get.setReturnRawInputStream(true);
     String range = apiaryReadRequest.getByteRangeSpec().getHttpRangeHeader();
     if (range != null) {
       get.getRequestHeaders().setRange(range);
@@ -288,7 +280,7 @@ private static String getHeaderValue(@NonNull HttpHeaders headers, @NonNull Stri
       if (list.isEmpty()) {
         return null;
       } else {
-        return list.get(0);
+        return list.get(0).trim().toLowerCase(Locale.ENGLISH);
       }
     } else if (o instanceof String) {
       return (String) o;
@@ -303,27 +295,32 @@ private static String getHeaderValue(@NonNull HttpHeaders headers, @NonNull Stri
   static final class ApiaryReadRequest implements Serializable {
     private static final long serialVersionUID = -4059435314115374448L;
     private static final Gson gson = new Gson();
-    private transient StorageObject object;
-    private final Map<StorageRpc.Option, ?> options;
-    private final ByteRangeSpec byteRangeSpec;
+    @NonNull private transient StorageObject object;
+    @NonNull private final Map<StorageRpc.Option, ?> options;
+    @NonNull private final ByteRangeSpec byteRangeSpec;
 
     private volatile String objectJson;
 
     ApiaryReadRequest(
-        StorageObject object, Map<StorageRpc.Option, ?> options, ByteRangeSpec byteRangeSpec) {
-      this.object = object;
-      this.options = options;
-      this.byteRangeSpec = byteRangeSpec;
+        @NonNull StorageObject object,
+        @NonNull Map<StorageRpc.Option, ?> options,
+        @NonNull ByteRangeSpec byteRangeSpec) {
+      this.object = requireNonNull(object, "object must be non null");
+      this.options = requireNonNull(options, "options must be non null");
+      this.byteRangeSpec = requireNonNull(byteRangeSpec, "byteRangeSpec must be non null");
     }
 
+    @NonNull
     StorageObject getObject() {
       return object;
     }
 
+    @NonNull
     Map<StorageRpc.Option, ?> getOptions() {
       return options;
     }
 
+    @NonNull
     ByteRangeSpec getByteRangeSpec() {
       return byteRangeSpec;
     }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/BaseStorageReadChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/BaseStorageReadChannel.java
@@ -120,7 +120,7 @@ public final synchronized int read(ByteBuffer dst) throws IOException {
   }
 
   @Override
-  public ApiFuture<BlobInfo> getObject() {
+  public final ApiFuture<BlobInfo> getObject() {
     return ApiFutures.transform(result, objectDecoder::decode, MoreExecutors.directExecutor());
   }
 
@@ -136,7 +136,7 @@ protected final int getChunkSize() {
   }
 
   @Nullable
-  protected T getResolvedObject() {
+  protected final T getResolvedObject() {
     if (result.isDone()) {
       return StorageException.wrapFutureGet(result);
     } else {
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobReadChannelV2.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobReadChannelV2.java
@@ -32,6 +32,7 @@ final class BlobReadChannelV2 extends BaseStorageReadChannel<StorageObject> {
   private final StorageObject storageObject;
   private final Map<StorageRpc.Option, ?> opts;
   private final BlobReadChannelContext blobReadChannelContext;
+  private final boolean autoGzipDecompression;
 
   BlobReadChannelV2(
       StorageObject storageObject,
@@ -41,10 +42,15 @@ final class BlobReadChannelV2 extends BaseStorageReadChannel<StorageObject> {
     this.storageObject = storageObject;
     this.opts = opts;
     this.blobReadChannelContext = blobReadChannelContext;
+    this.autoGzipDecompression =
+        // RETURN_RAW_INPUT_STREAM means do not add GZIPInputStream to the pipeline. Meaning, if
+        // RETURN_RAW_INPUT_STREAM is false, automatically attempt to decompress if Content-Encoding
+        // gzip.
+        Boolean.FALSE.equals(opts.get(StorageRpc.Option.RETURN_RAW_INPUT_STREAM));
   }
 
   @Override
-  public RestorableState<ReadChannel> capture() {
+  public synchronized RestorableState<ReadChannel> capture() {
     ApiaryReadRequest apiaryReadRequest = getApiaryReadRequest();
     return new BlobReadChannelV2State(
         apiaryReadRequest, blobReadChannelContext.getStorageOptions(), getChunkSize());
@@ -56,6 +62,7 @@ protected LazyReadChannel<StorageObject> newLazyReadChannel() {
             ResumableMedia.http()
                 .read()
                 .byteChannel(blobReadChannelContext)
+                .setAutoGzipDecompression(autoGzipDecompression)
                 .buffered(getBufferHandle())
                 .setApiaryReadRequest(getApiaryReadRequest())
                 .build());
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedReadableByteChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/DefaultBufferedReadableByteChannel.java
@@ -41,7 +41,7 @@ public int read(ByteBuffer dst) throws IOException {
     if (retEOF) {
       retEOF = false;
       return -1;
-    } else if (!channel.isOpen()) {
+    } else if (!enqueuedBytes() && !channel.isOpen()) {
       throw new ClosedChannelException();
     }
 
@@ -133,7 +133,7 @@ public int read(ByteBuffer dst) throws IOException {
 
   @Override
   public boolean isOpen() {
-    return !retEOF && channel.isOpen();
+    return enqueuedBytes() || (!retEOF && channel.isOpen());
   }
 
   @Override
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicDownloadSessionBuilder.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicDownloadSessionBuilder.java
@@ -23,6 +23,7 @@
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.cloud.storage.BufferedReadableByteChannelSession.BufferedReadableByteChannel;
 import com.google.cloud.storage.UnbufferedReadableByteChannelSession.UnbufferedReadableByteChannel;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.google.storage.v2.Object;
 import com.google.storage.v2.ReadObjectRequest;
 import com.google.storage.v2.ReadObjectResponse;
@@ -99,7 +100,9 @@ public UnbufferedReadableByteChannelSessionBuilder unbuffered() {
       return (object, resultFuture) -> {
         if (autoGzipDecompression) {
           return new GzipReadableByteChannel(
-              new GapicUnbufferedReadableByteChannel(resultFuture, read, object, hasher));
+              new GapicUnbufferedReadableByteChannel(resultFuture, read, object, hasher),
+              ApiFutures.transform(
+                  resultFuture, Object::getContentEncoding, MoreExecutors.directExecutor()));
         } else {
           return new GapicUnbufferedReadableByteChannel(resultFuture, read, object, hasher);
         }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/GzipReadableByteChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/GzipReadableByteChannel.java
@@ -16,9 +16,13 @@
 
 package com.google.cloud.storage;
 
+import com.google.api.core.ApiFuture;
 import com.google.cloud.storage.UnbufferedReadableByteChannelSession.UnbufferedReadableByteChannel;
-import com.google.storage.v2.Object;
-import java.io.*;
+import java.io.ByteArrayInputStream;
+import java.io.FilterInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.SequenceInputStream;
 import java.nio.ByteBuffer;
 import java.nio.channels.Channels;
 import java.nio.channels.ReadableByteChannel;
@@ -27,13 +31,15 @@
 import java.util.zip.GZIPInputStream;
 
 final class GzipReadableByteChannel implements UnbufferedReadableByteChannel {
-  private final GapicUnbufferedReadableByteChannel source;
+  private final UnbufferedReadableByteChannel source;
+  private final ApiFuture<String> contentEncoding;
 
   private boolean retEOF = false;
   private ScatteringByteChannel delegate;
 
-  GzipReadableByteChannel(GapicUnbufferedReadableByteChannel source) {
+  GzipReadableByteChannel(UnbufferedReadableByteChannel source, ApiFuture<String> contentEncoding) {
     this.source = source;
+    this.contentEncoding = contentEncoding;
   }
 
   @Override
@@ -54,10 +60,10 @@ public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
       source.read(wrap);
       try {
         // Step 2: wait for the object metadata, this is populated in the first message from GCS
-        Object object = source.getResult().get();
+        String contentEncoding = this.contentEncoding.get();
         // if the Content-Encoding is gzip, Step 3: wire gzip decompression into the byte path
         //   this will have a copy impact as we are no longer controlling all the buffers
-        if ("gzip".equals(object.getContentEncoding())) {
+        if ("gzip".equals(contentEncoding) || "x-gzip".equals(contentEncoding)) {
           // to wire gzip decompression into the byte path:
           // Create an input stream of the first4 bytes we already read
           ByteArrayInputStream first4again = new ByteArrayInputStream(first4);
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/HttpDownloadSessionBuilder.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/HttpDownloadSessionBuilder.java
@@ -25,6 +25,7 @@
 import com.google.cloud.storage.BlobReadChannelV2.BlobReadChannelContext;
 import com.google.cloud.storage.BufferedReadableByteChannelSession.BufferedReadableByteChannel;
 import com.google.cloud.storage.UnbufferedReadableByteChannelSession.UnbufferedReadableByteChannel;
+import com.google.common.util.concurrent.MoreExecutors;
 import java.nio.ByteBuffer;
 import java.util.function.BiFunction;
 import javax.annotation.concurrent.Immutable;
@@ -51,10 +52,18 @@ public ReadableByteChannelSessionBuilder byteChannel(
   public static final class ReadableByteChannelSessionBuilder {
 
     private final BlobReadChannelContext blobReadChannelContext;
+    private boolean autoGzipDecompression;
     // private Hasher hasher; // TODO: wire in Hasher
 
     private ReadableByteChannelSessionBuilder(BlobReadChannelContext blobReadChannelContext) {
       this.blobReadChannelContext = blobReadChannelContext;
+      this.autoGzipDecompression = false;
+    }
+
+    public ReadableByteChannelSessionBuilder setAutoGzipDecompression(
+        boolean autoGzipDecompression) {
+      this.autoGzipDecompression = autoGzipDecompression;
+      return this;
     }
 
     public BufferedReadableByteChannelSessionBuilder buffered() {
@@ -77,13 +86,27 @@ public UnbufferedReadableByteChannelSessionBuilder unbuffered() {
             ApiaryReadRequest, SettableApiFuture<StorageObject>, UnbufferedReadableByteChannel>
         bindFunction() {
       // for any non-final value, create a reference to the value at this point in time
-      return (request, resultFuture) ->
-          new ApiaryUnbufferedReadableByteChannel(
+      boolean autoGzipDecompression = this.autoGzipDecompression;
+      return (request, resultFuture) -> {
+        if (autoGzipDecompression) {
+          return new GzipReadableByteChannel(
+              new ApiaryUnbufferedReadableByteChannel(
+                  request,
+                  blobReadChannelContext.getApiaryClient(),
+                  resultFuture,
+                  blobReadChannelContext.getStorageOptions(),
+                  blobReadChannelContext.getRetryAlgorithmManager().idempotent()),
+              ApiFutures.transform(
+                  resultFuture, StorageObject::getContentEncoding, MoreExecutors.directExecutor()));
+        } else {
+          return new ApiaryUnbufferedReadableByteChannel(
               request,
               blobReadChannelContext.getApiaryClient(),
               resultFuture,
               blobReadChannelContext.getStorageOptions(),
               blobReadChannelContext.getRetryAlgorithmManager().idempotent());
+        }
+      };
     }
 
     public static final class BufferedReadableByteChannelSessionBuilder {
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITBlobReadChannelTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITBlobReadChannelTest.java
@@ -427,7 +427,7 @@ public void limitAfterReadWorks() throws IOException {
       String xxd1 = xxd(bytes1);
       assertThat(xxd1).isEqualTo(xxdExpected1);
 
-      // seek forward to a new offset
+      // change the limit
       reader.limit(10);
 
       // read again
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITBlobReadChannelV2RetryTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITBlobReadChannelV2RetryTest.java
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/it/runner/registry/TestBench.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/it/runner/registry/TestBench.java
