Prototype implementation of zero-copy ReadRows decoding

google-cloud-bigquerystorage/src/main/java/com/google/api/gax/rpc/ZeroCopyMessageMarshaller.java

@@ -0,0 +1,140 @@
+/*
+ * Copyright 2020 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
+ *
+ *     https://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.api.gax.rpc;
+
+import com.google.protobuf.ByteString;
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.MessageLite;
+import com.google.protobuf.Parser;
+import com.google.protobuf.UnsafeByteOperations;
+import io.grpc.Detachable;
+import io.grpc.HasByteBuffer;
+import io.grpc.KnownLength;
+import io.grpc.MethodDescriptor.PrototypeMarshaller;
+import io.grpc.Status;
+import io.grpc.protobuf.lite.ProtoLiteUtils;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.List;
+import java.util.Map;
+
+public class ZeroCopyMessageMarshaller<T extends MessageLite> implements PrototypeMarshaller<T> {
+  private final Map<T, InputStream> unclosedStreams =
+      Collections.synchronizedMap(new IdentityHashMap<>());
+  private final Parser<T> parser;
+  private final PrototypeMarshaller<T> marshaller;
+
+  @SuppressWarnings("unchecked")
+  public ZeroCopyMessageMarshaller(T defaultInstance) {
+    parser = (Parser<T>) defaultInstance.getParserForType();
+    marshaller = (PrototypeMarshaller<T>) ProtoLiteUtils.marshaller(defaultInstance);
+  }
+
+  @Override
+  public Class<T> getMessageClass() {
+    return marshaller.getMessageClass();
+  }
+
+  @Override
+  public T getMessagePrototype() {
+    return marshaller.getMessagePrototype();
+  }
+
+  @Override
+  public InputStream stream(T value) {
+    return marshaller.stream(value);
+  }
+
+  @Override
+  public T parse(InputStream stream) {
+    try {
+      if (stream instanceof KnownLength
+          && stream instanceof Detachable
+          && stream instanceof HasByteBuffer
+          && ((HasByteBuffer) stream).byteBufferSupported()) {
+        int size = stream.available();
+        // Stream is now detached here and should be closed later.
+        InputStream detachedStream = ((Detachable) stream).detach();
+        try {
+          // This call is to keep buffer while traversing buffers using skip.
+          detachedStream.mark(size);
+          List<ByteString> byteStrings = new ArrayList<>();
+          while (detachedStream.available() != 0) {
+            ByteBuffer buffer = ((HasByteBuffer) detachedStream).getByteBuffer();
+            byteStrings.add(UnsafeByteOperations.unsafeWrap(buffer));
+            detachedStream.skip(buffer.remaining());
+          }
+          detachedStream.reset();
+          CodedInputStream codedInputStream = ByteString.copyFrom(byteStrings).newCodedInput();
+          codedInputStream.enableAliasing(true);
+          codedInputStream.setSizeLimit(Integer.MAX_VALUE);
+          // Fast path (no memory copy).
+          T message;
+          try {
+            message = parseFrom(codedInputStream);
+          } catch (InvalidProtocolBufferException ipbe) {
+            throw Status.INTERNAL
+                .withDescription("Invalid protobuf byte sequence")
+                .withCause(ipbe)
+                .asRuntimeException();
+          }
+          unclosedStreams.put(message, detachedStream);
+          detachedStream = null;
+          return message;
+        } finally {
+          if (detachedStream != null) {
+            detachedStream.close();
+          }
+        }
+      }
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+    // Slow path.
+    return marshaller.parse(stream);
+  }
+
+  private T parseFrom(CodedInputStream stream) throws InvalidProtocolBufferException {
+    T message = parser.parseFrom(stream);
+    try {
+      stream.checkLastTagWas(0);
+      return message;
+    } catch (InvalidProtocolBufferException e) {
+      e.setUnfinishedMessage(message);
+      throw e;
+    }
+  }
+
+  /**
+   * Application code must call this method to claim the stream for the message and must call
+   * stream.close() in order to return it to the pool.
+   */
+  public InputStream popStream(T message) {
+    return unclosedStreams.remove(message);
+  }
+
+  public List<InputStream> popAllStreams() {
+    List<InputStream> streams = new ArrayList<>(unclosedStreams.values());
+    unclosedStreams.clear();
+    return streams;
+  }
+}

google-cloud-bigquerystorage/src/main/java/com/google/api/gax/rpc/ZeroCopyQueuingResponseObserver.java

@@ -0,0 +1,107 @@
+/*
+ * Copyright 2020 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
+ *
+ *     https://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.api.gax.rpc;
+
+import com.google.common.collect.Queues;
+import com.google.protobuf.MessageLite;
+import java.io.InputStream;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.BlockingQueue;
+
+/**
+ * This class should extend QueuingResponseObserver, but it is declared final.
+ *
+ * @param <T>
+ */
+public class ZeroCopyQueuingResponseObserver<T extends MessageLite>
+    extends StateCheckingResponseObserver<T> {
+  static final Object EOF_MARKER = new Object();
+
+  private final ZeroCopyMessageMarshaller<T> marshaller;
+  private final Map<T, InputStream> unclosedStreams =
+      Collections.synchronizedMap(new IdentityHashMap<>());
+  private final BlockingQueue<Object> buffer = Queues.newArrayBlockingQueue(2);
+  private StreamController controller;
+  private boolean isCancelled;
+
+  ZeroCopyQueuingResponseObserver(ZeroCopyMessageMarshaller<T> marshaller) {
+    this.marshaller = marshaller;
+  }
+
+  void request() {
+    controller.request(1);
+  }
+
+  Object getNext() throws InterruptedException {
+    if (isCancelled) {
+      return EOF_MARKER;
+    }
+    return buffer.take();
+  }
+
+  boolean isReady() {
+    return isCancelled || !buffer.isEmpty();
+  }
+
+  void cancel() {
+    isCancelled = true;
+    controller.cancel();
+  }
+
+  @Override
+  protected void onStartImpl(StreamController controller) {
+    this.controller = controller;
+    controller.disableAutoInboundFlowControl();
+    controller.request(1);
+  }
+
+  @Override
+  protected void onResponseImpl(T response) {
+    // Claim ownership of the stream before inserting the response.
+    if (marshaller != null) {
+      InputStream stream = marshaller.popStream(response);
+      if (stream != null) {
+        unclosedStreams.put(response, stream);
+      }
+    }
+    buffer.add(response);
+  }
+
+  @Override
+  protected void onErrorImpl(Throwable t) {
+    buffer.add(t);
+  }
+
+  @Override
+  protected void onCompleteImpl() {
+    buffer.add(EOF_MARKER);
+  }
+
+  public InputStream popStream(T response) {
+    return unclosedStreams.remove(response);
+  }
+
+  public List<InputStream> popAllStreams() {
+    List<InputStream> streams = new ArrayList<>(unclosedStreams.values());
+    unclosedStreams.clear();
+    return streams;
+  }
+}

google-cloud-bigquerystorage/src/main/java/com/google/api/gax/rpc/ZeroCopyServerStream.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2020 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
+ *
+ *     https://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.api.gax.rpc;
+
+import com.google.api.core.InternalApi;
+import com.google.protobuf.MessageLite;
+import java.io.InputStream;
+import java.util.Iterator;
+import java.util.List;
+import javax.annotation.Nonnull;
+
+/**
+ * This class is identical to {@link ServerStream}, except it returns a zero-copy iterator.
+ *
+ * @param <T>
+ */
+public class ZeroCopyServerStream<T extends MessageLite> extends ServerStream<T> {
+  private final ZeroCopyQueuingResponseObserver<T> observer;
+  private final ZeroCopyServerStreamIterator<T> iterator;
+  private boolean consumed;
+
+  public ZeroCopyServerStream(ZeroCopyMessageMarshaller<T> marshaller) {
+    observer = new ZeroCopyQueuingResponseObserver<>(marshaller);
+    iterator = new ZeroCopyServerStreamIterator<>(observer);
+  }
+
+  @Override
+  @InternalApi
+  ResponseObserver<T> observer() {
+    return observer;
+  }
+
+  @Override
+  @Nonnull
+  public Iterator<T> iterator() {
+    if (consumed) {
+      throw new IllegalStateException("Iterator already consumed");
+    }
+    consumed = true;
+    return iterator;
+  }
+
+  @Override
+  public boolean isReceiveReady() {
+    return iterator.isReady();
+  }
+
+  @Override
+  public void cancel() {
+    observer.cancel();
+  }
+
+  public InputStream popStream(T response) {
+    return observer.popStream(response);
+  }
+
+  public List<InputStream> popAllStreams() {
+    return observer.popAllStreams();
+  }
+}