Repository: googleapis/java-storage
Pull request #1180 (2021-12-16 20:28:57)

feat: allow limiting ReadChannel
The ReadChannel returned from storage.reader and blob.reader now also implemented LimitableReadChannel which allows limiting the max number of bytes the channel will make available for read.

This can be used in conjunction with seek to allow for range gets of objects independent of any buffer or chunk sizes.

See the integration tests for an example of how this could be used.


Patch:
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobReadChannel.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobReadChannel.java
index 8d51ab3b..678f138e 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobReadChannel.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobReadChannel.java
@@ -18,6 +18,7 @@ package com.google.cloud.storage;
 
 import static com.google.cloud.RetryHelper.runWithRetries;
 
+import com.google.api.client.util.Preconditions;
 import com.google.api.gax.retrying.ResultRetryAlgorithm;
 import com.google.api.services.storage.model.StorageObject;
 import com.google.cloud.ReadChannel;
@@ -52,6 +53,7 @@ class BlobReadChannel implements ReadChannel {
   private final StorageObject storageObject;
   private int bufferPos;
   private byte[] buffer;
+  private long limit;
 
   BlobReadChannel(
       StorageOptions serviceOptions, BlobId blob, Map<StorageRpc.Option, ?> requestOptions) {
@@ -62,6 +64,7 @@ class BlobReadChannel implements ReadChannel {
     isOpen = true;
     storageRpc = serviceOptions.getStorageRpcV1();
     storageObject = blob.toPb();
+    this.limit = Long.MAX_VALUE;
   }
 
   @Override
@@ -71,7 +74,8 @@ class BlobReadChannel implements ReadChannel {
             .setPosition(position)
             .setIsOpen(isOpen)
             .setEndOfStream(endOfStream)
-            .setChunkSize(chunkSize);
+            .setChunkSize(chunkSize)
+            .setLimit(limit);
     if (buffer != null) {
       builder.setPosition(position + bufferPos);
       builder.setEndOfStream(false);
@@ -119,7 +123,8 @@ class BlobReadChannel implements ReadChannel {
       if (endOfStream) {
         return -1;
       }
-      final int toRead = Math.max(byteBuffer.remaining(), chunkSize);
+      final int toRead =
+          Math.toIntExact(Math.min(limit - position, Math.max(byteBuffer.remaining(), chunkSize)));
       try {
         ResultRetryAlgorithm<?> algorithm =
             retryAlgorithmManager.getForObjectsGet(storageObject, requestOptions);
@@ -158,6 +163,18 @@ class BlobReadChannel implements ReadChannel {
     return toWrite;
   }
 
+  @Override
+  public ReadChannel limit(long limit) {
+    Preconditions.checkArgument(limit >= 0, "Limit must be >= 0");
+    this.limit = limit;
+    return this;
+  }
+
+  @Override
+  public long limit() {
+    return limit;
+  }
+
   static class StateImpl implements RestorableState<ReadChannel>, Serializable {
 
     private static final long serialVersionUID = 3889420316004453706L;
@@ -170,6 +187,7 @@ class BlobReadChannel implements ReadChannel {
     private final boolean isOpen;
     private final boolean endOfStream;
     private final int chunkSize;
+    private final long limit;
 
     StateImpl(Builder builder) {
       this.serviceOptions = builder.serviceOptions;
@@ -180,6 +198,7 @@ class BlobReadChannel implements ReadChannel {
       this.isOpen = builder.isOpen;
       this.endOfStream = builder.endOfStream;
       this.chunkSize = builder.chunkSize;
+      this.limit = builder.limit;
     }
 
     static class Builder {
@@ -191,6 +210,7 @@ class BlobReadChannel implements ReadChannel {
       private boolean isOpen;
       private boolean endOfStream;
       private int chunkSize;
+      private long limit;
 
       private Builder(StorageOptions options, BlobId blob, Map<StorageRpc.Option, ?> reqOptions) {
         this.serviceOptions = options;
@@ -223,6 +243,11 @@ class BlobReadChannel implements ReadChannel {
         return this;
       }
 
+      Builder setLimit(long limit) {
+        this.limit = limit;
+        return this;
+      }
+
       RestorableState<ReadChannel> build() {
         return new StateImpl(this);
       }
@@ -241,13 +266,22 @@ class BlobReadChannel implements ReadChannel {
       channel.isOpen = isOpen;
       channel.endOfStream = endOfStream;
       channel.chunkSize = chunkSize;
+      channel.limit = limit;
       return channel;
     }
 
     @Override
     public int hashCode() {
       return Objects.hash(
-          serviceOptions, blob, requestOptions, lastEtag, position, isOpen, endOfStream, chunkSize);
+          serviceOptions,
+          blob,
+          requestOptions,
+          lastEtag,
+          position,
+          isOpen,
+          endOfStream,
+          chunkSize,
+          limit);
     }
 
     @Override
@@ -266,7 +300,8 @@ class BlobReadChannel implements ReadChannel {
           && this.position == other.position
           && this.isOpen == other.isOpen
           && this.endOfStream == other.endOfStream
-          && this.chunkSize == other.chunkSize;
+          && this.chunkSize == other.chunkSize
+          && this.limit == other.limit;
     }
 
     @Override
@@ -276,6 +311,7 @@ class BlobReadChannel implements ReadChannel {
           .add("position", position)
           .add("isOpen", isOpen)
           .add("endOfStream", endOfStream)
+          .add("limit", limit)
           .toString();
     }
   }
diff --git a/samples/native-image-sample/pom.xml b/samples/native-image-sample/pom.xml
index 207facce..47585024 100644
--- a/samples/native-image-sample/pom.xml
+++ b/samples/native-image-sample/pom.xml
@@ -139,7 +139,7 @@
           <plugin>
             <groupId>org.graalvm.buildtools</groupId>
             <artifactId>native-maven-plugin</artifactId>
-            <version>0.9.9</version>
+            <version>0.9.10</version>
             <extensions>true</extensions>
             <configuration>
               <mainClass>com.example.storage.NativeImageStorageSample</mainClass>


Test patch:
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/BlobReadChannelTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/BlobReadChannelTest.java
index dafb0b2a..6eb2b8aa 100644
--- a/google-cloud-storage/src/test/java/com/google/cloud/storage/BlobReadChannelTest.java
+++ b/google-cloud-storage/src/test/java/com/google/cloud/storage/BlobReadChannelTest.java
@@ -223,13 +223,19 @@ public class BlobReadChannelTest {
   public void testStateEquals() {
     replay(storageRpcMock);
     reader = new BlobReadChannel(options, BLOB_ID, EMPTY_RPC_OPTIONS);
+    int limit = 342;
+    reader.limit(limit);
     @SuppressWarnings("resource") // avoid closing when you don't want partial writes to GCS
     ReadChannel secondReader = new BlobReadChannel(options, BLOB_ID, EMPTY_RPC_OPTIONS);
+    secondReader.limit(limit);
     RestorableState<ReadChannel> state = reader.capture();
     RestorableState<ReadChannel> secondState = secondReader.capture();
     assertEquals(state, secondState);
     assertEquals(state.hashCode(), secondState.hashCode());
     assertEquals(state.toString(), secondState.toString());
+
+    ReadChannel restore = secondState.restore();
+    assertEquals(limit, restore.limit());
   }
 
   private static byte[] randomByteArray(int size) {
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
new file mode 100644
index 00000000..3f1470a0
--- /dev/null
+++ b/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITBlobReadChannelTest.java
@@ -0,0 +1,107 @@
+/*
+ * 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.it;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.cloud.NoCredentials;
+import com.google.cloud.ReadChannel;
+import com.google.cloud.WriteChannel;
+import com.google.cloud.storage.BlobInfo;
+import com.google.cloud.storage.Bucket;
+import com.google.cloud.storage.BucketInfo;
+import com.google.cloud.storage.DataGeneration;
+import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageOptions;
+import com.google.cloud.storage.conformance.retry.TestBench;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Random;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TestName;
+
+public final class ITBlobReadChannelTest {
+
+  private static final int _16MiB = 16 * 1024 * 1024;
+  private static final int _256KiB = 256 * 1024;
+
+  @ClassRule
+  public static final TestBench testBench =
+      TestBench.newBuilder().setContainerName("blob-read-channel-test").build();
+
+  @Rule public final TestName testName = new TestName();
+
+  @Rule public final DataGeneration dataGeneration = new DataGeneration(new Random(872364872));
+
+  @Test
+  public void testLimit_smallerThanOneChunk() throws IOException {
+    int srcContentSize = _256KiB;
+    int rangeBegin = 57;
+    int rangeEnd = 2384;
+    int chunkSize = _16MiB;
+    doLimitTest(srcContentSize, rangeBegin, rangeEnd, chunkSize);
+  }
+
+  @Test
+  public void testLimit_largerThanOneChunk() throws IOException {
+    int srcContentSize = _16MiB + (_256KiB * 3);
+    int rangeBegin = 384;
+    int rangeEnd = rangeBegin + _16MiB;
+    int chunkSize = _16MiB;
+
+    doLimitTest(srcContentSize, rangeBegin, rangeEnd, chunkSize);
+  }
+
+  private void doLimitTest(int srcContentSize, int rangeBegin, int rangeEnd, int chunkSize)
+      throws IOException {
+    Storage s =
+        StorageOptions.newBuilder()
+            .setProjectId("blob-read-channel-test")
+            .setHost(testBench.getBaseUri())
+            .setCredentials(NoCredentials.getInstance())
+            .build()
+            .getService();
+
+    String testNameMethodName = testName.getMethodName();
+    String bucketName = String.format("bucket-%s", testNameMethodName.toLowerCase());
+    String blobName = String.format("%s/src", testNameMethodName);
+
+    Bucket bucket = s.create(BucketInfo.of(bucketName));
+    BlobInfo src = BlobInfo.newBuilder(bucket, blobName).build();
+    ByteBuffer content = dataGeneration.randByteBuffer(srcContentSize);
+    ByteBuffer expectedSubContent = content.duplicate();
+    expectedSubContent.position(rangeBegin);
+    expectedSubContent.limit(rangeEnd);
+    try (WriteChannel writer = s.writer(src)) {
+      writer.write(content);
+    }
+
+    ByteBuffer actual = ByteBuffer.allocate(rangeEnd - rangeBegin);
+
+    try (ReadChannel reader = s.reader(src.getBlobId())) {
+      reader.setChunkSize(chunkSize);
+      reader.seek(rangeBegin);
+      reader.limit(rangeEnd);
+      reader.read(actual);
+      actual.flip();
+    }
+
+    assertThat(actual).isEqualTo(expectedSubContent);
+  }
+}