feat: Adds a ZeroCopy response marshaller for grpc ReadObject handling (#2489)

* feat: Adds a ZeroCopy response marshaller for grpc ReadObject handling

* Update google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcStorageOptions.java

Co-authored-by: BenWhitehead <BenWhitehead@users.noreply.github.com>

* add copyright header

* Apply suggestions from code review

Co-authored-by: BenWhitehead <BenWhitehead@users.noreply.github.com>

* Extract classes to own files, more StorageClient initialization

* copyright headers on new files

* formatter

* one more lint issue

* fix: improve GrpcStorageOptions.ReadObjectResponseZeroCopyMessageMarshaller#close() handling of multiple IOExceptions

* clean up ZeroCopyMarshallerTest

* add gprc core to pom

---------

Co-authored-by: BenWhitehead <BenWhitehead@users.noreply.github.com>

Author: JesseLovelace

google-cloud-storage/pom.xml

@@ -96,6 +96,14 @@
       <groupId>com.google.protobuf</groupId>
       <artifactId>protobuf-java-util</artifactId>
     </dependency>
+    <dependency>
+      <groupId>io.grpc</groupId>
+      <artifactId>grpc-core</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.grpc</groupId>
+      <artifactId>grpc-protobuf</artifactId>
+    </dependency>
     <dependency>
       <groupId>com.google.api.grpc</groupId>
       <artifactId>proto-google-common-protos</artifactId>

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

@@ -48,19 +48,23 @@ public static GapicDownloadSessionBuilder create() {
    * ultimately produced channel will not do any retries of its own.
    */
   public ReadableByteChannelSessionBuilder byteChannel(
-      ServerStreamingCallable<ReadObjectRequest, ReadObjectResponse> read) {
-    return new ReadableByteChannelSessionBuilder(read);
+      ServerStreamingCallable<ReadObjectRequest, ReadObjectResponse> read,
+      ResponseContentLifecycleManager responseContentLifecycleManager) {
+    return new ReadableByteChannelSessionBuilder(read, responseContentLifecycleManager);
   }
 
   public static final class ReadableByteChannelSessionBuilder {
 
     private final ServerStreamingCallable<ReadObjectRequest, ReadObjectResponse> read;
+    private final ResponseContentLifecycleManager responseContentLifecycleManager;
     private boolean autoGzipDecompression;
     private Hasher hasher;
 
     private ReadableByteChannelSessionBuilder(
-        ServerStreamingCallable<ReadObjectRequest, ReadObjectResponse> read) {
+        ServerStreamingCallable<ReadObjectRequest, ReadObjectResponse> read,
+        ResponseContentLifecycleManager responseContentLifecycleManager) {
       this.read = read;
+      this.responseContentLifecycleManager = responseContentLifecycleManager;
       this.hasher = Hasher.noop();
       this.autoGzipDecompression = false;
     }
@@ -100,11 +104,13 @@ public UnbufferedReadableByteChannelSessionBuilder unbuffered() {
       return (object, resultFuture) -> {
         if (autoGzipDecompression) {
           return new GzipReadableByteChannel(
-              new GapicUnbufferedReadableByteChannel(resultFuture, read, object, hasher),
+              new GapicUnbufferedReadableByteChannel(
+                  resultFuture, read, object, hasher, responseContentLifecycleManager),
               ApiFutures.transform(
                   resultFuture, Object::getContentEncoding, MoreExecutors.directExecutor()));
         } else {
-          return new GapicUnbufferedReadableByteChannel(resultFuture, read, object, hasher);
+          return new GapicUnbufferedReadableByteChannel(
+              resultFuture, read, object, hasher, responseContentLifecycleManager);
         }
       };
     }

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

@@ -16,15 +16,14 @@
 
 package com.google.cloud.storage;
 
-import static com.google.common.base.Preconditions.checkArgument;
-
 import com.google.api.client.http.HttpStatusCodes;
 import com.google.api.core.ApiFuture;
 import com.google.api.core.SettableApiFuture;
 import com.google.api.gax.rpc.ServerStream;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.cloud.storage.Crc32cValue.Crc32cLengthKnown;
 import com.google.cloud.storage.UnbufferedReadableByteChannelSession.UnbufferedReadableByteChannel;
+import com.google.protobuf.ByteString;
 import com.google.storage.v2.ChecksummedData;
 import com.google.storage.v2.Object;
 import com.google.storage.v2.ReadObjectRequest;
@@ -46,25 +45,28 @@ final class GapicUnbufferedReadableByteChannel
   private final ReadObjectRequest req;
   private final Hasher hasher;
   private final LazyServerStreamIterator iter;
+  private final ResponseContentLifecycleManager rclm;
 
   private boolean open = true;
   private boolean complete = false;
   private long blobOffset;
 
   private Object metadata;
 
-  private ByteBuffer leftovers;
+  private ResponseContentLifecycleHandle leftovers;
 
   GapicUnbufferedReadableByteChannel(
       SettableApiFuture<Object> result,
       ServerStreamingCallable<ReadObjectRequest, ReadObjectResponse> read,
       ReadObjectRequest req,
-      Hasher hasher) {
+      Hasher hasher,
+      ResponseContentLifecycleManager rclm) {
     this.result = result;
     this.read = read;
     this.req = req;
     this.hasher = hasher;
     this.blobOffset = req.getReadOffset();
+    this.rclm = rclm;
     this.iter = new LazyServerStreamIterator();
   }
 
@@ -82,15 +84,17 @@ public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
     ReadCursor c = new ReadCursor(blobOffset, blobOffset + totalBufferCapacity);
     while (c.hasRemaining()) {
       if (leftovers != null) {
-        copy(c, leftovers, dsts, offset, length);
+        leftovers.copy(c, dsts, offset, length);
         if (!leftovers.hasRemaining()) {
+          leftovers.close();
           leftovers = null;
         }
         continue;
       }
 
       if (iter.hasNext()) {
         ReadObjectResponse resp = iter.next();
+        ResponseContentLifecycleHandle handle = rclm.get(resp);
         if (resp.hasMetadata()) {
           Object respMetadata = resp.getMetadata();
           if (metadata == null) {
@@ -107,22 +111,24 @@ public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
           }
         }
         ChecksummedData checksummedData = resp.getChecksummedData();
-        ByteBuffer content = checksummedData.getContent().asReadOnlyByteBuffer();
-        // very important to know whether a crc32c value is set. Without checking, protobuf will
+        ByteString content = checksummedData.getContent();
+        int contentSize = content.size();
+        // Very important to know whether a crc32c value is set. Without checking, protobuf will
         // happily return 0, which is a valid crc32c value.
         if (checksummedData.hasCrc32C()) {
-          Crc32cLengthKnown expected =
-              Crc32cValue.of(checksummedData.getCrc32C(), checksummedData.getContent().size());
+          Crc32cLengthKnown expected = Crc32cValue.of(checksummedData.getCrc32C(), contentSize);
           try {
-            hasher.validate(expected, content::duplicate);
+            hasher.validate(expected, content.asReadOnlyByteBufferList());
           } catch (IOException e) {
             close();
             throw e;
           }
         }
-        copy(c, content, dsts, offset, length);
-        if (content.hasRemaining()) {
-          leftovers = content;
+        handle.copy(c, dsts, offset, length);
+        if (handle.hasRemaining()) {
+          leftovers = handle;
+        } else {
+          handle.close();
         }
       } else {
         complete = true;
@@ -144,59 +150,26 @@ public boolean isOpen() {
   @Override
   public void close() throws IOException {
     open = false;
-    iter.close();
+    try {
+      if (leftovers != null) {
+        leftovers.close();
+      }
+    } finally {
+      iter.close();
+    }
   }
 
   ApiFuture<Object> getResult() {
     return result;
   }
 
-  private void copy(ReadCursor c, ByteBuffer content, ByteBuffer[] dsts, int offset, int length) {
-    long copiedBytes = Buffers.copy(content, dsts, offset, length);
-    c.advance(copiedBytes);
-  }
-
   private IOException closeWithError(String message) throws IOException {
     close();
     StorageException cause =
         new StorageException(HttpStatusCodes.STATUS_CODE_PRECONDITION_FAILED, message);
     throw new IOException(message, cause);
   }
 
-  /**
-   * Shrink wraps a beginning, offset and limit for tracking state of an individual invocation of
-   * {@link #read}
-   */
-  private static final class ReadCursor {
-    private final long beginning;
-    private long offset;
-    private final long limit;
-
-    private ReadCursor(long beginning, long limit) {
-      this.limit = limit;
-      this.beginning = beginning;
-      this.offset = beginning;
-    }
-
-    public boolean hasRemaining() {
-      return limit - offset > 0;
-    }
-
-    public void advance(long incr) {
-      checkArgument(incr >= 0);
-      offset += incr;
-    }
-
-    public long read() {
-      return offset - beginning;
-    }
-
-    @Override
-    public String toString() {
-      return String.format("ReadCursor{begin=%d, offset=%d, limit=%d}", beginning, offset, limit);
-    }
-  }
-
   private final class LazyServerStreamIterator implements Iterator<ReadObjectResponse>, Closeable {
     private ServerStream<ReadObjectResponse> serverStream;
     private Iterator<ReadObjectResponse> responseIterator;

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

@@ -28,15 +28,18 @@
 final class GrpcBlobReadChannel extends BaseStorageReadChannel<Object> {
 
   private final ServerStreamingCallable<ReadObjectRequest, ReadObjectResponse> read;
+  private final ResponseContentLifecycleManager responseContentLifecycleManager;
   private final ReadObjectRequest request;
   private final boolean autoGzipDecompression;
 
   GrpcBlobReadChannel(
       ServerStreamingCallable<ReadObjectRequest, ReadObjectResponse> read,
+      ResponseContentLifecycleManager responseContentLifecycleManager,
       ReadObjectRequest request,
       boolean autoGzipDecompression) {
     super(Conversions.grpc().blobInfo());
     this.read = read;
+    this.responseContentLifecycleManager = responseContentLifecycleManager;
     this.request = request;
     this.autoGzipDecompression = autoGzipDecompression;
   }
@@ -53,7 +56,7 @@ protected LazyReadChannel<?, Object> newLazyReadChannel() {
           ReadableByteChannelSessionBuilder b =
               ResumableMedia.gapic()
                   .read()
-                  .byteChannel(read)
+                  .byteChannel(read, responseContentLifecycleManager)
                   .setHasher(Hasher.noop())
                   .setAutoGzipDecompression(autoGzipDecompression);
           BufferHandle bufferHandle = getBufferHandle();

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

@@ -180,6 +180,7 @@ final class GrpcStorageImpl extends BaseService<StorageOptions>
                   .collect(ImmutableSet.toImmutableSet())));
 
   final StorageClient storageClient;
+  final ResponseContentLifecycleManager responseContentLifecycleManager;
   final WriterFactory writerFactory;
   final GrpcConversions codecs;
   final GrpcRetryAlgorithmManager retryAlgorithmManager;
@@ -192,10 +193,12 @@ final class GrpcStorageImpl extends BaseService<StorageOptions>
   GrpcStorageImpl(
       GrpcStorageOptions options,
       StorageClient storageClient,
+      ResponseContentLifecycleManager responseContentLifecycleManager,
       WriterFactory writerFactory,
       Opts<UserProject> defaultOpts) {
     super(options);
     this.storageClient = storageClient;
+    this.responseContentLifecycleManager = responseContentLifecycleManager;
     this.writerFactory = writerFactory;
     this.defaultOpts = defaultOpts;
     this.codecs = Conversions.grpc();
@@ -716,8 +719,10 @@ public GrpcBlobReadChannel reader(BlobId blob, BlobSourceOption... options) {
     ReadObjectRequest request = getReadObjectRequest(blob, opts);
     Set<StatusCode.Code> codes = resultRetryAlgorithmToCodes(retryAlgorithmManager.getFor(request));
     GrpcCallContext grpcCallContext = Retrying.newCallContext().withRetryableCodes(codes);
+
     return new GrpcBlobReadChannel(
         storageClient.readObjectCallable().withDefaultCallContext(grpcCallContext),
+        responseContentLifecycleManager,
         request,
         !opts.autoGzipDecompression());
   }
@@ -1868,7 +1873,9 @@ private UnbufferedReadableByteChannelSession<Object> unbufferedReadSession(
         opts.grpcMetadataMapper().apply(Retrying.newCallContext().withRetryableCodes(codes));
     return ResumableMedia.gapic()
         .read()
-        .byteChannel(storageClient.readObjectCallable().withDefaultCallContext(grpcCallContext))
+        .byteChannel(
+            storageClient.readObjectCallable().withDefaultCallContext(grpcCallContext),
+            responseContentLifecycleManager)
         .setAutoGzipDecompression(!opts.autoGzipDecompression())
         .unbuffered()
         .setReadObjectRequest(readObjectRequest)