chore: start plumbing the zero copy~ness closer to the public api

Author: BenWhitehead

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

@@ -17,11 +17,23 @@
 package com.google.cloud.storage;
 
 import com.google.api.core.ApiFuture;
+import com.google.protobuf.ByteString;
+import java.io.Closeable;
+import java.io.IOException;
 
 /** Blob Descriptor is to blob, what File Descriptor is to a file */
 public interface BlobDescriptor {
 
   BlobInfo getBlobInfo();
 
   ApiFuture<byte[]> readRangeAsBytes(ByteRangeSpec range);
+
+  interface ZeroCopySupport {
+    interface DisposableByteString extends AutoCloseable, Closeable {
+      ByteString byteString();
+
+      @Override
+      void close() throws IOException;
+    }
+  }
 }

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

@@ -34,8 +34,9 @@
 import com.google.storage.v2.Object;
 import com.google.storage.v2.ObjectRangeData;
 import com.google.storage.v2.ReadRange;
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -244,7 +245,7 @@ protected void onStartImpl(StreamController controller) {
     @Override
     protected void onResponseImpl(BidiReadObjectResponse response) {
       controller.request(1);
-      try (ResponseContentLifecycleHandle handle =
+      try (ResponseContentLifecycleHandle<BidiReadObjectResponse> handle =
           bidiResponseContentLifecycleManager.get(response)) {
         if (response.hasMetadata()) {
           state.metadata.set(response.getMetadata());
@@ -256,21 +257,32 @@ protected void onResponseImpl(BidiReadObjectResponse response) {
         if (rangeData.isEmpty()) {
           return;
         }
-        for (ObjectRangeData d : rangeData) {
+        for (int i = 0; i < rangeData.size(); i++) {
+          ObjectRangeData d = rangeData.get(i);
           long id = d.getReadRange().getReadId();
           OutstandingReadToArray read = state.outstandingReads.get(id);
           if (read == null) {
             continue;
           }
-          ByteString content = d.getChecksummedData().getContent();
-          ChildRef childRef = handle.borrow();
-          read.accept(childRef, content);
+          final int idx = i;
+          //noinspection rawtypes
+          ChildRef childRef =
+              handle.borrow(r -> r.getObjectDataRanges(idx).getChecksummedData().getContent());
+          read.accept(childRef);
           if (d.getRangeEnd()) {
             // invoke eof on exec, the resolving future could have a downstream callback
             // that we don't want to block this grpc thread
-            exec.execute(read::eof);
-            // don't remove the outstanding read until the future has been resolved
-            state.outstandingReads.remove(id);
+            exec.execute(
+                () -> {
+                  try {
+                    read.eof();
+                    // don't remove the outstanding read until the future has been resolved
+                    state.outstandingReads.remove(id);
+                  } catch (IOException e) {
+                    // TODO: sync this up with stream restarts when the time comes
+                    throw StorageException.coalesce(e);
+                  }
+                });
           }
         }
       } catch (IOException e) {
@@ -290,28 +302,34 @@ protected void onCompleteImpl() {}
   static final class OutstandingReadToArray {
     private final long readId;
     private final ReadCursor readCursor;
-    private final ByteArrayOutputStream bytes;
+    private final List<ChildRef> childRefs;
     private final SettableApiFuture<byte[]> complete;
 
     @VisibleForTesting
     OutstandingReadToArray(
         long readId, long readOffset, long readLimit, SettableApiFuture<byte[]> complete) {
       this.readId = readId;
       this.readCursor = new ReadCursor(readOffset, readOffset + readLimit);
-      this.bytes = new ByteArrayOutputStream();
+      this.childRefs = Collections.synchronizedList(new ArrayList<>());
       this.complete = complete;
     }
 
-    public void accept(ChildRef childRef, ByteString bytes) throws IOException {
-      try (ChildRef autoclose = childRef) {
-        int size = bytes.size();
-        bytes.writeTo(this.bytes);
-        readCursor.advance(size);
-      }
+    public void accept(ChildRef childRef) throws IOException {
+      int size = childRef.byteString().size();
+      childRefs.add(childRef);
+      readCursor.advance(size);
     }
 
-    public void eof() {
-      complete.set(bytes.toByteArray());
+    public void eof() throws IOException {
+      try {
+        ByteString base = ByteString.empty();
+        for (ChildRef ref : childRefs) {
+          base = base.concat(ref.byteString());
+        }
+        complete.set(base.toByteArray());
+      } finally {
+        GrpcUtils.closeAll(childRefs);
+      }
     }
 
     public ReadRange makeReadRange() {

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

@@ -100,7 +100,6 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.IdentityHashMap;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -1133,40 +1132,14 @@ private Response parseFrom(CodedInputStream stream) throws InvalidProtocolBuffer
     }
 
     @Override
-    public ResponseContentLifecycleHandle get(Response response) {
+    public ResponseContentLifecycleHandle<Response> get(Response response) {
       InputStream stream = unclosedStreams.remove(response);
       return ResponseContentLifecycleHandle.create(response, toByteBuffersFunction, stream);
     }
 
     @Override
     public void close() throws IOException {
-      closeAllStreams(unclosedStreams.values());
-    }
-
-    /**
-     * In the event closing the streams results in multiple streams throwing IOExceptions, collect
-     * them all as suppressed exceptions on the first occurrence.
-     */
-    @VisibleForTesting
-    static void closeAllStreams(Iterable<InputStream> inputStreams) throws IOException {
-      Iterator<InputStream> iterator = inputStreams.iterator();
-      IOException ioException = null;
-      while (iterator.hasNext()) {
-        InputStream next = iterator.next();
-        try {
-          next.close();
-        } catch (IOException e) {
-          if (ioException == null) {
-            ioException = e;
-          } else if (ioException != e) {
-            ioException.addSuppressed(e);
-          }
-        }
-      }
-
-      if (ioException != null) {
-        throw ioException;
-      }
+      GrpcUtils.closeAll(unclosedStreams.values());
     }
   }
 

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

@@ -19,7 +19,11 @@
 import com.google.api.gax.grpc.GrpcCallContext;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Collection;
 import java.util.Locale;
+import java.util.Objects;
 
 final class GrpcUtils {
 
@@ -34,4 +38,38 @@ static GrpcCallContext contextWithBucketName(String bucketName, GrpcCallContext
     }
     return baseContext;
   }
+
+  /**
+   * In the event closing the streams results in multiple streams throwing IOExceptions, collect
+   * them all as suppressed exceptions on the first occurrence.
+   */
+  static <C extends Closeable> void closeAll(Collection<C> closeables) throws IOException {
+    IOException ioException =
+        closeables.stream()
+            .map(
+                stream -> {
+                  try {
+                    stream.close();
+                    return null;
+                  } catch (IOException e) {
+                    return e;
+                  }
+                })
+            .filter(Objects::nonNull)
+            .reduce(
+                null,
+                (l, r) -> {
+                  if (l != null) {
+                    l.addSuppressed(r);
+                    return l;
+                  } else {
+                    return r;
+                  }
+                },
+                (l, r) -> l);
+
+    if (ioException != null) {
+      throw ioException;
+    }
+  }
 }

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

@@ -15,8 +15,10 @@
  */
 package com.google.cloud.storage;
 
+import com.google.cloud.storage.BlobDescriptor.ZeroCopySupport.DisposableByteString;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Suppliers;
+import com.google.protobuf.ByteString;
 import java.io.Closeable;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -27,33 +29,37 @@
 import java.util.function.Supplier;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
-final class ResponseContentLifecycleHandle implements Closeable {
+final class ResponseContentLifecycleHandle<Response> implements Closeable {
+
+  private final Response response;
   @Nullable private final Closeable dispose;
 
   private final Supplier<List<ByteBuffer>> lazyBuffers;
   private final AtomicBoolean open;
   private final AtomicInteger refs;
 
   private ResponseContentLifecycleHandle(
-      Supplier<List<ByteBuffer>> lazyBuffers, @Nullable Closeable dispose) {
+      Response response,
+      Function<Response, List<ByteBuffer>> toBuffersFunction,
+      @Nullable Closeable dispose) {
+    this.response = response;
     this.dispose = dispose;
-    this.lazyBuffers = lazyBuffers;
+    this.lazyBuffers = Suppliers.memoize(() -> toBuffersFunction.apply(response));
     this.open = new AtomicBoolean(true);
     this.refs = new AtomicInteger(1);
   }
 
-  static <Response> ResponseContentLifecycleHandle create(
+  static <Response> ResponseContentLifecycleHandle<Response> create(
       Response response,
       Function<Response, List<ByteBuffer>> toBuffersFunction,
       @Nullable Closeable dispose) {
-    Supplier<List<ByteBuffer>> lazyBuffers =
-        Suppliers.memoize(() -> toBuffersFunction.apply(response));
-    return new ResponseContentLifecycleHandle(lazyBuffers, dispose);
+    return new ResponseContentLifecycleHandle<>(response, toBuffersFunction, dispose);
   }
 
-  ChildRef borrow() {
+  ChildRef borrow(Function<Response, ByteString> toByteStringFunction) {
     Preconditions.checkState(open.get(), "only able to borrow when open");
-    ChildRef childRef = new ChildRef();
+    Preconditions.checkNotNull(toByteStringFunction);
+    ChildRef childRef = new ChildRef(toByteStringFunction);
     refs.incrementAndGet();
     return childRef;
   }
@@ -91,7 +97,18 @@ private void dispose() throws IOException {
     }
   }
 
-  final class ChildRef implements Closeable {
+  final class ChildRef implements Closeable, DisposableByteString {
+
+    private final Function<Response, ByteString> toByteStringFunction;
+
+    private ChildRef(Function<Response, ByteString> toByteStringFunction) {
+      this.toByteStringFunction = toByteStringFunction;
+    }
+
+    @Override
+    public ByteString byteString() {
+      return toByteStringFunction.apply(response);
+    }
 
     @Override
     public void close() throws IOException {

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

@@ -21,7 +21,7 @@
 import java.io.IOException;
 
 interface ResponseContentLifecycleManager<Response> extends Closeable {
-  ResponseContentLifecycleHandle get(Response response);
+  ResponseContentLifecycleHandle<Response> get(Response response);
 
   @Override
   default void close() throws IOException {}