feat: improve throughput of http based storage#reader between 100 MiB/s and 200 MiB/s

google-cloud-storage/clirr-ignored-differences.xml

@@ -1,6 +1,18 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <!-- see https://www.mojohaus.org/clirr-maven-plugin/examples/ignored-differences.html -->
 <differences>
+  <!-- Not breaking, internal only interface and the new methods have default implementations -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/storage/UnbufferedReadableByteChannelSession$UnbufferedReadableByteChannel</className>
+    <method>* read(*)</method>
+  </difference>
+  <!-- Allow accessing the underlying Apiary instance -->
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/cloud/storage/spi/v1/StorageRpc</className>
+    <method>* getStorage()</method>
+  </difference>
 
   <difference>
     <differenceType>8001</differenceType>

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

@@ -0,0 +1,338 @@
+/*
+ * Copyright 2022 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;
+
+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.core.SettableApiFuture;
+import com.google.api.gax.retrying.ResultRetryAlgorithm;
+import com.google.api.services.storage.Storage;
+import com.google.api.services.storage.Storage.Objects;
+import com.google.api.services.storage.Storage.Objects.Get;
+import com.google.api.services.storage.model.StorageObject;
+import com.google.cloud.storage.UnbufferedReadableByteChannelSession.UnbufferedReadableByteChannel;
+import com.google.cloud.storage.spi.v1.StorageRpc;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.MoreObjects;
+import com.google.common.hash.HashFunction;
+import com.google.common.hash.Hashing;
+import com.google.common.io.BaseEncoding;
+import com.google.gson.Gson;
+import com.google.gson.stream.JsonReader;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import java.io.StringReader;
+import java.nio.ByteBuffer;
+import java.nio.channels.Channels;
+import java.nio.channels.ReadableByteChannel;
+import java.nio.channels.ScatteringByteChannel;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import javax.annotation.concurrent.Immutable;
+
+class ApiaryUnbufferedReadableByteChannel implements UnbufferedReadableByteChannel {
+
+  private final ApiaryReadRequest apiaryReadRequest;
+  private final Storage storage;
+  private final SettableApiFuture<StorageObject> result;
+  private final HttpStorageOptions options;
+  private final ResultRetryAlgorithm<?> resultRetryAlgorithm;
+  private final Consumer<StorageObject> resolvedObjectCallback;
+
+  private long position;
+  private ScatteringByteChannel sbc;
+  private boolean open;
+
+  // returned X-Goog-Generation header value
+  private Long xGoogGeneration;
+
+  ApiaryUnbufferedReadableByteChannel(
+      ApiaryReadRequest apiaryReadRequest,
+      Storage storage,
+      SettableApiFuture<StorageObject> result,
+      HttpStorageOptions options,
+      ResultRetryAlgorithm<?> resultRetryAlgorithm,
+      Consumer<StorageObject> resolvedObjectCallback) {
+    this.apiaryReadRequest = apiaryReadRequest;
+    this.storage = storage;
+    this.result = result;
+    this.options = options;
+    this.resultRetryAlgorithm = resultRetryAlgorithm;
+    this.resolvedObjectCallback = resolvedObjectCallback;
+    this.open = true;
+    this.position =
+        apiaryReadRequest.getByteRangeSpec() != null
+            ? apiaryReadRequest.getByteRangeSpec().beginOffset()
+            : 0;
+  }
+
+  @SuppressWarnings("UnnecessaryContinue")
+  @Override
+  public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
+    do {
+      if (sbc == null) {
+        sbc = Retrying.run(options, resultRetryAlgorithm, this::open, Function.identity());
+      }
+
+      try {
+        // According to the contract of Retrying#run it's possible for sbc to be null even after
+        // invocation. However, the function we provide is guaranteed to return non-null or throw
+        // an exception. So we suppress the warning from intellij here.
+        //noinspection ConstantConditions
+        long read = sbc.read(dsts, offset, length);
+        if (read == -1) {
+          open = false;
+        } else {
+          position += read;
+        }
+        return read;
+      } catch (Exception t) {
+        if (resultRetryAlgorithm.shouldRetry(t, null)) {
+          // 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;
+          }
+        } else {
+          throw new IOException(StorageException.coalesce(t));
+        }
+      }
+    } while (true);
+  }
+
+  @Override
+  public boolean isOpen() {
+    return open;
+  }
+
+  @Override
+  public void close() throws IOException {
+    open = false;
+    if (sbc != null) {
+      sbc.close();
+    }
+  }
+
+  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);
+
+      HttpResponse media = get.executeMedia();
+      InputStream content = media.getContent();
+      if (xGoogGeneration == null) {
+        HttpHeaders responseHeaders = media.getHeaders();
+        //noinspection unchecked
+        List<String> xGoogGenHeader = (List<String>) responseHeaders.get("x-goog-generation");
+        // TODO: wire in result metadata population
+        if (xGoogGenHeader != null && !xGoogGenHeader.isEmpty()) {
+          String s = xGoogGenHeader.get(0);
+          Long generation = Long.valueOf(s);
+          this.xGoogGeneration = generation;
+          resolvedObjectCallback.accept(
+              apiaryReadRequest.getObject().clone().setGeneration(generation));
+        }
+      }
+
+      ReadableByteChannel rbc = Channels.newChannel(content);
+      return StorageByteChannels.readable().asScatteringByteChannel(rbc);
+    } catch (HttpResponseException e) {
+      if (xGoogGeneration != null) {
+        int statusCode = e.getStatusCode();
+        if (statusCode == 404) {
+          throw new StorageException(404, "Failure while trying to resume download", e);
+        }
+      }
+      throw StorageException.translate(e);
+    } catch (IOException e) {
+      throw StorageException.translate(e);
+    } catch (Throwable t) {
+      throw StorageException.coalesce(t);
+    }
+  }
+
+  @VisibleForTesting
+  static Get createGetRequest(
+      ApiaryReadRequest apiaryReadRequest,
+      Objects objects,
+      Long xGoogGeneration,
+      boolean returnRawInputStream)
+      throws IOException {
+    StorageObject from = apiaryReadRequest.getObject();
+    Map<StorageRpc.Option, ?> options = apiaryReadRequest.getOptions();
+    Get get = objects.get(from.getBucket(), from.getName());
+    if (from.getGeneration() != null) {
+      get.setGeneration(from.getGeneration());
+    } else if (xGoogGeneration != null) {
+      get.setGeneration(xGoogGeneration);
+    }
+    ifNonNull(
+        options.get(StorageRpc.Option.IF_GENERATION_MATCH),
+        ApiaryUnbufferedReadableByteChannel::cast,
+        get::setIfGenerationMatch);
+    ifNonNull(
+        options.get(StorageRpc.Option.IF_GENERATION_NOT_MATCH),
+        ApiaryUnbufferedReadableByteChannel::cast,
+        get::setIfGenerationNotMatch);
+    ifNonNull(
+        options.get(StorageRpc.Option.IF_METAGENERATION_MATCH),
+        ApiaryUnbufferedReadableByteChannel::cast,
+        get::setIfMetagenerationMatch);
+    ifNonNull(
+        options.get(StorageRpc.Option.IF_METAGENERATION_NOT_MATCH),
+        ApiaryUnbufferedReadableByteChannel::cast,
+        get::setIfMetagenerationNotMatch);
+    ifNonNull(
+        options.get(StorageRpc.Option.USER_PROJECT),
+        ApiaryUnbufferedReadableByteChannel::cast,
+        get::setUserProject);
+    HttpHeaders headers = get.getRequestHeaders();
+    ifNonNull(
+        options.get(StorageRpc.Option.CUSTOMER_SUPPLIED_KEY),
+        ApiaryUnbufferedReadableByteChannel::cast,
+        (String key) -> {
+          BaseEncoding base64 = BaseEncoding.base64();
+          HashFunction hashFunction = Hashing.sha256();
+          headers.set("x-goog-encryption-algorithm", "AES256");
+          headers.set("x-goog-encryption-key", key);
+          headers.set(
+              "x-goog-encryption-key-sha256",
+              base64.encode(hashFunction.hashBytes(base64.decode(key)).asBytes()));
+        });
+
+    get.setReturnRawInputStream(returnRawInputStream);
+    String range = apiaryReadRequest.getByteRangeSpec().getHttpRangeHeader();
+    if (range != null) {
+      get.getRequestHeaders().setRange(range);
+    }
+    get.getMediaHttpDownloader().setDirectDownloadEnabled(true);
+
+    return get;
+  }
+
+  @SuppressWarnings("unchecked")
+  private static <T> T cast(Object o) {
+    return (T) o;
+  }
+
+  @Immutable
+  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;
+
+    private volatile String objectJson;
+
+    ApiaryReadRequest(
+        StorageObject object, Map<StorageRpc.Option, ?> options, ByteRangeSpec byteRangeSpec) {
+      this.object = object;
+      this.options = options;
+      this.byteRangeSpec = byteRangeSpec;
+    }
+
+    StorageObject getObject() {
+      return object;
+    }
+
+    Map<StorageRpc.Option, ?> getOptions() {
+      return options;
+    }
+
+    ByteRangeSpec getByteRangeSpec() {
+      return byteRangeSpec;
+    }
+
+    ApiaryReadRequest withNewBeginOffset(long beginOffset) {
+      if (beginOffset > 0 && beginOffset != byteRangeSpec.beginOffset()) {
+        return new ApiaryReadRequest(
+            object, options, byteRangeSpec.withNewBeginOffset(beginOffset));
+      } else {
+        return this;
+      }
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (!(o instanceof ApiaryReadRequest)) {
+        return false;
+      }
+      ApiaryReadRequest that = (ApiaryReadRequest) o;
+      return java.util.Objects.equals(object, that.object)
+          && java.util.Objects.equals(options, that.options)
+          && java.util.Objects.equals(byteRangeSpec, that.byteRangeSpec);
+    }
+
+    @Override
+    public int hashCode() {
+      return java.util.Objects.hash(object, options, byteRangeSpec);
+    }
+
+    @Override
+    public String toString() {
+      return MoreObjects.toStringHelper(this)
+          .add("byteRangeSpec", byteRangeSpec)
+          .add("options", options)
+          .add("object", getObjectJson())
+          .toString();
+    }
+
+    private String getObjectJson() {
+      if (objectJson == null) {
+        synchronized (this) {
+          if (objectJson == null) {
+            objectJson = gson.toJson(object);
+          }
+        }
+      }
+      return objectJson;
+    }
+
+    private void writeObject(ObjectOutputStream out) throws IOException {
+      String ignore = getObjectJson();
+      out.defaultWriteObject();
+    }
+
+    private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
+      in.defaultReadObject();
+      JsonReader jsonReader = gson.newJsonReader(new StringReader(this.objectJson));
+      this.object = gson.fromJson(jsonReader, StorageObject.class);
+    }
+  }
+}